submission 754952

- #!POPCORN leaderboard amd-mxfp4-mm
- #!POPCORN gpu MI355X
  """
- Fused BF16->MXFP4 quant + FP4 GEMM kernel for AMD MI355X.
- Optimized Triton kernel using:
- 1. v_cvt_scalef32_pk_fp4_bf16 hardware instruction for FP4 conversion
- 2. tl.dot_scaled for FP4 matrix multiplication
- 3. Per-shape tuned configurations
- 4. Split-K support with reduction kernel
+ V402: V398 + M4 cache_modifier=.cg
  """
  from task import input_t, output_t
  import torch
  ⋯ 2 unchanged lines
   
   
  @triton.jit
- def _bf16_to_fp4_hw(
+ def _mxfp4_quant_in_reg(
      x_bf16,
-     TILE_M: tl.constexpr,
-     TILE_K: tl.constexpr,
+     BLOCK_SIZE_M: tl.constexpr,
+     BLOCK_SIZE_K: tl.constexpr,
  ):
      """Quantize BF16 block to MXFP4 using HW v_cvt_scalef32_pk_fp4_bf16."""
-     FP4_GRP_SZ: tl.constexpr = 32
-     NUM_QBLK: tl.constexpr = TILE_K // FP4_GRP_SZ
+     MXFP4_QUANT_BLOCK_SIZE: tl.constexpr = 32
+     NUM_QUANT_BLOCKS: tl.constexpr = BLOCK_SIZE_K // MXFP4_QUANT_BLOCK_SIZE
   
      # Compute scales from FP32 values
-     x_fp32 = x_bf16.to(tl.float32).reshape(TILE_M, NUM_QBLK, FP4_GRP_SZ)
+     x_fp32 = x_bf16.to(tl.float32).reshape(BLOCK_SIZE_M, NUM_QUANT_BLOCKS, MXFP4_QUANT_BLOCK_SIZE)
   
      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_raw_i = log2_amax - 2
-     scale_e8m0_raw_i = tl.minimum(tl.maximum(scale_e8m0_raw_i, -127), 127)
-     bs_e8m0 = scale_e8m0_raw_i.to(tl.uint8) + 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_raw_i.to(tl.int32) + 127).to(tl.uint32) << 23
-     hw_scale = hw_scale_bits.to(tl.float32, bitcast=True)  # [M, NUM_QBLK, 1]
+     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]
   
      # Broadcast scale to per-pair granularity
