submission 755142

  ⋯ 3 unchanged lines
  # ///
  # leaderboard = "amd-mxfp4-mm"
   
- """
- v867: Hardcoded 6-shape dispatcher. Zero dynamic dispatch overhead.
- All configs pre-computed. No dicts, no cache, no probes, no if/elif range checks.
- Exact (M,N,K) tuple matching. Pre-allocated tensors on first call.
- """
- import os, sys
- os.environ["HIP_FORCE_DEV_KERNARG"] = "1"
-  
- # PATCH: Remove denormal handling from AITER's _mxfp4_quant_op (saves ~0.7μs on S5)
- # This patches the REFERENCE too, so all quant paths must use the same patch
- _qp = "/home/runner/aiter/aiter/ops/triton/_triton_kernels/quant/quant.py"
- try:
-     with open(_qp, 'r') as f: _qc = f.read()
-     # PATCH A: Denormal removal
-     if '(not saturate_mask) & (qx_fp32 < min_normal)' in _qc:
-         _qc = _qc.replace('(not saturate_mask) & (qx_fp32 < min_normal)',
-                           'saturate_mask & (not saturate_mask)  # PATCHED')
-     # PATCH B: Integer exponent extraction (replaces log2+floor, saves 2 transcendental ops)
-     # PATCH B+E: Skip pow2 round, extract exponent directly from amax, +1 to compensate
-     # Original: amax → int → (+0x200000)&0xFF800000 → float → log2 → floor → -2
-     # New: amax → int → shift → mask → float → +1 → -129
-     # Also skip the pow2 rounding step entirely (3 ops saved)
-     old_s = '    amax = tl.max(tl.abs(x), axis=-1, keep_dims=True)\n    amax = amax.to(tl.int32, bitcast=True)\n    amax = (amax + 0x200000).to(tl.uint32, bitcast=True) & 0xFF800000\n    amax = amax.to(tl.float32, bitcast=True)\n    scale_e8m0_unbiased = tl.log2(amax).floor() - 2'
-     new_s = '    amax = tl.max(tl.abs(x), axis=-1, keep_dims=True)\n    amax_i = amax.to(tl.int32, bitcast=True)\n    scale_e8m0_unbiased = ((amax_i >> 23) & 0xFF).to(tl.float32) - 128.0'
-     if old_s in _qc:
-         _qc = _qc.replace(old_s, new_s)
-     # PATCH C: Replace exp2(-scale) with integer float construction (saves 1 transcendental op)
-     old_exp = '    quant_scale = tl.exp2(-scale_e8m0_unbiased)'
-     new_exp = '    qs_exp = tl.clamp(-scale_e8m0_unbiased + 127.0, 1.0, 254.0).to(tl.uint32)\n    quant_scale = (qs_exp << 23).to(tl.float32, bitcast=True)'
-     if old_exp in _qc:
-         _qc = _qc.replace(old_exp, new_exp)
-     # PATCH D: Direct bs_e8m0 from exponent (skip float intermediate)
-     old_bs = '    bs_e8m0 = scale_e8m0_unbiased.to(tl.uint8) + 127'
-     new_bs = '    bs_e8m0 = tl.clamp(scale_e8m0_unbiased + 127.0, 0.0, 254.0).to(tl.uint8)'
-     if old_bs in _qc:
-         _qc = _qc.replace(old_bs, new_bs)
-     # PATCH E: Remove saturate+denormal merge (replace 3 where/full with direct assignment)
-     old_merge = '''    # Merge results
-     e2m1_value = tl.full(qx.type.get_block_shapes(), 0x7, dtype=tl.uint8)
-     e2m1_value = tl.where(normal_mask, normal_x, e2m1_value)
-     e2m1_value = tl.where(denormal_mask, denormal_x, e2m1_value)'''
-     new_merge = '''    # Merge results (saturate+denormal proven dead, direct assign)
-     e2m1_value = normal_x'''
-     if old_merge in _qc:
-         _qc = _qc.replace(old_merge, new_merge)
-     # PATCH F: Remove mant_odd (saves 2 ops/element, both sides match)
-     old_mant = '    # rounding bias part 2\n    normal_x += mant_odd'
-     new_mant = '    # mant_odd removed for speed'
-     if old_mant in _qc:
-         _qc = _qc.replace(old_mant, new_mant)
-     old_val = '((EXP_BIAS_FP4 - EXP_BIAS_FP32) << MBITS_F32) + (1 << 21) - 1'
-     new_val = '((EXP_BIAS_FP4 - EXP_BIAS_FP32) << MBITS_F32) + (1 << 21)'
-     if old_val in _qc:
-         _qc = _qc.replace(old_val, new_val)
-     with open(_qp, 'w') as f: f.write(_qc)
- except: pass
-  
- # PATCH 2: fast_math=True in preshuffle dot_scaled (reduces VALU ops)
- _kp = "/home/runner/aiter/aiter/ops/triton/_triton_kernels/gemm/basic/gemm_a16wfp4.py"
- try:
-     with open(_kp, 'r') as f: _kc = f.read()
-     if 'accumulator += tl.dot_scaled' in _kc:
-         _kc = _kc.replace(
-             'accumulator += tl.dot_scaled(a, a_scales, "e2m1", b, b_scales, "e2m1")',
-             'accumulator = tl.dot_scaled(a, a_scales, "e2m1", b, b_scales, "e2m1", accumulator, fast_math=True)'
-         )
-         with open(_kp, 'w') as f: f.write(_kc)
- except: pass
-  
  from task import input_t, output_t
  import torch
  import triton
  import triton.language as tl
- import aiter
- from aiter import dtypes as _dt
   
- _FP4X2 = _dt.fp4x2
- _E8M0 = _dt.fp8_e8m0
   
- # ── Preshuffle (shape 1) — lazy import ──
- _preshuffle = None
-  
- # ── Pre-computed configs (populated on first call per shape) ──
- _s = {}  # shape key -> pre-allocated state
-  
-  
- # ═══════════════════════════════════════════════════════════════════
- # Triton kernels — identical to v866, no changes
- # ═══════════════════════════════════════════════════════════════════
-  
  @triton.jit
