submission 745428

- #!POPCORN leaderboard amd-mxfp4-mm
- #!POPCORN gpu MI355X
-  
  """
- v211: M<=32 K<=1024 cache_modifier=None (from .cg).
-  
- For K=512 BSM=8 BSN=128 BSK=256, B data per block is 32KB FP4.
- Without .cg, L1 caching improves latency for 2 K-iterations.
- AMD library default uses null for this config.
+ V304: keep V300 logic, but narrow the `M=64,K=2048` fused config toward the
+ more conservative `storewt` occupancy hint.
  """
+ from task import input_t, output_t
  import torch
  import triton
  import triton.language as tl
- from aiter import dtypes
- from aiter.ops.triton._triton_kernels.quant.quant import _mxfp4_quant_op
- from aiter.ops.gemm_op_a4w4 import gemm_a4w4_asm
- from aiter.ops.triton._triton_kernels.gemm.basic.gemm_a16wfp4 import (
-     _gemm_a16wfp4_preshuffle_kernel,
- )
- from aiter.ops.triton.gluon.gemm_afp4wfp4 import (
-     _gemm_afp4wfp4_reduce_kernel as _gluon_reduce_kernel,
- )
- from aiter.ops.triton._triton_kernels.gemm.basic.gemm_afp4wfp4 import (
-     _gemm_afp4wfp4_reduce_kernel,
- )
- from aiter.ops.triton.gemm.basic.gemm_afp4wfp4 import get_splitk
   
- from task import input_t, output_t
   
- # Pre-allocated buffers keyed by (M, K, N)
- _buffers = {}
+ @triton.jit
+ def _mxfp4_quant_in_reg(
+     x_bf16,
+     BLOCK_SIZE_M: tl.constexpr,
+     BLOCK_SIZE_K: tl.constexpr,
+ ):
+     """Quantize BF16 block to MXFP4 using HW v_cvt_scalef32_pk_fp4_bf16."""
+     MXFP4_QUANT_BLOCK_SIZE: tl.constexpr = 32
+     NUM_QUANT_BLOCKS: tl.constexpr = BLOCK_SIZE_K // MXFP4_QUANT_BLOCK_SIZE
   
- # ASM kernel name — 32x128 is optimal for all small-M shapes per tuned CSV analysis
- _ASM_KERNEL_32x128 = "_ZN5aiter41f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128E"
+     # Compute scales from FP32 values
+     x_fp32 = x_bf16.to(tl.float32).reshape(BLOCK_SIZE_M, NUM_QUANT_BLOCKS, MXFP4_QUANT_BLOCK_SIZE)
   
- # Threshold: use fused for M <= this value
- _FUSED_M_THRESHOLD = 64
+     amax = tl.max(tl.abs(x_fp32), axis=-1, keep_dims=True)
+     amax = amax.to(tl.int32, bitcast=True)
+     amax = (amax + 0x200000).to(tl.uint32, bitcast=True) & 0xFF800000
+     log2_amax = ((amax >> 23) & 0xFF).to(tl.int32) - 127
+     scale_e8m0_unbiased_i = log2_amax - 2
+     scale_e8m0_unbiased_i = tl.minimum(tl.maximum(scale_e8m0_unbiased_i, -127), 127)
+     bs_e8m0 = scale_e8m0_unbiased_i.to(tl.uint8) + 127
   
+     # HW instruction divides by scale: fp4 = convert(bf16 / hw_scale)
+     # hw_scale = 2^unbiased (reciprocal of SW quant_scale which is 2^(-unbiased))
+     hw_scale_bits = (scale_e8m0_unbiased_i.to(tl.int32) + 127).to(tl.uint32) << 23
+     hw_scale = hw_scale_bits.to(tl.float32, bitcast=True)  # [M, NUM_QB, 1]
   