-     hw_scale_flat = tl.broadcast_to(hw_scale, (TILE_M, NUM_QBLK, FP4_GRP_SZ))
-     hw_scale_flat = hw_scale_flat.reshape(TILE_M, TILE_K)
+     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(TILE_M, TILE_K // 2, 2)
+     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(TILE_M, TILE_K // 2)
+     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(TILE_M, TILE_K // 2, 2)
+     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(TILE_M, TILE_K // 2)
+     x_u32 = x_u32.reshape(BLOCK_SIZE_M, BLOCK_SIZE_K // 2)
   
      # HW FP4 conversion
      fp4_u32 = tl.inline_asm_elementwise(
  ⋯ 5 unchanged lines
          pack=1,
      )
      x_fp4 = (fp4_u32 & 0xFF).to(tl.uint8)
-     x_fp4 = x_fp4.reshape(TILE_M, TILE_K // 2)
+     x_fp4 = x_fp4.reshape(BLOCK_SIZE_M, BLOCK_SIZE_K // 2)
   
-     return x_fp4, bs_e8m0.reshape(TILE_M, NUM_QBLK)
+     return x_fp4, bs_e8m0.reshape(BLOCK_SIZE_M, NUM_QUANT_BLOCKS)
   
   
  @triton.jit
- def _quant_only_launcher(
+ 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,
-     TILE_M: tl.constexpr,
-     TILE_K: tl.constexpr,
+     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 * TILE_M + tl.arange(0, TILE_M)
-     offs_k = pid_k * TILE_K + tl.arange(0, TILE_K)
+     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 = _bf16_to_fp4_hw(a_bf16, TILE_M, TILE_K)
-     HALF_K: tl.constexpr = TILE_K // 2
+     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 = TILE_K // 32
+     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)))
  ⋯ 1 unchanged lines
   
  @triton.heuristics(
      {
-         "EVEN_K": lambda args: (args["K"] % (args["TILE_K"] // 2) == 0)
-         and (args["KSPLIT_TILE"] % args["TILE_K"] == 0)
-         and (args["K"] % (args["KSPLIT_TILE"] // 2) == 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 _quant_gemm_fused_launcher(
+ 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,
-     TILE_M: tl.constexpr,
-     TILE_N: tl.constexpr,
-     TILE_K: tl.constexpr,
-     M_CLUSTER: tl.constexpr,
+     BLOCK_SIZE_M: tl.constexpr,
+     BLOCK_SIZE_N: tl.constexpr,
+     BLOCK_SIZE_K: tl.constexpr,
+     GROUP_SIZE_M: tl.constexpr,
      NUM_KSPLIT: tl.constexpr,
-     KSPLIT_TILE: tl.constexpr,
+     SPLITK_BLOCK_SIZE: tl.constexpr,
      EVEN_K: tl.constexpr,
      num_warps: tl.constexpr,
      num_stages: tl.constexpr,
  ⋯ 10 unchanged lines
      tl.assume(stride_bsn > 0)
      tl.assume(stride_bsk > 0)
   
-     SCALE_GRP_SZ: tl.constexpr = 32
-     num_pid_m = tl.cdiv(M, TILE_M)
-     num_pid_n = tl.cdiv(N, TILE_N)
+     SCALE_GROUP_SIZE: tl.constexpr = 32
+     num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+     num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
   
      pid_unified = tl.program_id(axis=0)
      pid_k = pid_unified % NUM_KSPLIT
      pid = pid_unified // NUM_KSPLIT
   
      if NUM_KSPLIT == 1:
-         num_pid_in_group = M_CLUSTER * num_pid_n
+         num_pid_in_group = GROUP_SIZE_M * num_pid_n
          group_id = pid // num_pid_in_group
-         first_pid_m = group_id * M_CLUSTER
-         group_size_m = min(num_pid_m - first_pid_m, M_CLUSTER)
+         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:
  ⋯ 4 unchanged lines
      tl.assume(pid_n >= 0)
      tl.assume(pid_k >= 0)
   
-     if (pid_k * KSPLIT_TILE // 2) < K:
-         num_k_iter = tl.cdiv(KSPLIT_TILE // 2, TILE_K // 2)
+     if (pid_k * SPLITK_BLOCK_SIZE // 2) < K:
+         num_k_iter = tl.cdiv(SPLITK_BLOCK_SIZE // 2, BLOCK_SIZE_K // 2)
   
          # A: BF16 [M, 2*K]
-         offs_am = (pid_m * TILE_M + tl.arange(0, TILE_M)) % M
-         offs_ak = pid_k * KSPLIT_TILE + tl.arange(0, TILE_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)
   
          # B: pre-shuffled MXFP4 [N//16, K_packed*16]
-         offs_k_shuffle_arr = tl.arange(0, (TILE_K // 2) * 16)
-         offs_k_shuffle = pid_k * (KSPLIT_TILE // 2) * 16 + offs_k_shuffle_arr
-         offs_bn = (pid_n * (TILE_N // 16) + tl.arange(0, TILE_N // 16)) % (N // 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)
   
          # B_scale: shuffled E8M0 [N_pad, K_scale_pad]
          # Each group of 32 N values occupies 32 consecutive rows.
-         # Row index = pid_n * TILE_N + group_offset * 32
-         offs_bsn = (pid_n * TILE_N + tl.arange(0, TILE_N // 32) * 32)
-         offs_ks = (pid_k * (KSPLIT_TILE // SCALE_GRP_SZ) * 32) + tl.arange(
-             0, TILE_K // SCALE_GRP_SZ * 32
+         # 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
          )
   
-         acc = tl.zeros((TILE_M, TILE_N), dtype=tl.float32)
+         accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
   
          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
  ⋯ 2 unchanged lines
                  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) * TILE_K
+                 k_offset = (k_iter - pid_k * num_k_iter) * BLOCK_SIZE_K
                  a_bf16 = tl.load(
                      a_ptrs,
-                     mask=tl.arange(0, TILE_K)[None, :] < (2 * K - pid_k * KSPLIT_TILE - k_offset),
+                     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 * (KSPLIT_TILE // 2) + (k_iter - pid_k * num_k_iter) * (TILE_K // 2))) * 16),
+                     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 = _bf16_to_fp4_hw(a_bf16, TILE_M, TILE_K)
+             a_fp4, a_scales = _mxfp4_quant_in_reg(a_bf16, BLOCK_SIZE_M, BLOCK_SIZE_K)
   
              # Unshuffle B scales
              b_scales = (
                  b_scales_raw
                  .reshape(
-                     TILE_N // 32,
-                     TILE_K // SCALE_GRP_SZ // 8,
+                     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(TILE_N, TILE_K // SCALE_GRP_SZ)
+                 .reshape(BLOCK_SIZE_N, BLOCK_SIZE_K // SCALE_GROUP_SIZE)
              )
   
              # Unshuffle B data
              b = (
-                 b_raw.reshape(1, TILE_N // 16, TILE_K // 64, 2, 16, 16)
+                 b_raw.reshape(1, BLOCK_SIZE_N // 16, BLOCK_SIZE_K // 64, 2, 16, 16)
                  .permute(0, 1, 4, 2, 3, 5)
-                 .reshape(TILE_N, TILE_K // 2)
+                 .reshape(BLOCK_SIZE_N, BLOCK_SIZE_K // 2)
                  .trans(1, 0)
              )
   
-             acc = tl.dot_scaled(
-                 a_fp4, a_scales, "e2m1", b, b_scales, "e2m1", acc
+             accumulator = tl.dot_scaled(
+                 a_fp4, a_scales, "e2m1", b, b_scales, "e2m1", accumulator
              )
   
-             a_ptrs += TILE_K * stride_ak
-             b_ptrs += (TILE_K // 2) * 16 * stride_bk
-             b_scale_ptrs += TILE_K * stride_bsk
+             a_ptrs += BLOCK_SIZE_K * stride_ak
+             b_ptrs += (BLOCK_SIZE_K // 2) * 16 * stride_bk
+             b_scale_ptrs += BLOCK_SIZE_K * stride_bsk
   
-         c = acc.to(c_ptr.type.element_ty)
+         c = accumulator.to(c_ptr.type.element_ty)
   
-         offs_cm = pid_m * TILE_M + tl.arange(0, TILE_M).to(tl.int64)
-         offs_cn = pid_n * TILE_N + tl.arange(0, TILE_N).to(tl.int64)
+         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]
  ⋯ 5 unchanged lines
   
   
  @triton.jit
- def _splitk_sum_launcher(
+ 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,
-     TILE_M: tl.constexpr, TILE_N: tl.constexpr,
+     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 * TILE_M + tl.arange(0, TILE_M)) % M
-     offs_n = (pid_n * TILE_N + tl.arange(0, TILE_N)) % N
+     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 = (
  ⋯ 17 unchanged lines
   
  @triton.heuristics(
      {
-         "EVEN_K": lambda args: (args["K"] % (args["TILE_K"] // 2) == 0)
-         and (args["KSPLIT_TILE"] % args["TILE_K"] == 0)
-         and (args["K"] % (args["KSPLIT_TILE"] // 2) == 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 _fp4_gemm_only_launcher(
+ 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,
  ⋯ 1 unchanged lines
      stride_bn, stride_bk,
      stride_ck, stride_cm, stride_cn,
      stride_bsn, stride_bsk,
-     TILE_M: tl.constexpr,
-     TILE_N: tl.constexpr,
-     TILE_K: tl.constexpr,
-     M_CLUSTER: tl.constexpr,
+     BLOCK_SIZE_M: tl.constexpr,
+     BLOCK_SIZE_N: tl.constexpr,
+     BLOCK_SIZE_K: tl.constexpr,
+     GROUP_SIZE_M: tl.constexpr,
      NUM_KSPLIT: tl.constexpr,
-     KSPLIT_TILE: tl.constexpr,
+     SPLITK_BLOCK_SIZE: tl.constexpr,
      EVEN_K: tl.constexpr,
      num_warps: tl.constexpr,
      num_stages: tl.constexpr,
  ⋯ 12 unchanged lines
      tl.assume(stride_bsn > 0)
      tl.assume(stride_bsk > 0)
   
-     SCALE_GRP_SZ: tl.constexpr = 32
-     HALF_TK: tl.constexpr = TILE_K // 2
-     SCALE_TK: tl.constexpr = TILE_K // SCALE_GRP_SZ
-     num_pid_m = tl.cdiv(M, TILE_M)
-     num_pid_n = tl.cdiv(N, TILE_N)
+     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)
   
      pid_unified = tl.program_id(axis=0)
      pid_k = pid_unified % NUM_KSPLIT
      pid = pid_unified // NUM_KSPLIT
   
      if NUM_KSPLIT == 1:
-         num_pid_in_group = M_CLUSTER * num_pid_n
+         num_pid_in_group = GROUP_SIZE_M * num_pid_n
          group_id = pid // num_pid_in_group
-         first_pid_m = group_id * M_CLUSTER
-         group_size_m = min(num_pid_m - first_pid_m, M_CLUSTER)
+         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:
  ⋯ 4 unchanged lines
      tl.assume(pid_n >= 0)
      tl.assume(pid_k >= 0)
   
-     if (pid_k * KSPLIT_TILE // 2) < K:
-         num_k_iter = tl.cdiv(KSPLIT_TILE // 2, HALF_TK)
+     if (pid_k * SPLITK_BLOCK_SIZE // 2) < K:
+         num_k_iter = tl.cdiv(SPLITK_BLOCK_SIZE // 2, HALF_BK)
   
-         offs_am = (pid_m * TILE_M + tl.arange(0, TILE_M)) % M
-         offs_aqk = pid_k * (KSPLIT_TILE // 2) + tl.arange(0, HALF_TK)
+         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)
   
-         offs_ask = pid_k * (KSPLIT_TILE // SCALE_GRP_SZ) + tl.arange(0, SCALE_TK)
+         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)
   
-         offs_k_shuffle_arr = tl.arange(0, HALF_TK * 16)
-         offs_k_shuffle = pid_k * (KSPLIT_TILE // 2) * 16 + offs_k_shuffle_arr
-         offs_bn = (pid_n * (TILE_N // 16) + tl.arange(0, TILE_N // 16)) % (N // 16)
+         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)
   
-         offs_bsn = (pid_n * TILE_N + tl.arange(0, TILE_N // 32) * 32)
-         offs_ks = (pid_k * (KSPLIT_TILE // SCALE_GRP_SZ) * 32) + tl.arange(
-             0, SCALE_TK * 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, SCALE_BK * 32
          )
          b_scale_ptrs = (
              b_scales_ptr + offs_bsn[:, None] * stride_bsn + offs_ks[None, :] * stride_bsk
          )
   
-         acc = tl.zeros((TILE_M, TILE_N), dtype=tl.float32)
+         accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
   
          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_TK
-                 k_remain = K - (pid_k * (KSPLIT_TILE // 2) + k_off)
+                 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_TK)[None, :] < k_remain,
+                     mask=tl.arange(0, HALF_BK)[None, :] < k_remain,
                      other=0,
                      cache_modifier=cache_modifier,
                  )
-                 s_remain = (2 * K) // SCALE_GRP_SZ - (pid_k * (KSPLIT_TILE // SCALE_GRP_SZ) + (k_iter - pid_k * num_k_iter) * SCALE_TK)
+                 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_TK)[None, :] < s_remain,
+                     mask=tl.arange(0, SCALE_BK)[None, :] < s_remain,
                      other=0,
                      cache_modifier=cache_modifier,
                  )
  ⋯ 1 unchanged lines
              b_scales = (
                  tl.load(b_scale_ptrs, cache_modifier=cache_modifier)
                  .reshape(
-                     TILE_N // 32,
-                     SCALE_TK // 8,
+                     BLOCK_SIZE_N // 32,
+                     SCALE_BK // 8,
                      4, 16, 2, 2, 1,
                  )
                  .permute(0, 5, 3, 1, 4, 2, 6)
-                 .reshape(TILE_N, SCALE_TK)
+                 .reshape(BLOCK_SIZE_N, SCALE_BK)
              )
   
              if EVEN_K:
  ⋯ 1 unchanged lines
              else:
                  b = tl.load(
                      b_ptrs,
-                     mask=offs_k_shuffle_arr[None, :] < ((K - (pid_k * (KSPLIT_TILE // 2) + (k_iter - pid_k * num_k_iter) * HALF_TK)) * 16),
+                     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, TILE_N // 16, TILE_K // 64, 2, 16, 16)
+                 b.reshape(1, BLOCK_SIZE_N // 16, BLOCK_SIZE_K // 64, 2, 16, 16)
                  .permute(0, 1, 4, 2, 3, 5)
-                 .reshape(TILE_N, HALF_TK)
+                 .reshape(BLOCK_SIZE_N, HALF_BK)
                  .trans(1, 0)
              )
   
-             acc = tl.dot_scaled(
-                 a_fp4, a_scales, "e2m1", b, b_scales, "e2m1", acc
+             accumulator = tl.dot_scaled(
+                 a_fp4, a_scales, "e2m1", b, b_scales, "e2m1", accumulator
              )
   
-             a_fp4_ptrs += HALF_TK * stride_qk
-             a_scale_ptrs += SCALE_TK * stride_sk
-             b_ptrs += HALF_TK * 16 * stride_bk
-             b_scale_ptrs += TILE_K * stride_bsk
+             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 = acc.to(c_ptr.type.element_ty)
+         c = accumulator.to(c_ptr.type.element_ty)
   
-         offs_cm = pid_m * TILE_M + tl.arange(0, TILE_M).to(tl.int64)
-         offs_cn = pid_n * TILE_N + tl.arange(0, TILE_N).to(tl.int64)
+         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]
  ⋯ 4 unchanged lines
          tl.store(c_ptrs, c, mask=c_mask)
   
   
- def _calc_splitk_params(K, TILE_K, NUM_KSPLIT):
-     KSPLIT_TILE = (
-         triton.cdiv((2 * triton.cdiv(K, NUM_KSPLIT)), TILE_K) * TILE_K
+ 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 TILE_K > 16:
+     while NUM_KSPLIT > 1 and BLOCK_SIZE_K > 16:
          if (
-             K % (KSPLIT_TILE // 2) == 0
-             and KSPLIT_TILE % TILE_K == 0
-             and K % (TILE_K // 2) == 0
+             K % (SPLITK_BLOCK_SIZE // 2) == 0
+             and SPLITK_BLOCK_SIZE % BLOCK_SIZE_K == 0
+             and K % (BLOCK_SIZE_K // 2) == 0
          ):
              break
-         elif K % (KSPLIT_TILE // 2) != 0 and NUM_KSPLIT > 1:
+         elif K % (SPLITK_BLOCK_SIZE // 2) != 0 and NUM_KSPLIT > 1:
              NUM_KSPLIT = NUM_KSPLIT // 2
-         elif KSPLIT_TILE % TILE_K != 0:
+         elif SPLITK_BLOCK_SIZE % BLOCK_SIZE_K != 0:
              if NUM_KSPLIT > 1:
                  NUM_KSPLIT = NUM_KSPLIT // 2
-             elif TILE_K > 16:
-                 TILE_K = TILE_K // 2
-         elif K % (TILE_K // 2) != 0 and TILE_K > 16:
-             TILE_K = TILE_K // 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
-         KSPLIT_TILE = (
-             triton.cdiv((2 * triton.cdiv(K, NUM_KSPLIT)), TILE_K) * TILE_K
+         SPLITK_BLOCK_SIZE = (
+             triton.cdiv((2 * triton.cdiv(K, NUM_KSPLIT)), BLOCK_SIZE_K) * BLOCK_SIZE_K
          )
-     NUM_KSPLIT = triton.cdiv(K, (KSPLIT_TILE // 2))
-     return KSPLIT_TILE, TILE_K, NUM_KSPLIT
+     NUM_KSPLIT = triton.cdiv(K, (SPLITK_BLOCK_SIZE // 2))
+     return SPLITK_BLOCK_SIZE, BLOCK_SIZE_K, NUM_KSPLIT
   
   
- _TUNE_PARAMS = {
-     # V10: Conservative tuning based on V8 baseline
-     # M=4: Try larger TILE_N for better N parallelism
-     (4, 2880, 512): {"TILE_M": 16, "TILE_N": 64, "TILE_K": 256, "M_CLUSTER": 1, "num_warps": 2, "num_stages": 3, "waves_per_eu": 3, "matrix_instr_nonkdim": 16, "cache_modifier": ".cg", "NUM_KSPLIT": 1},
-     # M=16, K=7168: Keep Split-K=7, try num_stages=3
-     (16, 2112, 7168): {"TILE_M": 16, "TILE_N": 128, "TILE_K": 512, "M_CLUSTER": 1, "num_warps": 4, "num_stages": 3, "waves_per_eu": 3, "matrix_instr_nonkdim": 16, "cache_modifier": ".cg", "NUM_KSPLIT": 7},
-     # M=32: Keep V8 config
-     (32, 4096, 512): {"TILE_M": 16, "TILE_N": 32, "TILE_K": 256, "M_CLUSTER": 1, "num_warps": 4, "num_stages": 3, "waves_per_eu": 3, "matrix_instr_nonkdim": 16, "cache_modifier": ".cg", "NUM_KSPLIT": 1},
-     (32, 2880, 512): {"TILE_M": 16, "TILE_N": 32, "TILE_K": 256, "M_CLUSTER": 1, "num_warps": 4, "num_stages": 3, "waves_per_eu": 3, "matrix_instr_nonkdim": 16, "cache_modifier": ".cg", "NUM_KSPLIT": 1},
-     # M=64: Try waves_per_eu=2
-     (64, 7168, 2048): {"TILE_M": 16, "TILE_N": 128, "TILE_K": 256, "M_CLUSTER": 1, "num_warps": 4, "num_stages": 2, "waves_per_eu": 2, "matrix_instr_nonkdim": 16, "cache_modifier": ".cg", "NUM_KSPLIT": 1},
-     # M=256: Try waves_per_eu=3
-     (256, 3072, 1536): {"TILE_M": 16, "TILE_N": 256, "TILE_K": 512, "M_CLUSTER": 1, "num_warps": 8, "num_stages": 2, "waves_per_eu": 3, "matrix_instr_nonkdim": 16, "cache_modifier": ".cg", "NUM_KSPLIT": 1},
+ TUNE_CONFIGS = {
+     (4, 2880, 512): {'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': 1, '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": 14},
+     (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": 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},
+     (64, 7168, 2048): {"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": 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},
  }
   
- _DEF_PARAMS = {"TILE_M": 16, "TILE_N": 32, "TILE_K": 256, "M_CLUSTER": 1, "num_warps": 2, "num_stages": 2, "waves_per_eu": 0, "matrix_instr_nonkdim": 16, "cache_modifier": ".cg", "NUM_KSPLIT": 1}
+ DEFAULT_TUNE_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}
   
- _FP4_GEMM_PARAMS = {
-     (32, 4096, 512): {"TILE_M": 32, "TILE_N": 128, "TILE_K": 256, "M_CLUSTER": 4, "num_warps": 4, "num_stages": 2, "waves_per_eu": 0, "matrix_instr_nonkdim": 16, "cache_modifier": ".cg", "NUM_KSPLIT": 1},
-     (32, 2880, 512): {"TILE_M": 32, "TILE_N": 64, "TILE_K": 256, "M_CLUSTER": 4, "num_warps": 4, "num_stages": 2, "waves_per_eu": 0, "matrix_instr_nonkdim": 16, "cache_modifier": ".cg", "NUM_KSPLIT": 1},
-     (64, 7168, 2048): {"TILE_M": 16, "TILE_N": 128, "TILE_K": 512, "M_CLUSTER": 4, "num_warps": 4, "num_stages": 2, "waves_per_eu": 0, "matrix_instr_nonkdim": 16, "cache_modifier": ".cg", "NUM_KSPLIT": 2},
+ GEMM_TUNE_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},
  }
   
- _FP4_GEMM_DEF = {"TILE_M": 16, "TILE_N": 64, "TILE_K": 256, "M_CLUSTER": 4, "num_warps": 4, "num_stages": 2, "waves_per_eu": 0, "matrix_instr_nonkdim": 16, "cache_modifier": ".cg", "NUM_KSPLIT": 1}
+ 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}
   
   
- _mem_pool = {}
- _param_pool = {}
+ _buf_cache = {}
+ _config_cache = {}
   
- def _acquire_buffers(m, n, num_ksplit, device):
+ def _get_buffers(m, n, num_ksplit, device):
      key = (m, n, num_ksplit)
-     if key not in _mem_pool:
+     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
-         _mem_pool[key] = (y, y_pp)
-     return _mem_pool[key]
+         _buf_cache[key] = (y, y_pp)
+     return _buf_cache[key]
   
- def _acquire_config(m, n, k):
+ def _get_config(m, n, k):
      key = (m, n, k)
-     if key not in _param_pool:
-         config = _TUNE_PARAMS.get(key, _DEF_PARAMS).copy()
+     if key not in _config_cache:
+         config = TUNE_CONFIGS.get(key, DEFAULT_TUNE_CONFIG).copy()
          K_packed = k // 2
          if config["NUM_KSPLIT"] > 1:
-             KSPLIT_TILE, TILE_K, NUM_KSPLIT = _calc_splitk_params(
-                 K_packed, config["TILE_K"], config["NUM_KSPLIT"]
+             SPLITK_BLOCK_SIZE, BLOCK_SIZE_K, NUM_KSPLIT = get_splitk(
+                 K_packed, config["BLOCK_SIZE_K"], config["NUM_KSPLIT"]
              )
-             config["KSPLIT_TILE"] = KSPLIT_TILE
-             config["TILE_K"] = TILE_K
+             config["SPLITK_BLOCK_SIZE"] = SPLITK_BLOCK_SIZE
+             config["BLOCK_SIZE_K"] = BLOCK_SIZE_K
              config["NUM_KSPLIT"] = NUM_KSPLIT
          else:
-             config["KSPLIT_TILE"] = 2 * K_packed
+             config["SPLITK_BLOCK_SIZE"] = 2 * K_packed
              config["NUM_KSPLIT"] = 1
-         if config["TILE_K"] >= 2 * K_packed:
-             config["TILE_K"] = triton.next_power_of_2(2 * K_packed)
-             config["KSPLIT_TILE"] = 2 * K_packed
+         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["TILE_N"] = max(config["TILE_N"], 32)
-         _param_pool[key] = config
-     return _param_pool[key]
+         config["BLOCK_SIZE_N"] = max(config["BLOCK_SIZE_N"], 32)
+         _config_cache[key] = config
+     return _config_cache[key]
   
   
- _grid_pool = {}
+ _launch_cache = {}
   
- def _prepare_launch_meta(m, n, k, device):
+  
+ 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 = _acquire_config(m, n, k)
+     config = _get_config(m, n, k)
      K_packed = k // 2
      ks = config["NUM_KSPLIT"]
   
-     y, y_pp = _acquire_buffers(m, n, ks, device)
+     y, y_pp = _get_buffers(m, n, ks, device)
   
-     grid = (ks * triton.cdiv(m, config["TILE_M"]) * triton.cdiv(n, config["TILE_N"]),)
+     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:
  ⋯ 1 unchanged lines
      else:
          c_stride_k, c_stride_m, c_stride_n = y_pp.stride(0), y_pp.stride(1), y_pp.stride(2)
   
-     meta = {
+     params = {
          'config': config,
          'K_packed': K_packed,
          'grid': grid,
  ⋯ 4 unchanged lines
      }
   
      if ks > 1:
-         meta['reduce_grid'] = (triton.cdiv(m, 16), triton.cdiv(n, 64))
-         meta['actual_ksplit'] = triton.cdiv(K_packed, (config["KSPLIT_TILE"] // 2))
-         meta['max_ksplit'] = triton.next_power_of_2(ks)
+         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 meta
+     return params
   
   
  def fused_quant_gemm(A_bf16, B_shuffle, B_scale_sh, m, n, k):
      key = (m, n, k)
-     if key not in _grid_pool:
-         _grid_pool[key] = _prepare_launch_meta(m, n, k, A_bf16.device)
-     p = _grid_pool[key]
-     y, y_pp = _acquire_buffers(m, n, p['ks'], A_bf16.device)
+     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_shuffle.view(torch.uint8).reshape(n // 16, p['K_packed'] * 16)
-     b_scale_uint8 = B_scale_sh.view(torch.uint8)
+     b_reshaped, b_scale_uint8 = _prepare_b_views(
+         B_shuffle, B_scale_sh, n, p['K_packed']
+     )
   
-     _quant_gemm_fused_launcher[p['grid']](
+     _fused_quant_gemm_preshuffle_kernel[p['grid']](
          A_bf16, b_reshaped,
          y if p['ks'] == 1 else y_pp,
          b_scale_uint8,
  ⋯ 6 unchanged lines
      )
   
      if p['ks'] > 1:
-         _splitk_sum_launcher[p['reduce_grid']](
+         _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),
  ⋯ 14 unchanged lines
      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))
-     _quant_only_launcher[grid_quant](
+     _standalone_quant_kernel[grid_quant](
          A_bf16, A_fp4, A_scale,
          m, K_bf16,
          A_bf16.stride(0), A_bf16.stride(1),
  ⋯ 2 unchanged lines
          QUANT_BM, QUANT_BK,
      )
   
-     config = _FP4_GEMM_PARAMS.get((m, n, k), _FP4_GEMM_DEF).copy()
+     config = GEMM_TUNE_CONFIGS.get((m, n, k), GEMM_DEFAULT_CONFIG).copy()
   
      if config["NUM_KSPLIT"] > 1:
-         KSPLIT_TILE, TILE_K, NUM_KSPLIT = _calc_splitk_params(
-             K_packed, config["TILE_K"], config["NUM_KSPLIT"]
+         SPLITK_BLOCK_SIZE, BLOCK_SIZE_K, NUM_KSPLIT = get_splitk(
+             K_packed, config["BLOCK_SIZE_K"], config["NUM_KSPLIT"]
          )
-         config["KSPLIT_TILE"] = KSPLIT_TILE
-         config["TILE_K"] = TILE_K
+         config["SPLITK_BLOCK_SIZE"] = SPLITK_BLOCK_SIZE
+         config["BLOCK_SIZE_K"] = BLOCK_SIZE_K
          config["NUM_KSPLIT"] = NUM_KSPLIT
      else:
-         config["KSPLIT_TILE"] = 2 * K_packed
+         config["SPLITK_BLOCK_SIZE"] = 2 * K_packed
          config["NUM_KSPLIT"] = 1
   
-     if config["TILE_K"] >= 2 * K_packed:
-         config["TILE_K"] = triton.next_power_of_2(2 * K_packed)
-         config["KSPLIT_TILE"] = 2 * K_packed
+     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["TILE_N"] = max(config["TILE_N"], 32)
+     config["BLOCK_SIZE_N"] = max(config["BLOCK_SIZE_N"], 32)
   
      y = torch.empty((m, n), dtype=torch.bfloat16, device=A_bf16.device)
   
  ⋯ 4 unchanged lines
      else:
          y_pp = None
   
-     b_uint8 = B_shuffle.view(torch.uint8)
-     b_reshaped = b_uint8.reshape(n // 16, K_packed * 16)
-     b_scale_uint8 = B_scale_sh.view(torch.uint8)
+     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["TILE_M"])
-         * triton.cdiv(n, META["TILE_N"]),
+         * triton.cdiv(m, META["BLOCK_SIZE_M"])
+         * triton.cdiv(n, META["BLOCK_SIZE_N"]),
      )
   
-     _fp4_gemm_only_launcher[grid](
+     _gemm_only_preshuffle_kernel[grid](
          A_fp4, A_scale,
          b_reshaped,
          y if config["NUM_KSPLIT"] == 1 else y_pp,
  ⋯ 10 unchanged lines
      )
   
      if config["NUM_KSPLIT"] > 1:
-         REDUCE_TILE_M = 16
-         REDUCE_TILE_N = 64
-         ACTUAL_KSPLIT = triton.cdiv(K_packed, (config["KSPLIT_TILE"] // 2))
+         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_TILE_M),
-             triton.cdiv(n, REDUCE_TILE_N),
+             triton.cdiv(m, REDUCE_BLOCK_SIZE_M),
+             triton.cdiv(n, REDUCE_BLOCK_SIZE_N),
          )
-         _splitk_sum_launcher[grid_reduce](
+         _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_TILE_M, REDUCE_TILE_N,
+             REDUCE_BLOCK_SIZE_M, REDUCE_BLOCK_SIZE_N,
              ACTUAL_KSPLIT, triton.next_power_of_2(config["NUM_KSPLIT"]),
          )
   
  ⋯ 2 unchanged lines
   
  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])
   No newline at end of file
+     return fused_quant_gemm(A, data[3], data[4], A.shape[0], data[1].shape[0], A.shape[1])