- def _mxfp4_quant_op(x, BLOCK_K: tl.constexpr, BLOCK_M: tl.constexpr):
-     EXP_BIAS_FP32: tl.constexpr = 127; EXP_BIAS_FP4: tl.constexpr = 1
-     MBITS_F32: tl.constexpr = 23; MBITS_FP4: tl.constexpr = 1
-     EBITS_F32: tl.constexpr = 8; EBITS_FP4: tl.constexpr = 2
-     max_normal: tl.constexpr = 6; min_normal: tl.constexpr = 1
-     QUANT: tl.constexpr = 32; NUM_QB: tl.constexpr = BLOCK_K // QUANT
-     x = x.reshape(BLOCK_M, NUM_QB, QUANT)
-     amax = tl.max(tl.abs(x), axis=-1, keep_dims=True)
-     amax_i = amax.to(tl.int32, bitcast=True)
-     scale_unb = ((amax_i >> 23) & 0xFF).to(tl.float32) - 128.0
-     scale_unb = tl.clamp(scale_unb, min=-127, max=127)
-     bs = tl.clamp(scale_unb + 127.0, 0.0, 254.0).to(tl.uint8)
-     qs_exp = tl.clamp(-scale_unb + 127.0, 1.0, 254.0).to(tl.uint32)
-     qscale = (qs_exp << 23).to(tl.float32, bitcast=True)
-     qx = x * qscale; qx = qx.to(tl.uint32, bitcast=True)
-     s = qx & 0x80000000; qx = qx ^ s
-     normal_x = qx.to(tl.int32)
-     val_to_add: tl.constexpr = ((EXP_BIAS_FP4 - EXP_BIAS_FP32) << MBITS_F32) + (1 << 21)
-     normal_x += val_to_add
-     normal_x = normal_x >> (MBITS_F32 - MBITS_FP4); normal_x = normal_x.to(tl.uint8)
-     e2m1 = normal_x
-     sign_lp = s >> (MBITS_F32 + EBITS_F32 - MBITS_FP4 - EBITS_FP4); sign_lp = sign_lp.to(tl.uint8)
-     e2m1 = e2m1 | sign_lp
-     e2m1 = tl.reshape(e2m1, [BLOCK_M, NUM_QB, QUANT // 2, 2])
-     evens, odds = tl.split(e2m1); fp4 = evens | (odds << 4)
-     return fp4.reshape(BLOCK_M, BLOCK_K // 2), bs.reshape(BLOCK_M, NUM_QB)
+ def _hw_quant_fp4(
+     inp_bf16,
+     TM: tl.constexpr,
+     TK: tl.constexpr,
+ ):
+     GRP: tl.constexpr = 32
+     NB: tl.constexpr = TK // GRP
   
+     xf = inp_bf16.to(tl.float32).reshape(TM, NB, GRP)
+     peak = tl.max(tl.abs(xf), axis=-1, keep_dims=True)
+     peak = peak.to(tl.int32, bitcast=True)
+     peak = (peak + 0x200000).to(tl.uint32, bitcast=True) & 0xFF800000
+     exp_val = ((peak >> 23) & 0xFF).to(tl.int32) - 127
+     scale_unb = exp_val - 2
+     scale_unb = tl.minimum(tl.maximum(scale_unb, -127), 127)
+     e8m0 = scale_unb.to(tl.uint8) + 127
   
- @triton.jit
- def xcd_swizzle(pid, domain_size, XCD_SWIZZLE: tl.constexpr):
-     return (pid % XCD_SWIZZLE) * (domain_size // XCD_SWIZZLE) + tl.minimum(pid % XCD_SWIZZLE, domain_size % XCD_SWIZZLE) + pid // XCD_SWIZZLE
+     div_bits = (scale_unb.to(tl.int32) + 127).to(tl.uint32) << 23
+     div_scale = div_bits.to(tl.float32, bitcast=True)
   
+     div_bc = tl.broadcast_to(div_scale, (TM, NB, GRP)).reshape(TM, TK)
+     div_pairs = div_bc.reshape(TM, TK // 2, 2)
+     div_lo, _ = tl.split(div_pairs)
+     div_flat = div_lo.reshape(TM, TK // 2)
   
- @triton.jit
- def _fused_quant_gemm_kernel(
-     A_ptr, Bq_ptr, Bscale_sh_ptr, C_ptr, M, N, K: tl.constexpr,
-     stride_a_m, stride_a_k, stride_bq_n, stride_bq_k, SN_DIV8_MUL256: tl.constexpr,
-     stride_c_m, stride_c_n,
-     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
- ):
-     pid_m = tl.program_id(0); pid_n = tl.program_id(1)
-     offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M); offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
-     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
-     QUANT: tl.constexpr = 32; NSK: tl.constexpr = BLOCK_K // QUANT
-     for ki in tl.range(0, K, BLOCK_K):
-         a_offs_k = ki + tl.arange(0, BLOCK_K); a_ptrs = A_ptr + offs_m[:, None] * stride_a_m + a_offs_k[None, :] * stride_a_k; a_mask = (offs_m < M)[:, None] & (a_offs_k < K)[None, :]; a_tile = tl.load(a_ptrs, mask=a_mask, other=0.0).to(tl.float32); a_fp4, a_scale = _mxfp4_quant_op(a_tile, BLOCK_K, BLOCK_M)
-         b_offs_k = ki // 2 + tl.arange(0, BLOCK_K // 2); b_ptrs = Bq_ptr + offs_n[None, :] * stride_bq_n + b_offs_k[:, None] * stride_bq_k; b_mask = (offs_n < N)[None, :] & (b_offs_k < K // 2)[:, None]; b_tile = tl.load(b_ptrs, mask=b_mask, other=0, cache_modifier=".cg")
-         bs_row = offs_n; bs_col_base = ki // QUANT; bs_col_offs = tl.arange(0, NSK); row = bs_row[:, None]; col = (bs_col_base + bs_col_offs)[None, :]
-         shuf_idx = (row // 32) * SN_DIV8_MUL256 + (col // 8) * 256 + (col % 4) * 64 + (row % 16) * 4 + ((col % 8) // 4) * 2 + ((row % 32) // 16)
-         bs_mask = (offs_n[:, None] < N) & (bs_col_offs[None, :] < (K // QUANT - bs_col_base)); b_scale = tl.load(Bscale_sh_ptr + shuf_idx, mask=bs_mask, other=0, cache_modifier=".cg")
-         acc = tl.dot_scaled(a_fp4, a_scale, "e2m1", b_tile, b_scale, "e2m1", acc, fast_math=True)
-     c_ptrs = C_ptr + offs_m[:, None] * stride_c_m + offs_n[None, :] * stride_c_n; c_mask = (offs_m < M)[:, None] & (offs_n < N)[None, :]
-     tl.store(c_ptrs, acc.to(tl.bfloat16), mask=c_mask)
+     raw16 = inp_bf16.to(tl.uint16, bitcast=True).reshape(TM, TK // 2, 2)
+     p0, p1 = tl.split(raw16)
+     packed = p0.to(tl.uint32) | (p1.to(tl.uint32) << 16)
+     packed = packed.reshape(TM, TK // 2)
   
+     result = tl.inline_asm_elementwise(
+         "v_cvt_scalef32_pk_fp4_bf16 $0, $1, $2",
+         "=v, v, v",
+         [packed, div_flat],
+         dtype=tl.uint32,
+         is_pure=True,
+         pack=1,
+     )
+     fp4_out = (result & 0xFF).to(tl.uint8).reshape(TM, TK // 2)
+     return fp4_out, e8m0.reshape(TM, NB)
   
- @triton.jit
- def _fused_splitk_gemm(
-     A_ptr, Bq_ptr, Bscale_sh_ptr, Y_ptr, M, N, K: tl.constexpr,
-     stride_a_m, stride_a_k, stride_bq_n, stride_bq_k, SN_DIV8_MUL256: tl.constexpr,
-     stride_y_k, stride_y_m, stride_y_n, grid_m, grid_n,
-     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
-     SPLIT_K: tl.constexpr, XCD_SWIZZLE: tl.constexpr,
-     matrix_instr_nonkdim: tl.constexpr = 32,
- ):
-     pid = tl.program_id(0); total_tiles = grid_m * grid_n * SPLIT_K
-     if XCD_SWIZZLE > 1: pid = xcd_swizzle(pid, total_tiles, XCD_SWIZZLE)
-     pid_k = pid % SPLIT_K; pid_mn = pid // SPLIT_K; pid_m = pid_mn // grid_n; pid_n = pid_mn % grid_n
-     offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M); offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
-     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
-     QUANT: tl.constexpr = 32; NSK: tl.constexpr = BLOCK_K // QUANT
-     k_per_split = ((K + SPLIT_K - 1) // SPLIT_K + BLOCK_K - 1) // BLOCK_K * BLOCK_K; k_start = pid_k * k_per_split; k_end = min(k_start + k_per_split, K)
-     for ki in tl.range(k_start, k_end, BLOCK_K):
-         a_offs_k = ki + tl.arange(0, BLOCK_K); a_ptrs = A_ptr + offs_m[:, None] * stride_a_m + a_offs_k[None, :] * stride_a_k; a_mask = (offs_m < M)[:, None] & (a_offs_k < K)[None, :]; a_tile = tl.load(a_ptrs, mask=a_mask, other=0.0).to(tl.float32); a_fp4, a_scale = _mxfp4_quant_op(a_tile, BLOCK_K, BLOCK_M)
-         b_offs_k = ki // 2 + tl.arange(0, BLOCK_K // 2); b_ptrs = Bq_ptr + offs_n[None, :] * stride_bq_n + b_offs_k[:, None] * stride_bq_k; b_mask = (offs_n < N)[None, :] & (b_offs_k < K // 2)[:, None]; b_tile = tl.load(b_ptrs, mask=b_mask, other=0, cache_modifier=".cg")
-         bs_row = offs_n; bs_col_base = ki // QUANT; bs_col_offs = tl.arange(0, NSK); row = bs_row[:, None]; col = (bs_col_base + bs_col_offs)[None, :]
-         shuf_idx = (row // 32) * SN_DIV8_MUL256 + (col // 8) * 256 + (col % 4) * 64 + (row % 16) * 4 + ((col % 8) // 4) * 2 + ((row % 32) // 16)
-         bs_mask = (offs_n[:, None] < N) & (bs_col_offs[None, :] < (K // QUANT - bs_col_base)); b_scale = tl.load(Bscale_sh_ptr + shuf_idx, mask=bs_mask, other=0, cache_modifier=".cg")
-         acc = tl.dot_scaled(a_fp4, a_scale, "e2m1", b_tile, b_scale, "e2m1", acc, fast_math=True)
-     y_ptrs = Y_ptr + pid_k * stride_y_k + offs_m[:, None] * stride_y_m + offs_n[None, :] * stride_y_n; y_mask = (offs_m < M)[:, None] & (offs_n < N)[None, :]
-     if SPLIT_K > 1: tl.store(y_ptrs, acc, mask=y_mask)
-     else: tl.store(y_ptrs, acc.to(tl.bfloat16), mask=y_mask)
   
-  
+ @triton.heuristics({
+     "ALIGNED_K": lambda args: (args["K"] % (args["BK"] // 2) == 0)
+         and (args["SK_BLOCK"] % args["BK"] == 0)
+         and (args["K"] % (args["SK_BLOCK"] // 2) == 0),
+ })
  @triton.jit
- def _reduce_splitk(Y_ptr, Out_ptr, M, N, stride_y_k, stride_y_m, stride_y_n, stride_o_m, stride_o_n, SPLIT_K: tl.constexpr, BLOCK_N: tl.constexpr):
-     pid_m = tl.program_id(0); pid_n = tl.program_id(1); offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N); n_mask = offs_n < N
-     acc = tl.zeros([BLOCK_N], dtype=tl.float32)
-     for k in tl.range(0, SPLIT_K): acc += tl.load(Y_ptr + k * stride_y_k + pid_m * stride_y_m + offs_n * stride_y_n, mask=n_mask, other=0.0).to(tl.float32)
-     tl.store(Out_ptr + pid_m * stride_o_m + offs_n * stride_o_n, acc.to(tl.bfloat16), mask=n_mask)
-  
-  
- @triton.jit
- def _fused_quant_shuffle_kernel(
-     x_ptr, x_fp4_ptr, bs_shuf_ptr, stride_x_m, stride_x_n, stride_fp4_m, stride_fp4_n,
-     M, N, SN_DIV8_MUL256, SCALE_COLS,
-     BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
-     NUM_ITER: tl.constexpr, NUM_STAGES: tl.constexpr,
+ def _gemm_fused(
+     a_ptr, w_ptr, out_ptr, ws_ptr,
+     M, N, K,
+     s_am, s_ak, s_wn, s_wk,
+     s_ok, s_om, s_on, s_wsn, s_wsk,
+     BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
+     GSM: tl.constexpr, NSK: tl.constexpr, SK_BLOCK: tl.constexpr,
+     ALIGNED_K: tl.constexpr,
+     num_warps: tl.constexpr, num_stages: tl.constexpr,
+     waves_per_eu: tl.constexpr, matrix_instr_nonkdim: tl.constexpr,
+     cache_modifier: tl.constexpr,
  ):
-     pid_m = tl.program_id(0); start_n = tl.program_id(1) * NUM_ITER
-     QUANT: tl.constexpr = 32; NUM_QB: tl.constexpr = BLOCK_SIZE_N // QUANT
-     for pid_n in tl.range(start_n, min(start_n + NUM_ITER, N), num_stages=NUM_STAGES):
-         x_offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M); x_offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-         x_offs = x_offs_m[:, None] * stride_x_m + x_offs_n[None, :] * stride_x_n; x_mask = (x_offs_m < M)[:, None] & (x_offs_n < N)[None, :]
-         x = tl.load(x_ptr + x_offs, mask=x_mask, other=0.0, cache_modifier=".cg").to(tl.float32)
-         out_tensor, bs_e8m0 = _mxfp4_quant_op(x, BLOCK_SIZE_N, BLOCK_SIZE_M)
-         out_offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M); out_offs_n = pid_n * BLOCK_SIZE_N // 2 + tl.arange(0, BLOCK_SIZE_N // 2)
-         out_offs = out_offs_m[:, None] * stride_fp4_m + out_offs_n[None, :] * stride_fp4_n; out_mask = (out_offs_m < M)[:, None] & (out_offs_n < (N // 2))[None, :]
-         tl.store(x_fp4_ptr + out_offs, out_tensor, mask=out_mask, cache_modifier=".cg")
-         bs_row = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M); bs_col = pid_n * NUM_QB + tl.arange(0, NUM_QB)
-         row = bs_row[:, None]; col = bs_col[None, :]
-         shuf_idx = (row // 32) * SN_DIV8_MUL256 + (col // 8) * 256 + (col % 4) * 64 + (row % 16) * 4 + ((col % 8) // 4) * 2 + ((row % 32) // 16)
-         bs_mask = (bs_row[:, None] < M) & (bs_col[None, :] < SCALE_COLS)
-         tl.store(bs_shuf_ptr + shuf_idx, bs_e8m0, mask=bs_mask, cache_modifier=".cg")
+     tl.assume(s_am > 0); tl.assume(s_ak > 0)
+     tl.assume(s_wn > 0); tl.assume(s_wk > 0)
+     tl.assume(s_om > 0); tl.assume(s_on > 0)
+     tl.assume(s_wsn > 0); tl.assume(s_wsk > 0)
   
+     SG: tl.constexpr = 32
+     nm = tl.cdiv(M, BM); nn = tl.cdiv(N, BN)
+     pid_all = tl.program_id(0)
+     pk = pid_all % NSK; pid = pid_all // NSK
   
- # ═══════════════════════════════════════════════════════════════════
- # Hardcoded per-shape helpers — kernel name builder
- # ═══════════════════════════════════════════════════════════════════
+     if NSK == 1:
+         grp_sz = GSM * nn
+         gid = pid // grp_sz
+         fm = gid * GSM
+         gsm = min(nm - fm, GSM)
+         pm = fm + ((pid % grp_sz) % gsm)
+         pn = (pid % grp_sz) // gsm
+     else:
+         pm = pid // nn; pn = pid % nn
   
- def _knl_name(tile_m, tile_n):
-     base = f"f4gemm_bf16_per1x32Fp4_BpreShuffle_{tile_m}x{tile_n}"
-     return f"_ZN5aiter{len(base)}{base}E"
+     tl.assume(pm >= 0); tl.assume(pn >= 0); tl.assume(pk >= 0)
   
- _KNL_32x128 = _knl_name(32, 128)
+     if (pk * SK_BLOCK // 2) < K:
+         n_iters = tl.cdiv(SK_BLOCK // 2, BK // 2)
+         row_m = (pm * BM + tl.arange(0, BM)) % M
+         col_k = pk * SK_BLOCK + tl.arange(0, BK)
+         a_p = a_ptr + (row_m[:, None] * s_am + col_k[None, :] * s_ak)
   
+         shuf_range = tl.arange(0, (BK // 2) * 16)
+         shuf_off = pk * (SK_BLOCK // 2) * 16 + shuf_range
+         w_row = (pn * (BN // 16) + tl.arange(0, BN // 16)) % (N // 16)
+         w_p = w_ptr + (w_row[:, None] * s_wn + shuf_off[None, :] * s_wk)
   
- # ═══════════════════════════════════════════════════════════════════
- # Shape-specific init functions — called ONCE per shape
- # ═══════════════════════════════════════════════════════════════════
+         sc_row = (pn * BN + tl.arange(0, BN // 32) * 32)
+         sc_col = (pk * (SK_BLOCK // SG) * 32) + tl.arange(0, BK // SG * 32)
+         ws_p = ws_ptr + sc_row[:, None] * s_wsn + sc_col[None, :] * s_wsk
   
- def _init_shape1(dev):
-     """Shape 1: M=4, N=2880, K=512 — Preshuffle path"""
-     return {'ready': True}
+         acc = tl.zeros((BM, BN), dtype=tl.float32)
   
+         for ki in range(pk * n_iters, (pk + 1) * n_iters):
+             if ALIGNED_K:
+                 a_tile = tl.load(a_p)
+                 wsc_raw = tl.load(ws_p, cache_modifier=cache_modifier)
+                 w_raw = tl.load(w_p, cache_modifier=cache_modifier)
+             else:
+                 koff = (ki - pk * n_iters) * BK
+                 a_tile = tl.load(a_p, mask=tl.arange(0, BK)[None, :] < (2 * K - pk * SK_BLOCK - koff), other=0.0)
+                 wsc_raw = tl.load(ws_p, cache_modifier=cache_modifier)
+                 w_raw = tl.load(w_p, mask=shuf_range[None, :] < ((K - (pk * (SK_BLOCK // 2) + (ki - pk * n_iters) * (BK // 2))) * 16), other=0, cache_modifier=cache_modifier)
   
- def _init_shape2(dev):
-     """Shape 2: M=16, N=2112, K=7168 — SplitK path, SK=14"""
-     M, N, K = 16, 2112, 7168
-     BM, BN, BK = 16, 128, 512
-     m_tiles = 1  # ceil(16/16)
-     n_tiles = 17  # ceil(2112/128) = 16.5 → 17
-     SK = 14
-     total_wgs = m_tiles * n_tiles * SK  # 1 * 17 * 14 = 238
-     sn_div8 = (((K // 32) + 7) // 8)  # ceil(224/8) = 28
-     sn_div8_mul256 = sn_div8 * 256  # 7168
-     out = torch.empty(M, N, dtype=torch.bfloat16, device=dev)
-     scratch = torch.empty(SK, M, N, dtype=torch.float32, device=dev)
-     return {
-         'out': out, 'scratch': scratch,
-         'total_wgs': total_wgs, 'grid_m': m_tiles, 'grid_n': n_tiles,
-         'sn_div8_mul256': sn_div8_mul256,
-         'reduce_grid': (M, triton.cdiv(N, 128)),
-     }
+             a_q, a_sc = _hw_quant_fp4(a_tile, BM, BK)
   
+             wsc = (wsc_raw.reshape(BN // 32, BK // SG // 8, 4, 16, 2, 2, 1)
+                    .permute(0, 5, 3, 1, 4, 2, 6).reshape(BN, BK // SG))
   
- def _init_shape3(dev):
-     """Shape 3: M=32, N=4096, K=512 — Fused path"""
-     M, N, K = 32, 4096, 512
-     BM, BN = 16, 64
-     grid = (triton.cdiv(M, BM), triton.cdiv(N, BN))  # (2, 64)
-     total_wgs = grid[0] * grid[1]  # 128
-     sn_div8_mul256 = (((K // 32 + 7) // 8)) * 256  # ceil(16/8)*256 = 512
-     return {
-         'out': torch.empty(M, N, dtype=torch.bfloat16, device=dev),
-         'grid': grid, 'sn_div8_mul256': sn_div8_mul256,
-         'wpe': 2 if total_wgs > 256 else 1,
-     }
+             w_tile = (w_raw.reshape(1, BN // 16, BK // 64, 2, 16, 16)
+                       .permute(0, 1, 4, 2, 3, 5).reshape(BN, BK // 2).trans(1, 0))
   
+             acc = tl.dot_scaled(a_q, a_sc, "e2m1", w_tile, wsc, "e2m1", acc, fast_math=True)
+             a_p += BK * s_ak
+             w_p += (BK // 2) * 16 * s_wk
+             ws_p += BK * s_wsk
   
- def _init_shape4(dev):
-     """Shape 4: M=32, N=2880, K=512 — Fused path"""
-     M, N, K = 32, 2880, 512
-     BM, BN = 16, 64
-     grid = (triton.cdiv(M, BM), triton.cdiv(N, BN))  # (2, 45)
-     total_wgs = grid[0] * grid[1]  # 90
-     sn_div8_mul256 = (((K // 32 + 7) // 8)) * 256  # 512
-     return {
-         'out': torch.empty(M, N, dtype=torch.bfloat16, device=dev),
-         'grid': grid, 'sn_div8_mul256': sn_div8_mul256,
-         'wpe': 1,
-     }
+         res = acc.to(out_ptr.type.element_ty)
+         o_m = pm * BM + tl.arange(0, BM).to(tl.int64)
+         o_n = pn * BN + tl.arange(0, BN).to(tl.int64)
+         o_p = out_ptr + s_om * o_m[:, None] + s_on * o_n[None, :] + pk * s_ok
+         tl.store(o_p, res, mask=(o_m[:, None] < M) & (o_n[None, :] < N))
   
   
- def _init_shape5(dev):
-     """Shape 5: M=64, N=7168, K=2048 — Quant + CK ASM, l2ks=3"""
-     M, N, K = 64, 7168, 2048
-     sm = 256  # ((64+255)//256)*256
-     sc = K // 32  # 64
-     sn = ((sc + 7) // 8) * 8  # 64
-     sn_div8_mul256 = (sn // 8) * 256  # 2048
-     x_fp4 = torch.empty((M, K // 2), dtype=torch.uint8, device=dev)
-     bs_shuffled = torch.empty(sm, sn, dtype=torch.uint8, device=dev)
-     out = torch.empty(M, N, dtype=torch.bfloat16, device=dev)
-     # Quant grid: BSM=4, BSN=128, NI=1 → (64/4, 2048/128) = (16, 16)
-     qgrid = (16, 16)
-     return {
-         'x_fp4': x_fp4, 'bs_shuffled': bs_shuffled, 'out': out,
-         'x_fp4_v': x_fp4.view(_FP4X2), 'bs_shuf_v': bs_shuffled.view(_E8M0),
-         'sc': sc, 'sn_div8_mul256': sn_div8_mul256, 'qgrid': qgrid,
-     }
+ @triton.jit
+ def _sum_partials(
+     src, dst, M, N,
+     s_sk, s_sm, s_sn, s_dm, s_dn,
+     RM: tl.constexpr, RN: tl.constexpr,
+     ACTUAL_K: tl.constexpr, MAX_K: tl.constexpr,
+ ):
+     im = tl.program_id(0); jn = tl.program_id(1)
+     om = (im * RM + tl.arange(0, RM)) % M
+     on = (jn * RN + tl.arange(0, RN)) % N
+     base = src + om[:, None] * s_sm + on[None, :] * s_sn
+     total = tl.load(base).to(tl.float32)
+     for s in tl.static_range(1, MAX_K):
+         if s < ACTUAL_K:
+             total += tl.load(base + s * s_sk).to(tl.float32)
+     tl.store(dst + om[:, None] * s_dm + on[None, :] * s_dn, total.to(dst.type.element_ty))
   
   
- def _init_shape6(dev):
-     """Shape 6: M=256, N=3072, K=1536 — Quant + CK ASM, l2ks=2"""
-     M, N, K = 256, 3072, 1536
-     sm = 256  # ((256+255)//256)*256
-     sc = K // 32  # 48
-     sn = ((sc + 7) // 8) * 8  # 48
-     sn_div8_mul256 = (sn // 8) * 256  # 1536
-     x_fp4 = torch.empty((M, K // 2), dtype=torch.uint8, device=dev)
-     bs_shuffled = torch.empty(sm, sn, dtype=torch.uint8, device=dev)
-     out = torch.empty(M, N, dtype=torch.bfloat16, device=dev)
-     # Quant grid: BSM=16, BSN=64, NI=2 → (256/16, 1536/(64*2)) = (16, 12)
-     qgrid = (16, 12)
-     return {
-         'x_fp4': x_fp4, 'bs_shuffled': bs_shuffled, 'out': out,
-         'x_fp4_v': x_fp4.view(_FP4X2), 'bs_shuf_v': bs_shuffled.view(_E8M0),
-         'sc': sc, 'sn_div8_mul256': sn_div8_mul256, 'qgrid': qgrid,
-     }
+ def _compute_sk(K, BK, NSK):
+     SK_BLOCK = triton.cdiv((2 * triton.cdiv(K, NSK)), BK) * BK
+     while NSK > 1 and BK > 16:
+         if K % (SK_BLOCK // 2) == 0 and SK_BLOCK % BK == 0 and K % (BK // 2) == 0:
+             break
+         elif K % (SK_BLOCK // 2) != 0 and NSK > 1:
+             NSK //= 2
+         elif SK_BLOCK % BK != 0:
+             NSK = max(NSK // 2, 1) if NSK > 1 else NSK; BK = max(BK // 2, 16) if NSK <= 1 else BK
+         elif K % (BK // 2) != 0 and BK > 16:
+             BK //= 2
+         else:
+             break
+         SK_BLOCK = triton.cdiv((2 * triton.cdiv(K, NSK)), BK) * BK
+     NSK = triton.cdiv(K, (SK_BLOCK // 2))
+     return SK_BLOCK, BK, NSK
   
   
- # ═══════════════════════════════════════════════════════════════════
- # Shape-specific dispatch functions — ZERO overhead hot paths
- # ═══════════════════════════════════════════════════════════════════
+ _CFGS = {
+     (4, 2880, 512): dict(BM=4, BN=128, BK=256, GSM=1, nw=4, ns=2, wpe=1, nkd=16, cm=None, NSK=1),
+     (16, 2112, 7168): dict(BM=16, BN=128, BK=512, GSM=1, nw=4, ns=2, wpe=3, nkd=16, cm=".cg", NSK=14),
+     (32, 4096, 512): dict(BM=16, BN=32, BK=256, GSM=1, nw=4, ns=3, wpe=3, nkd=16, cm=".cg", NSK=1),
+     (32, 2880, 512): dict(BM=8, BN=128, BK=256, GSM=1, nw=4, ns=2, wpe=2, nkd=16, cm=None, NSK=1),
+     (64, 7168, 2048): dict(BM=16, BN=128, BK=256, GSM=1, nw=4, ns=2, wpe=3, nkd=16, cm=".cg", NSK=1),
+     (256, 3072, 1536): dict(BM=16, BN=256, BK=512, GSM=1, nw=8, ns=2, wpe=2, nkd=16, cm=None, NSK=1),
+ }
+ _DEF_CFG = dict(BM=16, BN=32, BK=256, GSM=1, nw=2, ns=2, wpe=0, nkd=16, cm=".cg", NSK=1)
   
- def _run_s1(A, B_shuffle, B_scale_sh):
-     """Shape 1: M=4, N=2880, K=512 — Preshuffle"""
-     global _preshuffle
-     if _preshuffle is None:
-         from aiter.ops.triton.gemm.basic.gemm_a16wfp4 import gemm_a16wfp4_preshuffle
-         _preshuffle = gemm_a16wfp4_preshuffle
-     K = 512; sc = 16; sn = 16; K_half = 256; N = 2880
-     padN = B_scale_sh.view(torch.uint8).shape[0]
-     bs_reshaped = B_scale_sh.view(torch.uint8).reshape(padN // 32, sn * 32)
-     b_shuf_reshaped = B_shuffle.view(torch.uint8).reshape(N // 16, K_half * 16)
-     # BN=64: 45 WGs (vs BN=128: 23 WGs). 2× CU utilization for S1.
-     config = {'BLOCK_SIZE_M': 4, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 1, 'NUM_KSPLIT': 1, 'SPLITK_BLOCK_SIZE': 512, 'matrix_instr_nonkdim': 16, 'num_warps': 4, 'num_stages': 2, 'waves_per_eu': 2, 'cache_modifier': '.cg'}
-     return _preshuffle(A, b_shuf_reshaped, bs_reshaped, prequant=True, dtype=torch.bfloat16, config=config)
+ _alloc = {}
+ _resolved = {}
+ _params = {}
   
   
- def _run_s2(A, B_q, B_scale_sh, c):
-     """Shape 2: M=16, N=2112, K=7168 — SplitK SK=14"""
-     Bq = B_q.view(torch.uint8); Bs = B_scale_sh.view(torch.uint8)
-     s = c['scratch']
-     _fused_splitk_gemm[(c['total_wgs'],)](
-         A, Bq, Bs, s, 16, 2112, 7168,
-         A.stride(0), A.stride(1), Bq.stride(0), Bq.stride(1),
-         c['sn_div8_mul256'], s.stride(0), s.stride(1), s.stride(2),
-         c['grid_m'], c['grid_n'],
-         BLOCK_M=16, BLOCK_N=128, BLOCK_K=512,
-         SPLIT_K=14, XCD_SWIZZLE=8,
-         matrix_instr_nonkdim=16,
-         num_warps=4, num_stages=1, waves_per_eu=2,
-     )
-     _reduce_splitk[c['reduce_grid']](
-         s, c['out'], 16, 2112,
-         s.stride(0), s.stride(1), s.stride(2),
-         c['out'].stride(0), c['out'].stride(1),
-         SPLIT_K=14, BLOCK_N=128, num_warps=4,
-     )
-     return c['out']
+ def _get_alloc(m, n, nsk, dev):
+     key = (m, n, nsk)
+     if key not in _alloc:
+         y = torch.empty((m, n), dtype=torch.bfloat16, device=dev)
+         pp = torch.empty((nsk, m, n), dtype=torch.float32, device=dev) if nsk > 1 else None
+         _alloc[key] = (y, pp)
+     return _alloc[key]
   
   
- def _run_s3(A, B_shuffle, B_scale_sh):
-     """Shape 3: M=32, N=4096, K=512 — Preshuffle BM=8 NW=8 (v996: -0.2μs)"""
-     global _preshuffle
-     if _preshuffle is None:
-         from aiter.ops.triton.gemm.basic.gemm_a16wfp4 import gemm_a16wfp4_preshuffle
-         _preshuffle = gemm_a16wfp4_preshuffle
-     K = 512; N = 4096; sc = K // 32; sn = ((sc+7)//8)*8; K_half = K // 2
-     padN = B_scale_sh.view(torch.uint8).shape[0]
-     bs_r = B_scale_sh.view(torch.uint8).reshape(padN // 32, sn * 32)
-     b_r = B_shuffle.view(torch.uint8).reshape(N // 16, K_half * 16)
-     # Try BM=4 NW=4 for S3: 32/4=8 M-tiles × 32 N-tiles = 256 WGs (perfect CU match!)
-     cfg = {"BLOCK_SIZE_M": 4, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 256, "GROUP_SIZE_M": 1, "NUM_KSPLIT": 1, "SPLITK_BLOCK_SIZE": 512, "matrix_instr_nonkdim": 16, "num_warps": 4, "num_stages": 2, "waves_per_eu": 2, "cache_modifier": ".cg"}
-     return _preshuffle(A, b_r, bs_r, prequant=True, dtype=torch.bfloat16, config=cfg)
+ def _resolve(m, n, k):
+     key = (m, n, k)
+     if key not in _resolved:
+         c = _CFGS.get(key, _DEF_CFG).copy()
+         kh = k // 2
+         if c["NSK"] > 1:
+             c["SK_BLOCK"], c["BK"], c["NSK"] = _compute_sk(kh, c["BK"], c["NSK"])
+         else:
+             c["SK_BLOCK"] = 2 * kh; c["NSK"] = 1
+         if c["BK"] >= 2 * kh:
+             c["BK"] = triton.next_power_of_2(2 * kh); c["SK_BLOCK"] = 2 * kh; c["NSK"] = 1
+         c["BN"] = max(c["BN"], 32)
+         _resolved[key] = c
+     return _resolved[key]
   
   
- def _run_s4(A, B_shuffle, B_scale_sh):
-     """Shape 4: M=32, N=2880, K=512 — Preshuffle BM=8 NW=4 (with fast_math patch!)"""
-     global _preshuffle
-     if _preshuffle is None:
-         from aiter.ops.triton.gemm.basic.gemm_a16wfp4 import gemm_a16wfp4_preshuffle
-         _preshuffle = gemm_a16wfp4_preshuffle
-     K = 512; N = 2880; K_half = K // 2; sc = K // 32; sn = ((sc+7)//8)*8
-     padN = B_scale_sh.view(torch.uint8).shape[0]
-     bs_r = B_scale_sh.view(torch.uint8).reshape(padN // 32, sn * 32)
-     b_r = B_shuffle.view(torch.uint8).reshape(N // 16, K_half * 16)
-     cfg = {"BLOCK_SIZE_M": 8, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 256, "GROUP_SIZE_M": 1, "NUM_KSPLIT": 1, "SPLITK_BLOCK_SIZE": 512, "matrix_instr_nonkdim": 16, "num_warps": 4, "num_stages": 2, "waves_per_eu": 2, "cache_modifier": ".cg"}
-     return _preshuffle(A, b_r, bs_r, prequant=True, dtype=torch.bfloat16, config=cfg)
+ def _prep_w(w_shuf, w_sc, n, kh):
+     return w_shuf.view(torch.uint8).reshape(n // 16, kh * 16), w_sc.view(torch.uint8)
   
   
- def _run_s5(A, B_shuffle, B_scale_sh, c):
-     """Shape 5: M=64, N=7168, K=2048 — Quant + CK ASM l2ks=3"""
-     # Use preshuffle (fused quant+GEMM, no separate quant kernel)
-     K = 2048; sc = K // 32; sn = ((sc+7)//8)*8; K_half = K // 2; N = 7168
-     padN = B_scale_sh.view(torch.uint8).shape[0]
-     bs_r = B_scale_sh.view(torch.uint8).reshape(padN // 32, sn * 32)
-     b_r = B_shuffle.view(torch.uint8).reshape(N // 16, K_half * 16)
-     global _preshuffle
-     if _preshuffle is None:
-         from aiter.ops.triton.gemm.basic.gemm_a16wfp4 import gemm_a16wfp4_preshuffle
-         _preshuffle = gemm_a16wfp4_preshuffle
-     cfg = {"BLOCK_SIZE_M": 8, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 256, "GROUP_SIZE_M": 1, "NUM_KSPLIT": 1, "SPLITK_BLOCK_SIZE": K, "matrix_instr_nonkdim": 16, "num_warps": 4, "num_stages": 2, "waves_per_eu": 2, "cache_modifier": ".cg"}
-     return _preshuffle(A, b_r, bs_r, prequant=True, dtype=torch.bfloat16, config=cfg)
+ def _setup(m, n, k, dev):
+     key = (m, n, k)
+     if key not in _params:
+         c = _resolve(m, n, k)
+         nsk = c["NSK"]; kh = k // 2
+         y, pp = _get_alloc(m, n, nsk, dev)
+         grid = (nsk * triton.cdiv(m, c["BM"]) * triton.cdiv(n, c["BN"]),)
+         if nsk == 1:
+             sk, sm, sn = 0, y.stride(0), y.stride(1)
+         else:
+             sk, sm, sn = pp.stride(0), pp.stride(1), pp.stride(2)
+         p = dict(cfg=c, kh=kh, grid=grid, nsk=nsk, sk=sk, sm=sm, sn=sn)
+         if nsk > 1:
+             p["rgrid"] = (triton.cdiv(m, 16), triton.cdiv(n, 64))
+             p["actual_k"] = triton.cdiv(kh, (c["SK_BLOCK"] // 2))
+             p["max_k"] = triton.next_power_of_2(nsk)
+         _params[key] = p
+     return _params[key]
   
   
- def _run_s6(A, B_shuffle, B_scale_sh, c):
-     """Shape 6: M=256, N=3072, K=1536 — PRESHUFFLE (faster than CK ASM with lean quant!)"""
-     K = 1536; N = 3072; sc = K // 32; sn = ((sc+7)//8)*8; K_half = K // 2
-     padN = B_scale_sh.view(torch.uint8).shape[0]
-     bs_r = B_scale_sh.view(torch.uint8).reshape(padN // 32, sn * 32)
-     b_r = B_shuffle.view(torch.uint8).reshape(N // 16, K_half * 16)
-     global _preshuffle
-     if _preshuffle is None:
-         from aiter.ops.triton.gemm.basic.gemm_a16wfp4 import gemm_a16wfp4_preshuffle
-         _preshuffle = gemm_a16wfp4_preshuffle
-     cfg = {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 256, "GROUP_SIZE_M": 1, "NUM_KSPLIT": 1, "SPLITK_BLOCK_SIZE": K, "matrix_instr_nonkdim": 16, "num_warps": 4, "num_stages": 2, "waves_per_eu": 2, "cache_modifier": ".cg"}
-     return _preshuffle(A, b_r, bs_r, prequant=True, dtype=torch.bfloat16, config=cfg)
+ def _run_fused(A, W_shuf, W_sc, m, n, k):
+     p = _setup(m, n, k, A.device)
+     c = p["cfg"]
+     y, pp = _get_alloc(m, n, p["nsk"], A.device)
+     wb, wsc = _prep_w(W_shuf, W_sc, n, p["kh"])
   
+     _gemm_fused[p["grid"]](
+         A, wb, y if p["nsk"] == 1 else pp, wsc,
+         m, n, p["kh"],
+         A.stride(0), A.stride(1), wb.stride(0), wb.stride(1),
+         p["sk"], p["sm"], p["sn"], wsc.stride(0), wsc.stride(1),
+         BM=c["BM"], BN=c["BN"], BK=c["BK"], GSM=c["GSM"],
+         NSK=c["NSK"], SK_BLOCK=c["SK_BLOCK"],
+         num_warps=c["nw"], num_stages=c["ns"],
+         waves_per_eu=c["wpe"], matrix_instr_nonkdim=c["nkd"],
+         cache_modifier=c["cm"],
+     )
   
- # ═══════════════════════════════════════════════════════════════════
- # General fallback for non-LB shapes (test mode uses different shapes)
- # ═══════════════════════════════════════════════════════════════════
-  
- def _init_general(m, k, n, device):
-     """General init for arbitrary shapes — used only in test mode."""
-     QUANT = 32; scale_cols = (k + QUANT - 1) // QUANT
-     sn = ((scale_cols + 7) // 8) * 8; sn_div8_mul256 = (sn // 8) * 256
-     CU_COUNT = 256
-     if k <= 1024:
-         BLOCK_K = max(128, triton.next_power_of_2(k)); BLOCK_M = 16 if m <= 32 else 32; BLOCK_N = 64; NW = 4
-         grid = (triton.cdiv(m, BLOCK_M), triton.cdiv(n, BLOCK_N)); total_wgs = grid[0] * grid[1]
-         wpe = 2 if total_wgs > CU_COUNT else 1
-         return {'mode': 'fused', 'out': torch.empty(m, n, dtype=torch.bfloat16, device=device),
-                 'grid': grid, 'BM': BLOCK_M, 'BN': BLOCK_N, 'BK': BLOCK_K, 'sn_div8_mul256': sn_div8_mul256, 'NW': NW, 'wpe': wpe}
-     elif m <= 32:
-         BLOCK_K = 512; BLOCK_N = 64; BLOCK_M = 16 if m <= 16 else 32
-         m_tiles = triton.cdiv(m, BLOCK_M); n_tiles = triton.cdiv(n, BLOCK_N)
-         k_iters = k // BLOCK_K if k % BLOCK_K == 0 else triton.cdiv(k, BLOCK_K); SPLIT_K = min(k_iters, 16)
-         total_wgs = m_tiles * n_tiles * SPLIT_K; XCD_SWIZZLE = 8 if total_wgs >= 16 else 1
-         wpe = 2 if total_wgs > CU_COUNT else 1
-         out = torch.empty(m, n, dtype=torch.bfloat16, device=device)
-         scratch = torch.empty(SPLIT_K, m, n, dtype=torch.float32, device=device) if SPLIT_K > 1 else None
-         reduce_grid = (m, triton.cdiv(n, 128)) if SPLIT_K > 1 else None
-         nonkdim = 16 if m <= 16 else 32
-         return {'mode': 'splitk', 'out': out, 'scratch': scratch, 'BM': BLOCK_M, 'BN': BLOCK_N, 'BK': BLOCK_K,
-                 'SPLIT_K': SPLIT_K, 'XCD_SWIZZLE': XCD_SWIZZLE, 'wpe': wpe, 'grid_m': m_tiles, 'grid_n': n_tiles,
-                 'total_wgs': total_wgs, 'sn_div8_mul256': sn_div8_mul256, 'reduce_grid': reduce_grid, 'nonkdim': nonkdim}
-     else:
-         sm = ((m + 255) // 256) * 256
-         x_fp4 = torch.empty((m, k // 2), dtype=torch.uint8, device=device)
-         bs_shuffled = torch.empty(sm, sn, dtype=torch.uint8, device=device)
-         out = torch.empty(m, n, dtype=torch.bfloat16, device=device)
-         if m <= 64:
-             BSM, NUM_ITER, BSN, NW, NS = 4, 1, 128, 4, 1; l2ks, quant_wpe = 3, 2
-         else:
-             NUM_ITER, BSM, BSN, NW, NS = 2, 16, 64, 2, 2; l2ks, quant_wpe = 2, 0
-         grid = (triton.cdiv(m, BSM), triton.cdiv(k, BSN * NUM_ITER))
-         x_fp4_v = x_fp4.view(_FP4X2); bs_shuf_v = bs_shuffled.view(_E8M0)
-         return {
-             'mode': 'asm', 'x_fp4': x_fp4, 'bs_shuffled': bs_shuffled, 'out': out,
-             'x_fp4_v': x_fp4_v, 'bs_shuf_v': bs_shuf_v,
-             'sc': scale_cols, 'sn_div8_mul256': sn_div8_mul256,
-             'grid': grid, 'BSM': BSM, 'BSN': BSN, 'NW': NW, 'NS': NS, 'NI': NUM_ITER,
-             'l2ks': l2ks, 'quant_wpe': quant_wpe,
-         }
-  
-  
- def _run_general(A, B_q, B_shuffle, B_scale_sh, c, m, n, k):
-     """General dispatch for arbitrary shapes — test mode only."""
-     if c['mode'] == 'fused':
-         Bq = B_q.view(torch.uint8); Bs = B_scale_sh.view(torch.uint8)
-         _fused_quant_gemm_kernel[c['grid']](
-             A, Bq, Bs, c['out'], m, n, k,
-             A.stride(0), A.stride(1), Bq.stride(0), Bq.stride(1),
-             c['sn_div8_mul256'], c['out'].stride(0), c['out'].stride(1),
-             BLOCK_M=c['BM'], BLOCK_N=c['BN'], BLOCK_K=c['BK'],
-             num_warps=c['NW'], num_stages=1, waves_per_eu=c['wpe'],
+     if p["nsk"] > 1:
+         _sum_partials[p["rgrid"]](
+             pp, y, m, n,
+             pp.stride(0), pp.stride(1), pp.stride(2),
+             y.stride(0), y.stride(1),
+             16, 64, p["actual_k"], p["max_k"],
          )
-         return c['out']
-     elif c['mode'] == 'splitk':
-         Bq = B_q.view(torch.uint8); Bs = B_scale_sh.view(torch.uint8); SK = c['SPLIT_K']
-         if SK == 1:
-             _fused_splitk_gemm[(c['total_wgs'],)](
-                 A, Bq, Bs, c['out'], m, n, k,
-                 A.stride(0), A.stride(1), Bq.stride(0), Bq.stride(1),
-                 c['sn_div8_mul256'], 0, c['out'].stride(0), c['out'].stride(1),
-                 c['grid_m'], c['grid_n'],
-                 BLOCK_M=c['BM'], BLOCK_N=c['BN'], BLOCK_K=c['BK'],
-                 SPLIT_K=1, XCD_SWIZZLE=c['XCD_SWIZZLE'],
-                 matrix_instr_nonkdim=c['nonkdim'],
-                 num_warps=4, num_stages=1, waves_per_eu=c['wpe'],
-             )
-         else:
-             s = c['scratch']
-             _fused_splitk_gemm[(c['total_wgs'],)](
-                 A, Bq, Bs, s, m, n, k,
-                 A.stride(0), A.stride(1), Bq.stride(0), Bq.stride(1),
-                 c['sn_div8_mul256'], s.stride(0), s.stride(1), s.stride(2),
-                 c['grid_m'], c['grid_n'],
-                 BLOCK_M=c['BM'], BLOCK_N=c['BN'], BLOCK_K=c['BK'],
-                 SPLIT_K=SK, XCD_SWIZZLE=c['XCD_SWIZZLE'],
-                 matrix_instr_nonkdim=c['nonkdim'],
-                 num_warps=4, num_stages=1, waves_per_eu=c['wpe'],
-             )
-             _reduce_splitk[c['reduce_grid']](
-                 s, c['out'], m, n,
-                 s.stride(0), s.stride(1), s.stride(2),
-                 c['out'].stride(0), c['out'].stride(1),
-                 SPLIT_K=SK, BLOCK_N=128, num_warps=4,
-             )
-         return c['out']
-     else:  # asm
-         _fused_quant_shuffle_kernel[c['grid']](
-             A, c['x_fp4'], c['bs_shuffled'],
-             A.stride(0), A.stride(1), c['x_fp4'].stride(0), c['x_fp4'].stride(1),
-             m, k, c['sn_div8_mul256'], c['sc'],
-             BLOCK_SIZE_M=c['BSM'], BLOCK_SIZE_N=c['BSN'],
-             NUM_ITER=c['NI'], NUM_STAGES=c['NS'],
-             num_warps=c['NW'], waves_per_eu=c['quant_wpe'], num_stages=1,
-         )
-         aiter.gemm_a4w4_asm(
-             c['x_fp4_v'], B_shuffle, c['bs_shuf_v'], B_scale_sh,
-             c['out'], _KNL_32x128, bpreshuffle=True, log2_k_split=c['l2ks'],
-         )
-         return c['out']
+     return y
   
   
- _gen_cache = {}
-  
-  
- # ═══════════════════════════════════════════════════════════════════
- # Main entry — hardcoded (M,N) dispatch for LB, general fallback for test
- # ═══════════════════════════════════════════════════════════════════
-  
  def custom_kernel(data: input_t) -> output_t:
-     A, B, B_q, B_shuffle, B_scale_sh = data
-     m, k = A.shape; n = B_q.shape[0]
-  
-     # Hardcoded 6-shape dispatch on (M, N) — unique for all 6 LB shapes
-     if m == 4 and n == 2880:
-         # Shape 1: (4, 2880, 512)
-         return _run_s1(A, B_shuffle, B_scale_sh)
-  
-     elif m == 16 and n == 2112:
-         # Shape 2: (16, 2112, 7168)
-         if 's2' not in _s: _s['s2'] = _init_shape2(A.device)
-         return _run_s2(A, B_q, B_scale_sh, _s['s2'])
-  
-     elif m == 32 and n == 4096:
-         # Shape 3: (32, 4096, 512) — preshuffle BM=8 NW=8
-         return _run_s3(A, B_shuffle, B_scale_sh)
-  
-     elif m == 32 and n == 2880:
-         # Shape 4: (32, 2880, 512) — preshuffle with fast_math patch
-         return _run_s4(A, B_shuffle, B_scale_sh)
-  
-     elif m == 64 and n == 7168:
-         # Shape 5: (64, 7168, 2048)
-         if 's5' not in _s: _s['s5'] = _init_shape5(A.device)
-         return _run_s5(A, B_shuffle, B_scale_sh, _s['s5'])
-  
-     elif m == 256 and n == 3072:
-         # Shape 6: (256, 3072, 1536)
-         if 's6' not in _s: _s['s6'] = _init_shape6(A.device)
-         return _run_s6(A, B_shuffle, B_scale_sh, _s['s6'])
-  
-     else:
-         # General fallback for test mode / unknown shapes
-         # Try preshuffle for small M K<=1024
-         if m <= 16 and k <= 1024:
-             global _preshuffle
-             if _preshuffle is None:
-                 from aiter.ops.triton.gemm.basic.gemm_a16wfp4 import gemm_a16wfp4_preshuffle
-                 _preshuffle = gemm_a16wfp4_preshuffle
-             try:
-                 sc = (k + 31) // 32; sn = ((sc + 7) // 8) * 8; K_half = k // 2
-                 padN = B_scale_sh.view(torch.uint8).shape[0]
-                 bs_reshaped = B_scale_sh.view(torch.uint8).reshape(padN // 32, sn * 32)
-                 b_shuf_reshaped = B_shuffle.view(torch.uint8).reshape(n // 16, K_half * 16)
-                 BM = 4 if m <= 8 else 8; NW = 4 if m <= 8 else 8
-                 config = {'BLOCK_SIZE_M': BM, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 1, 'NUM_KSPLIT': 1, 'SPLITK_BLOCK_SIZE': k, 'matrix_instr_nonkdim': 16, 'num_warps': NW, 'num_stages': 2, 'waves_per_eu': 2, 'cache_modifier': '.cg'}
-                 return _preshuffle(A, b_shuf_reshaped, bs_reshaped, prequant=True, dtype=torch.bfloat16, config=config)
-             except Exception:
-                 pass
-         key = (m, k, n)
-         if key not in _gen_cache:
-             _gen_cache[key] = _init_general(m, k, n, A.device)
-         return _run_general(A, B_q, B_shuffle, B_scale_sh, _gen_cache[key], m, n, k)
+     A = data[0]
+     return _run_fused(A, data[3], data[4], A.shape[0], data[1].shape[0], A.shape[1])