- def _get_fused_config(M, N, K):
-     """Get shape-specific config for fused quant+GEMM path.
-     All configs use BSK=256 num_stages=2 for Triton software pipelining.
-     """
-     if K > 4096:
-         # Custom split-K=7 BSK=256 for large-K shapes (e.g., 16x2112x7168)
-         # BSM=8: 238 blocks (0.93 waves) vs BSM=16: 119 blocks (0.46 waves)
-         # waves_per_eu=2: tuned JSON uses this for M>=16 shapes
-         return {
-             "BLOCK_SIZE_M": 8,
-             "BLOCK_SIZE_N": 128,
-             "BLOCK_SIZE_K": 256,
-             "GROUP_SIZE_M": 1,
-             "num_warps": 4,
-             "num_stages": 2,
-             "waves_per_eu": 2,
-             "matrix_instr_nonkdim": 16,
-             "cache_modifier": ".cg",
-             "NUM_KSPLIT": 7,
-         }
-     if M <= 4:
-         return {
-             "BLOCK_SIZE_M": 4,
-             "BLOCK_SIZE_N": 128,
-             "BLOCK_SIZE_K": 256,
-             "GROUP_SIZE_M": 1,
-             "num_warps": 4,
-             "num_stages": 2,
-             "waves_per_eu": 0,
-             "matrix_instr_nonkdim": 16,
-             "cache_modifier": ".cg",
-             "NUM_KSPLIT": 1,
-         }
-     elif M <= 8:
-         return {
-             "BLOCK_SIZE_M": 8,
-             "BLOCK_SIZE_N": 128,
-             "BLOCK_SIZE_K": 256,
-             "GROUP_SIZE_M": 1,
-             "num_warps": 4,
-             "num_stages": 2,
-             "waves_per_eu": 0,
-             "matrix_instr_nonkdim": 16,
-             "cache_modifier": ".cg",
-             "NUM_KSPLIT": 1,
-         }
-     elif M <= 32 and K <= 1024:
-         return {
-             "BLOCK_SIZE_M": 8,
-             "BLOCK_SIZE_N": 128,
-             "BLOCK_SIZE_K": 256,
-             "GROUP_SIZE_M": 1,
-             "num_warps": 4,
-             "num_stages": 2,
-             "waves_per_eu": 2,
-             "matrix_instr_nonkdim": 16,
-             "cache_modifier": None,
-             "NUM_KSPLIT": 1,
-         }
-     elif M <= 32:
-         return {
-             "BLOCK_SIZE_M": 32,
-             "BLOCK_SIZE_N": 64,
-             "BLOCK_SIZE_K": 512,
-             "GROUP_SIZE_M": 1,
-             "num_warps": 8,
-             "num_stages": 1,
-             "waves_per_eu": 2,
-             "matrix_instr_nonkdim": 16,
-             "cache_modifier": None,
-             "NUM_KSPLIT": 1,
-         }
-     else:
-         # M=64 (64x7168x2048): BSM=16 BSN=128 BSK=256 NW=4 NS=2
-         # 4*56=224 blocks, 8 K-iters with pipelining
-         # waves_per_eu=2: hint for higher occupancy per EU
-         return {
-             "BLOCK_SIZE_M": 16,
-             "BLOCK_SIZE_N": 128,
-             "BLOCK_SIZE_K": 256,
-             "GROUP_SIZE_M": 1,
-             "num_warps": 4,
-             "num_stages": 2,
-             "waves_per_eu": 2,
-             "matrix_instr_nonkdim": 16,
-             "cache_modifier": ".cg",
-             "NUM_KSPLIT": 1,
-         }
+     # Broadcast scale to per-pair granularity
+     hw_scale_flat = tl.broadcast_to(hw_scale, (BLOCK_SIZE_M, NUM_QUANT_BLOCKS, MXFP4_QUANT_BLOCK_SIZE))
+     hw_scale_flat = hw_scale_flat.reshape(BLOCK_SIZE_M, BLOCK_SIZE_K)
+     # Take scale for even element of each pair (both share same scale within 32-group)
+     hw_scale_pairs = hw_scale_flat.reshape(BLOCK_SIZE_M, BLOCK_SIZE_K // 2, 2)
+     hw_scale_even, _ = tl.split(hw_scale_pairs)
+     hw_scale_pair = hw_scale_even.reshape(BLOCK_SIZE_M, BLOCK_SIZE_K // 2)
   
+     # Pack BF16 pairs into uint32 for HW instruction
+     x_u16 = x_bf16.to(tl.uint16, bitcast=True).reshape(BLOCK_SIZE_M, BLOCK_SIZE_K // 2, 2)
+     lo_u16, hi_u16 = tl.split(x_u16)
+     x_u32 = lo_u16.to(tl.uint32) | (hi_u16.to(tl.uint32) << 16)
+     x_u32 = x_u32.reshape(BLOCK_SIZE_M, BLOCK_SIZE_K // 2)
   
+     # HW FP4 conversion
+     fp4_u32 = tl.inline_asm_elementwise(
+         "v_cvt_scalef32_pk_fp4_bf16 $0, $1, $2",
+         "=v, v, v",
+         [x_u32, hw_scale_pair],
+         dtype=tl.uint32,
+         is_pure=True,
+         pack=1,
+     )
+     x_fp4 = (fp4_u32 & 0xFF).to(tl.uint8)
+     x_fp4 = x_fp4.reshape(BLOCK_SIZE_M, BLOCK_SIZE_K // 2)
+  
+     return x_fp4, bs_e8m0.reshape(BLOCK_SIZE_M, NUM_QUANT_BLOCKS)
+  
+  
+ @triton.jit
+ def _standalone_quant_kernel(
+     a_ptr, a_fp4_ptr, a_scale_ptr,
+     M, K,
+     stride_am, stride_ak,
+     stride_qm, stride_qk,
+     stride_sm, stride_sk,
+     BLOCK_SIZE_M: tl.constexpr,
+     BLOCK_SIZE_K: tl.constexpr,
+ ):
+     pid_m = tl.program_id(0)
+     pid_k = tl.program_id(1)
+     offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+     offs_k = pid_k * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
+     a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
+     m_mask = offs_m[:, None] < M
+     k_mask = offs_k[None, :] < K
+     a_bf16 = tl.load(a_ptrs, mask=m_mask & k_mask, other=0.0)
+     a_fp4, a_scales = _mxfp4_quant_in_reg(a_bf16, BLOCK_SIZE_M, BLOCK_SIZE_K)
+     HALF_K: tl.constexpr = BLOCK_SIZE_K // 2
+     offs_qk = pid_k * HALF_K + tl.arange(0, HALF_K)
+     q_ptrs = a_fp4_ptr + offs_m[:, None] * stride_qm + offs_qk[None, :] * stride_qk
+     tl.store(q_ptrs, a_fp4, mask=m_mask & (offs_qk[None, :] < (K // 2)))
+     SCALE_K: tl.constexpr = BLOCK_SIZE_K // 32
+     offs_sk = pid_k * SCALE_K + tl.arange(0, SCALE_K)
+     s_ptrs = a_scale_ptr + offs_m[:, None] * stride_sm + offs_sk[None, :] * stride_sk
+     tl.store(s_ptrs, a_scales, mask=m_mask & (offs_sk[None, :] < (K // 32)))
+  
+  
  @triton.heuristics(
      {
-         "EVEN_M_N": lambda args: args["M"] % args["BLOCK_SIZE_M"] == 0
-         and args["N"] % (args["BLOCK_SIZE_N"] * args["NUM_ITER"]) == 0,
+         "EVEN_K": lambda args: (args["K"] % (args["BLOCK_SIZE_K"] // 2) == 0)
+         and (args["SPLITK_BLOCK_SIZE"] % args["BLOCK_SIZE_K"] == 0)
+         and (args["K"] % (args["SPLITK_BLOCK_SIZE"] // 2) == 0),
      }
  )
  @triton.jit
- def _fused_mxfp4_quant_shuffle_kernel(
-     x_ptr,
-     x_fp4_ptr,
-     bs_ptr,
-     stride_x_m_in,
-     stride_x_n_in,
-     stride_x_fp4_m_in,
-     stride_x_fp4_n_in,
-     M,
-     N,
+ def _fused_quant_gemm_preshuffle_kernel(
+     a_ptr, b_ptr, c_ptr, b_scales_ptr,
+     M, N, K,
+     stride_am, stride_ak,
+     stride_bn, stride_bk,
+     stride_ck, stride_cm, stride_cn,
+     stride_bsn, stride_bsk,
      BLOCK_SIZE_M: tl.constexpr,
      BLOCK_SIZE_N: tl.constexpr,
-     NUM_ITER: tl.constexpr,
-     NUM_STAGES: tl.constexpr,
-     MXFP4_QUANT_BLOCK_SIZE: tl.constexpr,
-     EVEN_M_N: tl.constexpr,
-     SCALING_MODE: tl.constexpr,
-     SCALE_N_PAD: tl.constexpr,
+     BLOCK_SIZE_K: tl.constexpr,
+     GROUP_SIZE_M: tl.constexpr,
+     NUM_KSPLIT: tl.constexpr,
+     SPLITK_BLOCK_SIZE: tl.constexpr,
+     EVEN_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
-     stride_x_m = tl.cast(stride_x_m_in, tl.int64)
-     stride_x_n = tl.cast(stride_x_n_in, tl.int64)
-     stride_x_fp4_m = tl.cast(stride_x_fp4_m_in, tl.int64)
-     stride_x_fp4_n = tl.cast(stride_x_fp4_n_in, tl.int64)
+     tl.assume(stride_am > 0)
+     tl.assume(stride_ak > 0)
+     tl.assume(stride_bn > 0)
+     tl.assume(stride_bk > 0)
+     tl.assume(stride_cm > 0)
+     tl.assume(stride_cn > 0)
+     tl.assume(stride_bsn > 0)
+     tl.assume(stride_bsk > 0)
   
-     NUM_QUANT_BLOCKS: tl.constexpr = BLOCK_SIZE_N // MXFP4_QUANT_BLOCK_SIZE
+     SCALE_GROUP_SIZE: tl.constexpr = 32
+     num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+     num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
   
-     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
+     pid_unified = tl.program_id(axis=0)
+     pid_k = pid_unified % NUM_KSPLIT
+     pid = pid_unified // NUM_KSPLIT
   
-         if EVEN_M_N:
-             x = tl.load(x_ptr + x_offs, cache_modifier=".cg").to(tl.float32)
-         else:
-             x_mask = (x_offs_m < M)[:, None] & (x_offs_n < N)[None, :]
-             x = tl.load(x_ptr + x_offs, mask=x_mask, cache_modifier=".cg").to(
-                 tl.float32
-             )
+     if NUM_KSPLIT == 1:
+         num_pid_in_group = GROUP_SIZE_M * num_pid_n
+         group_id = pid // num_pid_in_group
+         first_pid_m = group_id * GROUP_SIZE_M
+         group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+         pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+         pid_n = (pid % num_pid_in_group) // group_size_m
+     else:
+         pid_m = pid // num_pid_n
+         pid_n = pid % num_pid_n
   
-         out_tensor, bs_e8m0 = _mxfp4_quant_op(
-             x, BLOCK_SIZE_N, BLOCK_SIZE_M, MXFP4_QUANT_BLOCK_SIZE
-         )
+     tl.assume(pid_m >= 0)
+     tl.assume(pid_n >= 0)
+     tl.assume(pid_k >= 0)
   
-         # Store fp4 output
-         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_x_fp4_m + out_offs_n[None, :] * stride_x_fp4_n
-         )
+     if (pid_k * SPLITK_BLOCK_SIZE // 2) < K:
+         num_k_iter = tl.cdiv(SPLITK_BLOCK_SIZE // 2, BLOCK_SIZE_K // 2)
   
-         if EVEN_M_N:
-             tl.store(x_fp4_ptr + out_offs, out_tensor, cache_modifier=".wt")
-         else:
-             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=".wt")
+         # A: BF16 [M, 2*K]
+         offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+         offs_ak = pid_k * SPLITK_BLOCK_SIZE + tl.arange(0, BLOCK_SIZE_K)
+         a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_ak[None, :] * stride_ak)
   
-         # Store scales with inline shuffle permutation
-         bs_offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
-         bs_offs_n = pid_n * NUM_QUANT_BLOCKS + tl.arange(0, NUM_QUANT_BLOCKS)
-         num_bs_cols = (N + MXFP4_QUANT_BLOCK_SIZE - 1) // MXFP4_QUANT_BLOCK_SIZE
+         # B: pre-shuffled MXFP4 [N//16, K_packed*16]
+         offs_k_shuffle_arr = tl.arange(0, (BLOCK_SIZE_K // 2) * 16)
+         offs_k_shuffle = pid_k * (SPLITK_BLOCK_SIZE // 2) * 16 + offs_k_shuffle_arr
+         offs_bn = (pid_n * (BLOCK_SIZE_N // 16) + tl.arange(0, BLOCK_SIZE_N // 16)) % (N // 16)
+         b_ptrs = b_ptr + (offs_bn[:, None] * stride_bn + offs_k_shuffle[None, :] * stride_bk)
   
-         bs_offs_0 = bs_offs_m[:, None] // 32
-         bs_offs_1 = bs_offs_m[:, None] % 32
-         bs_offs_2 = bs_offs_1 % 16
-         bs_offs_1 = bs_offs_1 // 16
-         bs_offs_3 = bs_offs_n[None, :] // 8
-         bs_offs_4 = bs_offs_n[None, :] % 8
-         bs_offs_5 = bs_offs_4 % 4
-         bs_offs_4 = bs_offs_4 // 4
-         bs_offs = (
-             bs_offs_1
-             + bs_offs_4 * 2
-             + bs_offs_2 * 2 * 2
-             + bs_offs_5 * 2 * 2 * 16
-             + bs_offs_3 * 2 * 2 * 16 * 4
-             + bs_offs_0 * 2 * 16 * SCALE_N_PAD
+         # B_scale: shuffled E8M0 [N_pad, K_scale_pad]
+         # Each group of 32 N values occupies 32 consecutive rows.
+         # Row index = pid_n * BLOCK_SIZE_N + group_offset * 32
+         offs_bsn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N // 32) * 32)
+         offs_ks = (pid_k * (SPLITK_BLOCK_SIZE // SCALE_GROUP_SIZE) * 32) + tl.arange(
+             0, BLOCK_SIZE_K // SCALE_GROUP_SIZE * 32
          )
+         b_scale_ptrs = (
+             b_scales_ptr + offs_bsn[:, None] * stride_bsn + offs_ks[None, :] * stride_bsk
+         )
   
-         bs_mask_valid = (bs_offs_m < M)[:, None] & (bs_offs_n < num_bs_cols)[None, :]
-         bs_e8m0 = tl.where(bs_mask_valid, bs_e8m0, 127)
+         accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
   
-         SCALE_M_PAD = (M + 255) // 256 * 256
-         bs_mask = (bs_offs_m < SCALE_M_PAD)[:, None] & (bs_offs_n < SCALE_N_PAD)[
-             None, :
-         ]
-         tl.store(
-             bs_ptr + bs_offs,
-             bs_e8m0.to(tl.uint8),
-             mask=bs_mask,
-             cache_modifier=".wt",
-         )
+         for k_iter in range(pid_k * num_k_iter, (pid_k + 1) * num_k_iter):
+             # Fire all loads first for better memory-level parallelism
+             if EVEN_K:
+                 a_bf16 = tl.load(a_ptrs)
+                 b_scales_raw = tl.load(b_scale_ptrs, cache_modifier=cache_modifier)
+                 b_raw = tl.load(b_ptrs, cache_modifier=cache_modifier)
+             else:
+                 k_offset = (k_iter - pid_k * num_k_iter) * BLOCK_SIZE_K
+                 a_bf16 = tl.load(
+                     a_ptrs,
+                     mask=tl.arange(0, BLOCK_SIZE_K)[None, :] < (2 * K - pid_k * SPLITK_BLOCK_SIZE - k_offset),
+                     other=0.0,
+                 )
+                 b_scales_raw = tl.load(b_scale_ptrs, cache_modifier=cache_modifier)
+                 b_raw = tl.load(
+                     b_ptrs,
+                     mask=offs_k_shuffle_arr[None, :] < ((K - (pid_k * (SPLITK_BLOCK_SIZE // 2) + (k_iter - pid_k * num_k_iter) * (BLOCK_SIZE_K // 2))) * 16),
+                     other=0,
+                     cache_modifier=cache_modifier,
+                 )
   
+             # Quantize A in registers
+             a_fp4, a_scales = _mxfp4_quant_in_reg(a_bf16, BLOCK_SIZE_M, BLOCK_SIZE_K)
   
- def _prepare_splitk_dispatch(M, N, K, config, device):
-     """Pre-compute all params for split-K direct dispatch (16x2112x7168)."""
-     K_kernel = K // 2
-     BSK = config["BLOCK_SIZE_K"]
-     NUM_KSPLIT = config["NUM_KSPLIT"]
+             # Unshuffle B scales
+             b_scales = (
+                 b_scales_raw
+                 .reshape(
+                     BLOCK_SIZE_N // 32,
+                     BLOCK_SIZE_K // SCALE_GROUP_SIZE // 8,
+                     4, 16, 2, 2, 1,
+                 )
+                 .permute(0, 5, 3, 1, 4, 2, 6)
+                 .reshape(BLOCK_SIZE_N, BLOCK_SIZE_K // SCALE_GROUP_SIZE)
+             )
   
-     SPLITK_BLOCK_SIZE, BSK, NUM_KSPLIT = get_splitk(K_kernel, BSK, NUM_KSPLIT)
+             # Unshuffle B data
+             b = (
+                 b_raw.reshape(1, BLOCK_SIZE_N // 16, BLOCK_SIZE_K // 64, 2, 16, 16)
+                 .permute(0, 1, 4, 2, 3, 5)
+                 .reshape(BLOCK_SIZE_N, BLOCK_SIZE_K // 2)
+                 .trans(1, 0)
+             )
   
-     BSN = max(config["BLOCK_SIZE_N"], 32)
-     BSM = config["BLOCK_SIZE_M"]
+             accumulator = tl.dot_scaled(
+                 a_fp4, a_scales, "e2m1", b, b_scales, "e2m1", accumulator
+             )
   
-     grid_size = NUM_KSPLIT * triton.cdiv(M, BSM) * triton.cdiv(N, BSN)
+             a_ptrs += BLOCK_SIZE_K * stride_ak
+             b_ptrs += (BLOCK_SIZE_K // 2) * 16 * stride_bk
+             b_scale_ptrs += BLOCK_SIZE_K * stride_bsk
   
-     # Pre-allocate y_pp
-     y_pp = torch.empty((NUM_KSPLIT, M, N), dtype=torch.float32, device=device)
+         c = accumulator.to(c_ptr.type.element_ty)
   
-     # Reduce kernel params — gluon version uses BSN=64 for fp32 partials
-     REDUCE_BSM = 16
-     REDUCE_BSN = 64  # Gluon default for fp32 partials
-     ACTUAL_KSPLIT = triton.cdiv(K_kernel, (SPLITK_BLOCK_SIZE // 2))
-     reduce_grid = (triton.cdiv(M, REDUCE_BSM), triton.cdiv(N, REDUCE_BSN))
+         offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
+         offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)
+         c_ptrs = (
+             c_ptr
+             + stride_cm * offs_cm[:, None]
+             + stride_cn * offs_cn[None, :]
+             + pid_k * stride_ck
+         )
+         c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
+         tl.store(c_ptrs, c, mask=c_mask)
   
-     return {
-         'BLOCK_SIZE_M': BSM,
-         'BLOCK_SIZE_N': BSN,
-         'BLOCK_SIZE_K': BSK,
-         'GROUP_SIZE_M': config["GROUP_SIZE_M"],
-         'NUM_KSPLIT': NUM_KSPLIT,
-         'SPLITK_BLOCK_SIZE': SPLITK_BLOCK_SIZE,
-         'num_warps': config["num_warps"],
-         'num_stages': config["num_stages"],
-         'waves_per_eu': config["waves_per_eu"],
-         'matrix_instr_nonkdim': config["matrix_instr_nonkdim"],
-         'cache_modifier': config["cache_modifier"],
-         'grid_size': grid_size,
-         'K_kernel': K_kernel,
-         'y_pp': y_pp,
-         'reduce_grid': reduce_grid,
-         'REDUCE_BSM': REDUCE_BSM,
-         'REDUCE_BSN': REDUCE_BSN,
-         'ACTUAL_KSPLIT': ACTUAL_KSPLIT,
-         'MAX_KSPLIT': triton.next_power_of_2(NUM_KSPLIT),
+  
+ @triton.jit
+ def _reduce_kernel(
+     c_in_ptr, c_out_ptr, M, N,
+     stride_c_in_k, stride_c_in_m, stride_c_in_n,
+     stride_c_out_m, stride_c_out_n,
+     BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
+     ACTUAL_KSPLIT: tl.constexpr, MAX_KSPLIT: tl.constexpr,
+ ):
+     pid_m = tl.program_id(axis=0)
+     pid_n = tl.program_id(axis=1)
+     offs_m = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+     offs_n = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
+  
+     # Sequential accumulation: load one partial at a time (fewer registers)
+     base_ptrs = (
+         c_in_ptr
+         + (offs_m[:, None] * stride_c_in_m)
+         + (offs_n[None, :] * stride_c_in_n)
+     )
+     acc = tl.load(base_ptrs).to(tl.float32)
+     for ks in tl.static_range(1, MAX_KSPLIT):
+         if ks < ACTUAL_KSPLIT:
+             acc += tl.load(base_ptrs + ks * stride_c_in_k).to(tl.float32)
+  
+     c = acc.to(c_out_ptr.type.element_ty)
+     c_out_ptrs = (
+         c_out_ptr
+         + (offs_m[:, None] * stride_c_out_m)
+         + (offs_n[None, :] * stride_c_out_n)
+     )
+     tl.store(c_out_ptrs, c)
+  
+  
+ @triton.heuristics(
+     {
+         "EVEN_K": lambda args: (args["K"] % (args["BLOCK_SIZE_K"] // 2) == 0)
+         and (args["SPLITK_BLOCK_SIZE"] % args["BLOCK_SIZE_K"] == 0)
+         and (args["K"] % (args["SPLITK_BLOCK_SIZE"] // 2) == 0),
      }
+ )
+ @triton.jit
+ def _gemm_only_preshuffle_kernel(
+     a_fp4_ptr, a_scale_ptr, b_ptr, c_ptr, b_scales_ptr,
+     M, N, K,
+     stride_qm, stride_qk,
+     stride_sm, stride_sk,
+     stride_bn, stride_bk,
+     stride_ck, stride_cm, stride_cn,
+     stride_bsn, stride_bsk,
+     BLOCK_SIZE_M: tl.constexpr,
+     BLOCK_SIZE_N: tl.constexpr,
+     BLOCK_SIZE_K: tl.constexpr,
+     GROUP_SIZE_M: tl.constexpr,
+     NUM_KSPLIT: tl.constexpr,
+     SPLITK_BLOCK_SIZE: tl.constexpr,
+     EVEN_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,
+ ):
+     tl.assume(stride_qm > 0)
+     tl.assume(stride_qk > 0)
+     tl.assume(stride_sm > 0)
+     tl.assume(stride_sk > 0)
+     tl.assume(stride_bn > 0)
+     tl.assume(stride_bk > 0)
+     tl.assume(stride_cm > 0)
+     tl.assume(stride_cn > 0)
+     tl.assume(stride_bsn > 0)
+     tl.assume(stride_bsk > 0)
   
+     SCALE_GROUP_SIZE: tl.constexpr = 32
+     HALF_BK: tl.constexpr = BLOCK_SIZE_K // 2
+     SCALE_BK: tl.constexpr = BLOCK_SIZE_K // SCALE_GROUP_SIZE
+     num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+     num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
   
- def _get_or_create_buffers(M, K, N, device):
-     """Get pre-allocated buffers for given shape."""
-     key = (M, K, N)
-     if key not in _buffers:
-         if M <= _FUSED_M_THRESHOLD:
-             config = _get_fused_config(M, N, K)
-             if config["NUM_KSPLIT"] > 1:
-                 # Split-K path: use direct dispatch with tuned reduce kernel
-                 splitk_params = _prepare_splitk_dispatch(M, N, K, config, device)
-                 _buffers[key] = {
-                     'mode': 'fused_splitk',
-                     'out': torch.empty((M, N), dtype=torch.bfloat16, device=device),
-                     'B_w': None,
-                     'B_sc': None,
-                     'splitk_params': splitk_params,
-                 }
-             else:
-                 # Non-split-K: direct dispatch (bypass wrapper overhead)
-                 K_kernel = K // 2
-                 BSK = config["BLOCK_SIZE_K"]
-                 BSN = max(config["BLOCK_SIZE_N"], 32)
-                 BSM = config["BLOCK_SIZE_M"]
-                 SPLITK_BLOCK_SIZE = 2 * K_kernel  # No split-K
+     pid_unified = tl.program_id(axis=0)
+     pid_k = pid_unified % NUM_KSPLIT
+     pid = pid_unified // NUM_KSPLIT
   
-                 grid_size = triton.cdiv(M, BSM) * triton.cdiv(N, BSN)
+     if NUM_KSPLIT == 1:
+         num_pid_in_group = GROUP_SIZE_M * num_pid_n
+         group_id = pid // num_pid_in_group
+         first_pid_m = group_id * GROUP_SIZE_M
+         group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+         pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+         pid_n = (pid % num_pid_in_group) // group_size_m
+     else:
+         pid_m = pid // num_pid_n
+         pid_n = pid % num_pid_n
   
-                 _buffers[key] = {
-                     'mode': 'fused_direct',
-                     'out': torch.empty((M, N), dtype=torch.bfloat16, device=device),
-                     'B_w': None,
-                     'B_sc': None,
-                     'grid_size': grid_size,
-                     'K_kernel': K_kernel,
-                     'BLOCK_SIZE_M': BSM,
-                     'BLOCK_SIZE_N': BSN,
-                     'BLOCK_SIZE_K': BSK,
-                     'SPLITK_BLOCK_SIZE': SPLITK_BLOCK_SIZE,
-                     'GROUP_SIZE_M': config["GROUP_SIZE_M"],
-                     'NUM_KSPLIT': 1,
-                     'num_warps': config["num_warps"],
-                     'num_stages': config["num_stages"],
-                     'waves_per_eu': config["waves_per_eu"],
-                     'matrix_instr_nonkdim': config["matrix_instr_nonkdim"],
-                     'cache_modifier': config["cache_modifier"],
-                 }
-         else:
-             MXFP4_QUANT_BLOCK_SIZE = 32
-             SCALE_N_valid = triton.cdiv(K, MXFP4_QUANT_BLOCK_SIZE)
-             SCALE_M = triton.cdiv(M, 256) * 256
-             SCALE_N = triton.cdiv(SCALE_N_valid, 8) * 8
+     tl.assume(pid_m >= 0)
+     tl.assume(pid_n >= 0)
+     tl.assume(pid_k >= 0)
   
-             NUM_ITER = 1
-             BLOCK_SIZE_M = min(32, triton.next_power_of_2(M))
-             BLOCK_SIZE_N = 64
-             NUM_WARPS = 2
-             NUM_STAGES = 1
+     if (pid_k * SPLITK_BLOCK_SIZE // 2) < K:
+         num_k_iter = tl.cdiv(SPLITK_BLOCK_SIZE // 2, HALF_BK)
   
-             BLOCK_SIZE_M = triton.cdiv(BLOCK_SIZE_M, 32) * 32
-             BLOCK_SIZE_N = triton.cdiv(BLOCK_SIZE_N, 32) * 32
+         offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+         offs_aqk = pid_k * (SPLITK_BLOCK_SIZE // 2) + tl.arange(0, HALF_BK)
+         a_fp4_ptrs = a_fp4_ptr + (offs_am[:, None] * stride_qm + offs_aqk[None, :] * stride_qk)
   
-             grid = (
-                 triton.cdiv(M, BLOCK_SIZE_M),
-                 triton.cdiv(K, BLOCK_SIZE_N * NUM_ITER),
-             )
+         offs_ask = pid_k * (SPLITK_BLOCK_SIZE // SCALE_GROUP_SIZE) + tl.arange(0, SCALE_BK)
+         a_scale_ptrs = a_scale_ptr + (offs_am[:, None] * stride_sm + offs_ask[None, :] * stride_sk)
   
-             padded_M = (M + 31) // 32 * 32
+         offs_k_shuffle_arr = tl.arange(0, HALF_BK * 16)
+         offs_k_shuffle = pid_k * (SPLITK_BLOCK_SIZE // 2) * 16 + offs_k_shuffle_arr
+         offs_bn = (pid_n * (BLOCK_SIZE_N // 16) + tl.arange(0, BLOCK_SIZE_N // 16)) % (N // 16)
+         b_ptrs = b_ptr + (offs_bn[:, None] * stride_bn + offs_k_shuffle[None, :] * stride_bk)
   
-             _buffers[key] = {
-                 'mode': 'two_phase',
-                 'x_fp4': torch.empty((M, K // 2), dtype=torch.uint8, device=device),
-                 'blockscale': torch.empty((SCALE_M, SCALE_N), dtype=torch.uint8, device=device),
-                 'gemm_out': torch.empty((padded_M, N), dtype=torch.bfloat16, device=device),
-                 'SCALE_N': SCALE_N,
-                 'BLOCK_SIZE_M': BLOCK_SIZE_M,
-                 'BLOCK_SIZE_N': BLOCK_SIZE_N,
-                 'NUM_ITER': NUM_ITER,
-                 'NUM_STAGES': NUM_STAGES,
-                 'NUM_WARPS': NUM_WARPS,
-                 'grid': grid,
-                 'M': M,
-             }
-     return _buffers[key]
+         offs_bsn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N // 32) * 32)
+         offs_ks = (pid_k * (SPLITK_BLOCK_SIZE // SCALE_GROUP_SIZE) * 32) + tl.arange(
+             0, SCALE_BK * 32
+         )
+         b_scale_ptrs = (
+             b_scales_ptr + offs_bsn[:, None] * stride_bsn + offs_ks[None, :] * stride_bsk
+         )
   
+         accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
   
- def custom_kernel(data: input_t) -> output_t:
-     A, _, _, B_shuffle, B_scale_sh = data
-     M, K = A.shape
-     N = B_shuffle.shape[0]
+         for k_iter in range(pid_k * num_k_iter, (pid_k + 1) * num_k_iter):
+             if EVEN_K:
+                 a_fp4 = tl.load(a_fp4_ptrs, cache_modifier=cache_modifier)
+                 a_scales = tl.load(a_scale_ptrs, cache_modifier=cache_modifier)
+             else:
+                 k_off = (k_iter - pid_k * num_k_iter) * HALF_BK
+                 k_remain = K - (pid_k * (SPLITK_BLOCK_SIZE // 2) + k_off)
+                 a_fp4 = tl.load(
+                     a_fp4_ptrs,
+                     mask=tl.arange(0, HALF_BK)[None, :] < k_remain,
+                     other=0,
+                     cache_modifier=cache_modifier,
+                 )
+                 s_remain = (2 * K) // SCALE_GROUP_SIZE - (pid_k * (SPLITK_BLOCK_SIZE // SCALE_GROUP_SIZE) + (k_iter - pid_k * num_k_iter) * SCALE_BK)
+                 a_scales = tl.load(
+                     a_scale_ptrs,
+                     mask=tl.arange(0, SCALE_BK)[None, :] < s_remain,
+                     other=0,
+                     cache_modifier=cache_modifier,
+                 )
   
-     buf = _get_or_create_buffers(M, K, N, A.device)
+             b_scales = (
+                 tl.load(b_scale_ptrs, cache_modifier=cache_modifier)
+                 .reshape(
+                     BLOCK_SIZE_N // 32,
+                     SCALE_BK // 8,
+                     4, 16, 2, 2, 1,
+                 )
+                 .permute(0, 5, 3, 1, 4, 2, 6)
+                 .reshape(BLOCK_SIZE_N, SCALE_BK)
+             )
   
-     if buf['mode'] == 'fused_splitk':
-         # Split-K path with tuned reduce kernel (REDUCE_BSN=16)
-         b_ptr = B_shuffle.data_ptr()
-         if buf['B_w'] is None or buf.get('_b_ptr') != b_ptr:
-             buf['B_w'] = B_shuffle.view(torch.uint8).reshape(N // 16, (K // 2) * 16)
-             bs_shape = B_scale_sh.shape
-             buf['B_sc'] = B_scale_sh.view(torch.uint8).reshape(
-                 bs_shape[0] // 32, bs_shape[1] * 32
+             if EVEN_K:
+                 b = tl.load(b_ptrs, cache_modifier=cache_modifier)
+             else:
+                 b = tl.load(
+                     b_ptrs,
+                     mask=offs_k_shuffle_arr[None, :] < ((K - (pid_k * (SPLITK_BLOCK_SIZE // 2) + (k_iter - pid_k * num_k_iter) * HALF_BK)) * 16),
+                     other=0,
+                     cache_modifier=cache_modifier,
+                 )
+  
+             b = (
+                 b.reshape(1, BLOCK_SIZE_N // 16, BLOCK_SIZE_K // 64, 2, 16, 16)
+                 .permute(0, 1, 4, 2, 3, 5)
+                 .reshape(BLOCK_SIZE_N, HALF_BK)
+                 .trans(1, 0)
              )
-             buf['_b_ptr'] = b_ptr
   
-         kp = buf['splitk_params']
-         out = buf['out']
-         y_pp = kp['y_pp']
+             accumulator = tl.dot_scaled(
+                 a_fp4, a_scales, "e2m1", b, b_scales, "e2m1", accumulator
+             )
   
-         _gemm_a16wfp4_preshuffle_kernel[(kp['grid_size'],)](
-             A,
-             buf['B_w'],
-             y_pp,
-             buf['B_sc'],
-             M,
-             N,
-             kp['K_kernel'],
-             A.stride(0),
-             A.stride(1),
-             buf['B_w'].stride(0),
-             buf['B_w'].stride(1),
-             y_pp.stride(0),
-             y_pp.stride(1),
-             y_pp.stride(2),
-             buf['B_sc'].stride(0),
-             buf['B_sc'].stride(1),
-             BLOCK_SIZE_M=kp['BLOCK_SIZE_M'],
-             BLOCK_SIZE_N=kp['BLOCK_SIZE_N'],
-             BLOCK_SIZE_K=kp['BLOCK_SIZE_K'],
-             GROUP_SIZE_M=kp['GROUP_SIZE_M'],
-             NUM_KSPLIT=kp['NUM_KSPLIT'],
-             SPLITK_BLOCK_SIZE=kp['SPLITK_BLOCK_SIZE'],
-             num_warps=kp['num_warps'],
-             num_stages=kp['num_stages'],
-             waves_per_eu=kp['waves_per_eu'],
-             matrix_instr_nonkdim=kp['matrix_instr_nonkdim'],
-             PREQUANT=True,
-             cache_modifier=kp['cache_modifier'],
+             a_fp4_ptrs += HALF_BK * stride_qk
+             a_scale_ptrs += SCALE_BK * stride_sk
+             b_ptrs += HALF_BK * 16 * stride_bk
+             b_scale_ptrs += BLOCK_SIZE_K * stride_bsk
+  
+         c = accumulator.to(c_ptr.type.element_ty)
+  
+         offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
+         offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)
+         c_ptrs = (
+             c_ptr
+             + stride_cm * offs_cm[:, None]
+             + stride_cn * offs_cn[None, :]
+             + pid_k * stride_ck
          )
+         c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
+         tl.store(c_ptrs, c, mask=c_mask)
   
-         _gluon_reduce_kernel[kp['reduce_grid']](
-             y_pp,
-             out,
-             M,
-             N,
-             y_pp.stride(0),
-             y_pp.stride(1),
-             y_pp.stride(2),
-             out.stride(0),
-             out.stride(1),
-             kp['REDUCE_BSM'],
-             kp['REDUCE_BSN'],
-             kp['ACTUAL_KSPLIT'],
-             kp['MAX_KSPLIT'],
+  
+ def get_splitk(K, BLOCK_SIZE_K, NUM_KSPLIT):
+     SPLITK_BLOCK_SIZE = (
+         triton.cdiv((2 * triton.cdiv(K, NUM_KSPLIT)), BLOCK_SIZE_K) * BLOCK_SIZE_K
+     )
+     while NUM_KSPLIT > 1 and BLOCK_SIZE_K > 16:
+         if (
+             K % (SPLITK_BLOCK_SIZE // 2) == 0
+             and SPLITK_BLOCK_SIZE % BLOCK_SIZE_K == 0
+             and K % (BLOCK_SIZE_K // 2) == 0
+         ):
+             break
+         elif K % (SPLITK_BLOCK_SIZE // 2) != 0 and NUM_KSPLIT > 1:
+             NUM_KSPLIT = NUM_KSPLIT // 2
+         elif SPLITK_BLOCK_SIZE % BLOCK_SIZE_K != 0:
+             if NUM_KSPLIT > 1:
+                 NUM_KSPLIT = NUM_KSPLIT // 2
+             elif BLOCK_SIZE_K > 16:
+                 BLOCK_SIZE_K = BLOCK_SIZE_K // 2
+         elif K % (BLOCK_SIZE_K // 2) != 0 and BLOCK_SIZE_K > 16:
+             BLOCK_SIZE_K = BLOCK_SIZE_K // 2
+         else:
+             break
+         SPLITK_BLOCK_SIZE = (
+             triton.cdiv((2 * triton.cdiv(K, NUM_KSPLIT)), BLOCK_SIZE_K) * BLOCK_SIZE_K
          )
+     NUM_KSPLIT = triton.cdiv(K, (SPLITK_BLOCK_SIZE // 2))
+     return SPLITK_BLOCK_SIZE, BLOCK_SIZE_K, NUM_KSPLIT
   
-         return out
   
-     elif buf['mode'] == 'fused_direct':
-         # Non-split-K fused path: direct kernel dispatch (bypass wrapper)
-         b_ptr = B_shuffle.data_ptr()
-         if buf['B_w'] is None or buf.get('_b_ptr') != b_ptr:
-             buf['B_w'] = B_shuffle.view(torch.uint8).reshape(N // 16, (K // 2) * 16)
-             bs_shape = B_scale_sh.shape
-             buf['B_sc'] = B_scale_sh.view(torch.uint8).reshape(
-                 bs_shape[0] // 32, bs_shape[1] * 32
+ CONFIGS = {
+     (4, 2880, 512): {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 256, "GROUP_SIZE_M": 1, "num_warps": 2, "num_stages": 3, "waves_per_eu": 3, "matrix_instr_nonkdim": 16, "cache_modifier": ".cg", "NUM_KSPLIT": 1},
+     (16, 2112, 7168): {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 512, "GROUP_SIZE_M": 1, "num_warps": 4, "num_stages": 2, "waves_per_eu": 3, "matrix_instr_nonkdim": 16, "cache_modifier": ".cg", "NUM_KSPLIT": 7},
+     (32, 4096, 512): {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 256, "GROUP_SIZE_M": 1, "num_warps": 4, "num_stages": 3, "waves_per_eu": 3, "matrix_instr_nonkdim": 16, "cache_modifier": ".cg", "NUM_KSPLIT": 1},
+     (32, 2880, 512): {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 256, "GROUP_SIZE_M": 1, "num_warps": 4, "num_stages": 3, "waves_per_eu": 3, "matrix_instr_nonkdim": 16, "cache_modifier": ".cg", "NUM_KSPLIT": 1},
+     (64, 7168, 2048): {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 256, "GROUP_SIZE_M": 1, "num_warps": 4, "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, "cache_modifier": ".cg", "NUM_KSPLIT": 1},
+     (256, 3072, 1536): {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 512, "GROUP_SIZE_M": 1, "num_warps": 8, "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, "cache_modifier": None, "NUM_KSPLIT": 1},
+ }
+  
+ DEFAULT_CONFIG = {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 256, "GROUP_SIZE_M": 1, "num_warps": 2, "num_stages": 2, "waves_per_eu": 0, "matrix_instr_nonkdim": 16, "cache_modifier": ".cg", "NUM_KSPLIT": 1}
+  
+ GEMM_CONFIGS = {
+     (32, 4096, 512): {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 256, "GROUP_SIZE_M": 4, "num_warps": 4, "num_stages": 2, "waves_per_eu": 0, "matrix_instr_nonkdim": 16, "cache_modifier": ".cg", "NUM_KSPLIT": 1},
+     (32, 2880, 512): {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 256, "GROUP_SIZE_M": 4, "num_warps": 4, "num_stages": 2, "waves_per_eu": 0, "matrix_instr_nonkdim": 16, "cache_modifier": ".cg", "NUM_KSPLIT": 1},
+     (64, 7168, 2048): {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 512, "GROUP_SIZE_M": 4, "num_warps": 4, "num_stages": 2, "waves_per_eu": 0, "matrix_instr_nonkdim": 16, "cache_modifier": ".cg", "NUM_KSPLIT": 2},
+ }
+  
+ GEMM_DEFAULT_CONFIG = {"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 256, "GROUP_SIZE_M": 4, "num_warps": 4, "num_stages": 2, "waves_per_eu": 0, "matrix_instr_nonkdim": 16, "cache_modifier": ".cg", "NUM_KSPLIT": 1}
+  
+  
+ _buf_cache = {}
+ _config_cache = {}
+  
+ def _get_buffers(m, n, num_ksplit, device):
+     key = (m, n, num_ksplit)
+     if key not in _buf_cache:
+         y = torch.empty((m, n), dtype=torch.bfloat16, device=device)
+         y_pp = torch.empty((num_ksplit, m, n), dtype=torch.float32, device=device) if num_ksplit > 1 else None
+         _buf_cache[key] = (y, y_pp)
+     return _buf_cache[key]
+  
+ def _get_config(m, n, k):
+     key = (m, n, k)
+     if key not in _config_cache:
+         config = CONFIGS.get(key, DEFAULT_CONFIG).copy()
+         K_packed = k // 2
+         if config["NUM_KSPLIT"] > 1:
+             SPLITK_BLOCK_SIZE, BLOCK_SIZE_K, NUM_KSPLIT = get_splitk(
+                 K_packed, config["BLOCK_SIZE_K"], config["NUM_KSPLIT"]
              )
-             buf['_b_ptr'] = b_ptr
+             config["SPLITK_BLOCK_SIZE"] = SPLITK_BLOCK_SIZE
+             config["BLOCK_SIZE_K"] = BLOCK_SIZE_K
+             config["NUM_KSPLIT"] = NUM_KSPLIT
+         else:
+             config["SPLITK_BLOCK_SIZE"] = 2 * K_packed
+             config["NUM_KSPLIT"] = 1
+         if config["BLOCK_SIZE_K"] >= 2 * K_packed:
+             config["BLOCK_SIZE_K"] = triton.next_power_of_2(2 * K_packed)
+             config["SPLITK_BLOCK_SIZE"] = 2 * K_packed
+             config["NUM_KSPLIT"] = 1
+         config["BLOCK_SIZE_N"] = max(config["BLOCK_SIZE_N"], 32)
+         _config_cache[key] = config
+     return _config_cache[key]
   
-         out = buf['out']
   
-         _gemm_a16wfp4_preshuffle_kernel[(buf['grid_size'],)](
-             A,
-             buf['B_w'],
-             out,
-             buf['B_sc'],
-             M,
-             N,
-             buf['K_kernel'],
-             A.stride(0),
-             A.stride(1),
-             buf['B_w'].stride(0),
-             buf['B_w'].stride(1),
-             0,  # stride_ck (no split-K)
-             out.stride(0),
-             out.stride(1),
-             buf['B_sc'].stride(0),
-             buf['B_sc'].stride(1),
-             BLOCK_SIZE_M=buf['BLOCK_SIZE_M'],
-             BLOCK_SIZE_N=buf['BLOCK_SIZE_N'],
-             BLOCK_SIZE_K=buf['BLOCK_SIZE_K'],
-             GROUP_SIZE_M=buf['GROUP_SIZE_M'],
-             NUM_KSPLIT=buf['NUM_KSPLIT'],
-             SPLITK_BLOCK_SIZE=buf['SPLITK_BLOCK_SIZE'],
-             num_warps=buf['num_warps'],
-             num_stages=buf['num_stages'],
-             waves_per_eu=buf['waves_per_eu'],
-             matrix_instr_nonkdim=buf['matrix_instr_nonkdim'],
-             PREQUANT=True,
-             cache_modifier=buf['cache_modifier'],
+ _launch_cache = {}
+  
+  
+ def _prepare_b_views(B_shuffle, B_scale_sh, n, k_packed):
+     b_reshaped = B_shuffle.view(torch.uint8).reshape(n // 16, k_packed * 16)
+     b_scale_uint8 = B_scale_sh.view(torch.uint8)
+     return b_reshaped, b_scale_uint8
+  
+  
+ def _build_launch_params(m, n, k, device):
+     """Precompute ALL launch parameters once per shape."""
+     config = _get_config(m, n, k)
+     K_packed = k // 2
+     ks = config["NUM_KSPLIT"]
+  
+     y, y_pp = _get_buffers(m, n, ks, device)
+  
+     grid = (ks * triton.cdiv(m, config["BLOCK_SIZE_M"]) * triton.cdiv(n, config["BLOCK_SIZE_N"]),)
+  
+     # Pre-store strides for y/y_pp
+     if ks == 1:
+         c_stride_k, c_stride_m, c_stride_n = 0, y.stride(0), y.stride(1)
+     else:
+         c_stride_k, c_stride_m, c_stride_n = y_pp.stride(0), y_pp.stride(1), y_pp.stride(2)
+  
+     params = {
+         'config': config,
+         'K_packed': K_packed,
+         'grid': grid,
+         'ks': ks,
+         'c_stride_k': c_stride_k,
+         'c_stride_m': c_stride_m,
+         'c_stride_n': c_stride_n,
+     }
+  
+     if ks > 1:
+         params['reduce_grid'] = (triton.cdiv(m, 16), triton.cdiv(n, 64))
+         params['actual_ksplit'] = triton.cdiv(K_packed, (config["SPLITK_BLOCK_SIZE"] // 2))
+         params['max_ksplit'] = triton.next_power_of_2(ks)
+  
+     return params
+  
+  
+ def fused_quant_gemm(A_bf16, B_shuffle, B_scale_sh, m, n, k):
+     key = (m, n, k)
+     if key not in _launch_cache:
+         _launch_cache[key] = _build_launch_params(m, n, k, A_bf16.device)
+     p = _launch_cache[key]
+     y, y_pp = _get_buffers(m, n, p['ks'], A_bf16.device)
+  
+     b_reshaped, b_scale_uint8 = _prepare_b_views(
+         B_shuffle, B_scale_sh, n, p['K_packed']
+     )
+  
+     _fused_quant_gemm_preshuffle_kernel[p['grid']](
+         A_bf16, b_reshaped,
+         y if p['ks'] == 1 else y_pp,
+         b_scale_uint8,
+         m, n, p['K_packed'],
+         A_bf16.stride(0), A_bf16.stride(1),
+         b_reshaped.stride(0), b_reshaped.stride(1),
+         p['c_stride_k'], p['c_stride_m'], p['c_stride_n'],
+         b_scale_uint8.stride(0), b_scale_uint8.stride(1),
+         **p['config'],
+     )
+  
+     if p['ks'] > 1:
+         _reduce_kernel[p['reduce_grid']](
+             y_pp, y, m, n,
+             y_pp.stride(0), y_pp.stride(1), y_pp.stride(2),
+             y.stride(0), y.stride(1),
+             16, 64,
+             p['actual_ksplit'], p['max_ksplit'],
          )
   
-         return out
+     return y
+  
+  
+ def separate_quant_gemm(A_bf16, B_shuffle, B_scale_sh, m, n, k):
+     K_packed = k // 2
+     K_bf16 = k
+  
+     QUANT_BM = 16
+     QUANT_BK = 256
+     A_fp4 = torch.empty((m, K_packed), dtype=torch.uint8, device=A_bf16.device)
+     A_scale = torch.empty((m, K_bf16 // 32), dtype=torch.uint8, device=A_bf16.device)
+  
+     grid_quant = (triton.cdiv(m, QUANT_BM), triton.cdiv(K_bf16, QUANT_BK))
+     _standalone_quant_kernel[grid_quant](
+         A_bf16, A_fp4, A_scale,
+         m, K_bf16,
+         A_bf16.stride(0), A_bf16.stride(1),
+         A_fp4.stride(0), A_fp4.stride(1),
+         A_scale.stride(0), A_scale.stride(1),
+         QUANT_BM, QUANT_BK,
+     )
+  
+     config = GEMM_CONFIGS.get((m, n, k), GEMM_DEFAULT_CONFIG).copy()
+  
+     if config["NUM_KSPLIT"] > 1:
+         SPLITK_BLOCK_SIZE, BLOCK_SIZE_K, NUM_KSPLIT = get_splitk(
+             K_packed, config["BLOCK_SIZE_K"], config["NUM_KSPLIT"]
+         )
+         config["SPLITK_BLOCK_SIZE"] = SPLITK_BLOCK_SIZE
+         config["BLOCK_SIZE_K"] = BLOCK_SIZE_K
+         config["NUM_KSPLIT"] = NUM_KSPLIT
      else:
-         _fused_mxfp4_quant_shuffle_kernel[buf['grid']](
-             A,
-             buf['x_fp4'],
-             buf['blockscale'],
-             *A.stride(),
-             *buf['x_fp4'].stride(),
-             M=M,
-             N=K,
-             BLOCK_SIZE_M=buf['BLOCK_SIZE_M'],
-             BLOCK_SIZE_N=buf['BLOCK_SIZE_N'],
-             NUM_ITER=buf['NUM_ITER'],
-             NUM_STAGES=buf['NUM_STAGES'],
-             MXFP4_QUANT_BLOCK_SIZE=32,
-             SCALING_MODE=0,
-             SCALE_N_PAD=buf['SCALE_N'],
-             num_warps=buf['NUM_WARPS'],
-             waves_per_eu=0,
-             num_stages=1,
+         config["SPLITK_BLOCK_SIZE"] = 2 * K_packed
+         config["NUM_KSPLIT"] = 1
+  
+     if config["BLOCK_SIZE_K"] >= 2 * K_packed:
+         config["BLOCK_SIZE_K"] = triton.next_power_of_2(2 * K_packed)
+         config["SPLITK_BLOCK_SIZE"] = 2 * K_packed
+         config["NUM_KSPLIT"] = 1
+  
+     config["BLOCK_SIZE_N"] = max(config["BLOCK_SIZE_N"], 32)
+  
+     y = torch.empty((m, n), dtype=torch.bfloat16, device=A_bf16.device)
+  
+     if config["NUM_KSPLIT"] > 1:
+         y_pp = torch.empty(
+             (config["NUM_KSPLIT"], m, n), dtype=torch.float32, device=A_bf16.device
          )
+     else:
+         y_pp = None
   
-         gemm_a4w4_asm(
-             buf['x_fp4'].view(dtypes.fp4x2),
-             B_shuffle,
-             buf['blockscale'].view(dtypes.fp8_e8m0),
-             B_scale_sh,
-             buf['gemm_out'],
-             _ASM_KERNEL_32x128,
-             None,
-             1.0,
-             0.0,
-             True,
-             log2_k_split=0,
+     b_reshaped, b_scale_uint8 = _prepare_b_views(
+         B_shuffle, B_scale_sh, n, K_packed
+     )
+  
+     grid = lambda META: (
+         META["NUM_KSPLIT"]
+         * triton.cdiv(m, META["BLOCK_SIZE_M"])
+         * triton.cdiv(n, META["BLOCK_SIZE_N"]),
+     )
+  
+     _gemm_only_preshuffle_kernel[grid](
+         A_fp4, A_scale,
+         b_reshaped,
+         y if config["NUM_KSPLIT"] == 1 else y_pp,
+         b_scale_uint8,
+         m, n, K_packed,
+         A_fp4.stride(0), A_fp4.stride(1),
+         A_scale.stride(0), A_scale.stride(1),
+         b_reshaped.stride(0), b_reshaped.stride(1),
+         0 if config["NUM_KSPLIT"] == 1 else y_pp.stride(0),
+         y.stride(0) if config["NUM_KSPLIT"] == 1 else y_pp.stride(1),
+         y.stride(1) if config["NUM_KSPLIT"] == 1 else y_pp.stride(2),
+         b_scale_uint8.stride(0), b_scale_uint8.stride(1),
+         **config,
+     )
+  
+     if config["NUM_KSPLIT"] > 1:
+         REDUCE_BLOCK_SIZE_M = 16
+         REDUCE_BLOCK_SIZE_N = 64
+         ACTUAL_KSPLIT = triton.cdiv(K_packed, (config["SPLITK_BLOCK_SIZE"] // 2))
+         grid_reduce = (
+             triton.cdiv(m, REDUCE_BLOCK_SIZE_M),
+             triton.cdiv(n, REDUCE_BLOCK_SIZE_N),
          )
+         _reduce_kernel[grid_reduce](
+             y_pp, y, m, n,
+             y_pp.stride(0), y_pp.stride(1), y_pp.stride(2),
+             y.stride(0), y.stride(1),
+             REDUCE_BLOCK_SIZE_M, REDUCE_BLOCK_SIZE_N,
+             ACTUAL_KSPLIT, triton.next_power_of_2(config["NUM_KSPLIT"]),
+         )
   
-         return buf['gemm_out'][:M]
+     return y
+  
+  
+ def custom_kernel(data: input_t) -> output_t:
+     A = data[0]
+     return fused_quant_gemm(A, data[3], data[4], A.shape[0], data[1].shape[0], A.shape[1])