submission 539978

  ⋯ 1 unchanged lines
  #!POPCORN gpu MI355X
   
  """
- v46: Preshuffle-scales GEMM — reads B_scale_sh directly (zero unshuffle overhead).
- For M<=16, K>=2048: call _gemm_afp4wfp4_kernel_preshuffle_scales directly,
- bypassing the M>=32 wrapper assertion. The kernel handles M<32 with raw A scales
- and does in-register B scale unshuffle via reshape+permute (free, hidden by mem latency).
- This eliminates the ~5µs unshuffle that killed v45 ranked performance.
- For all others: ASM GEMM (v35 approach).
+ v49: Fused quant+GEMM for ALL M<=16 shapes (not just K>=2048).
+ Saves 1 kernel launch for M=4/K=512 shapes.
+ ASM GEMM for M>=32 (where ASM is faster).
+ XCD remap bug fix: pad grid to multiple of 8.
  """
  from task import input_t, output_t
   
  ⋯ 6 unchanged lines
  from aiter.ops.gemm_op_a4w4 import get_GEMM_config
  from aiter.ops.gemm_op_common import get_padded_m
   
- # Import preshuffle-scales kernel directly (bypasses M>=32 wrapper assert)
- try:
-     import aiter.ops.triton.gemm_afp4wfp4 as _gemm_mod
-     _ps_kernel = _gemm_mod._gemm_afp4wfp4_kernel_preshuffle_scales
-     _reduce_kernel = _gemm_mod._gemm_afp4wfp4_reduce_kernel
-     _get_splitk_fn = _gemm_mod.get_splitk
-     _HAS_PS = True
- except (ImportError, AttributeError):
-     _HAS_PS = False
+ import aiter.ops.triton.gemm_afp4wfp4 as _gemm_mod
+ _reduce_kernel = _gemm_mod._gemm_afp4wfp4_reduce_kernel
+ _get_splitk_fn = _gemm_mod.get_splitk
   
  _fp4x2 = dtypes.fp4x2
  _fp8_e8m0 = dtypes.fp8_e8m0
  ⋯ 1 unchanged lines
   
   
  @triton.jit
- def _quant_raw_kernel(
-     x_ptr, x_fp4_ptr, scale_ptr,
-     stride_x_m, stride_x_n,
-     stride_fp4_m, stride_fp4_n,
-     stride_sc_m, stride_sc_n,
-     M, K,
-     BLOCK_M: tl.constexpr, BLOCK_K: tl.constexpr,
+ def _remap_xcd(pid, num_pids, NUM_XCDS: tl.constexpr):
+     chunk_size = tl.cdiv(num_pids, NUM_XCDS)
+     xcd = pid % NUM_XCDS
+     pid_in_xcd = pid // NUM_XCDS
+     return xcd * chunk_size + pid_in_xcd
+  
+  
+ @triton.jit
+ def _pid_grid(pid, num_pid_m, num_pid_n, GROUP_SIZE_M: tl.constexpr):
+     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
+     return pid_m, pid_n
+  
+  
+ @triton.jit
+ def _fused_quant_gemm_kernel(
+     a_ptr, b_ptr, c_ptr, b_scales_ptr,
+     M, N, K_real,
+     stride_am, stride_ak,
+     stride_bk, stride_bn,
+     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,
      QUANT_BLOCK: tl.constexpr,
  ):
-     pid_m = tl.program_id(0)
-     pid_k = tl.program_id(1)
-     offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
-     offs_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K)
-     mask = (offs_m[:, None] < M) & (offs_k[None, :] < K)
-     x = tl.load(x_ptr + offs_m[:, None] * stride_x_m + offs_k[None, :] * stride_x_n,
-                 mask=mask, other=0.0).to(tl.float32)
+     SCALE_GROUP_SIZE: tl.constexpr = 32
+     K_packed = K_real // 2
+     GRID_MN = tl.cdiv(M, BLOCK_SIZE_M) * tl.cdiv(N, BLOCK_SIZE_N)
+     total_pids = GRID_MN * NUM_KSPLIT
+     # Pad to multiple of 8 for correct XCD remapping
+     total_pids_padded = ((total_pids + 7) // 8) * 8
   
-     out_fp4, scales_e8m0 = _mxfp4_quant_op(x, BLOCK_K, BLOCK_M, QUANT_BLOCK)
+     pid_unified = tl.program_id(axis=0)
+     pid_unified = _remap_xcd(pid_unified, total_pids_padded, NUM_XCDS=8)
   
-     fp4_offs_k = pid_k * BLOCK_K // 2 + tl.arange(0, BLOCK_K // 2)
-     fp4_mask = (offs_m[:, None] < M) & (fp4_offs_k[None, :] < K // 2)
-     tl.store(x_fp4_ptr + offs_m[:, None] * stride_fp4_m + fp4_offs_k[None, :] * stride_fp4_n,
-              out_fp4, mask=fp4_mask)
+     if pid_unified < total_pids:
+         pid_k = pid_unified % NUM_KSPLIT
+         pid = pid_unified // NUM_KSPLIT
+         num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+         num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
   
-     NUM_SC: tl.constexpr = BLOCK_K // QUANT_BLOCK
-     sc_offs_k = pid_k * NUM_SC + tl.arange(0, NUM_SC)
-     sc_mask = (offs_m[:, None] < M) & (sc_offs_k[None, :] < (K + QUANT_BLOCK - 1) // QUANT_BLOCK)
-     tl.store(scale_ptr + offs_m[:, None] * stride_sc_m + sc_offs_k[None, :] * stride_sc_n,
-              scales_e8m0, mask=sc_mask)
+         if NUM_KSPLIT == 1:
+             pid_m, pid_n = _pid_grid(pid, num_pid_m, num_pid_n, GROUP_SIZE_M=GROUP_SIZE_M)
+         else:
+             pid_m = pid // num_pid_n
+             pid_n = pid % num_pid_n
   
+         tl.assume(pid_m >= 0)
+         tl.assume(pid_n >= 0)
   
+         if (pid_k * SPLITK_BLOCK_SIZE) < K_real:
+             num_k_iter = tl.cdiv(SPLITK_BLOCK_SIZE, BLOCK_SIZE_K)
+  
+             offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+             offs_k = pid_k * SPLITK_BLOCK_SIZE + tl.arange(0, BLOCK_SIZE_K)
+             a_ptrs = a_ptr + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak
+  
+             offs_k_packed = pid_k * (SPLITK_BLOCK_SIZE // 2) + tl.arange(0, BLOCK_SIZE_K // 2)
+             offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
+             b_ptrs = b_ptr + offs_k_packed[:, None] * stride_bk + offs_bn[None, :] * stride_bn
+  
+             offs_bsn = (pid_n * (BLOCK_SIZE_N // 32) + tl.arange(0, BLOCK_SIZE_N // 32)) % N
+             offs_ks_scale = (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_scale[None, :] * stride_bsk
+  
+             accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+  
+             for k in tl.range(0, num_k_iter):
+                 a_bf16 = tl.load(a_ptrs, mask=offs_k[None, :] < K_real, other=0.0).to(tl.float32)
+                 a_fp4, a_scales = _mxfp4_quant_op(a_bf16, BLOCK_SIZE_K, BLOCK_SIZE_M, QUANT_BLOCK)
+  
+                 b_fp4 = tl.load(b_ptrs, mask=offs_k_packed[:, None] < K_packed, other=0)
+  
+                 b_scales = (
+                     tl.load(b_scale_ptrs)
+                     .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)
+                 )
+  
+                 accumulator = tl.dot_scaled(a_fp4, a_scales, "e2m1", b_fp4, b_scales, "e2m1", accumulator)
+  
+                 a_ptrs += BLOCK_SIZE_K * stride_ak
+                 offs_k += BLOCK_SIZE_K
+                 b_ptrs += (BLOCK_SIZE_K // 2) * stride_bk
+                 offs_k_packed += BLOCK_SIZE_K // 2
+                 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)
+  
+  
  @triton.jit
  def _fused_quant_shuffle_kernel(
      x_ptr, x_fp4_ptr, bs_ptr,
  ⋯ 46 unchanged lines
   
   
  _cache_asm = {}
- _cache_triton = {}
+ _cache_fused = {}
  _gemm_asm = None
  _warmup_done = False
   
  ⋯ 5 unchanged lines
      M, K = A.shape
      N = B_shuffle.shape[0]
   
-     use_triton = _HAS_PS and (M <= 16) and (K >= 2048)
+     # Fused for ALL M<=16 shapes (saves 1 kernel launch)
+     use_fused = (M <= 16)
   
-     # Warmup: use ASM path to initialize module
+     # Warmup: use ASM path to init aiter module
      if not _warmup_done:
          scale_n_valid = (K + 31) // 32
          SCALE_M = ((M + 255) // 256) * 256
  ⋯ 30 unchanged lines
                  pass
          return result
   
-     if use_triton:
-         # --- Preshuffle-scales GEMM: reads B_scale_sh directly ---
+     if use_fused:
+         # --- Fused quant+GEMM: single kernel launch ---
          key = (M, K, N)
-         c = _cache_triton.get(key)
+         c = _cache_fused.get(key)
          if c is None:
              K_packed = K // 2
              scale_n = (K + 31) // 32
              SCALE_N_B = ((scale_n + 7) // 8) * 8
   
-             # Quant config
-             BSM_q = triton.next_power_of_2(M)
-             BSK_q = 32
-             grid_q = (triton.cdiv(M, BSM_q), triton.cdiv(K, BSK_q))
-  
-             x_fp4 = torch.empty((M, K_packed), dtype=torch.uint8, device=A.device)
-             x_scales = torch.empty((M, scale_n), dtype=torch.uint8, device=A.device)
-  
-             # GEMM config
              BLOCK_SIZE_M = max(16, triton.next_power_of_2(M))
              BLOCK_SIZE_N = 128
-             BLOCK_SIZE_K = 256  # min for preshuffle_scales reshape
+             BLOCK_SIZE_K = 256
   
-             # SplitK for occupancy
              base_blocks = triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N)
              target_ksplit = max(1, 128 // max(1, base_blocks))
   
-             SPLITK_BLOCK_SIZE, BLOCK_SIZE_K, NUM_KSPLIT = _get_splitk_fn(
-                 K_packed, BLOCK_SIZE_K, target_ksplit
-             )
-  
-             # Ensure BLOCK_SIZE_K >= 256 for reshape
-             if BLOCK_SIZE_K < 256:
-                 BLOCK_SIZE_K = 256
-                 SPLITK_BLOCK_SIZE = 2 * K_packed
+             if target_ksplit > 1:
+                 SPLITK_BLOCK_SIZE, BLOCK_SIZE_K_adj, NUM_KSPLIT = _get_splitk_fn(
+                     K_packed, BLOCK_SIZE_K, target_ksplit
+                 )
+                 if BLOCK_SIZE_K_adj < 256:
+                     BLOCK_SIZE_K_adj = 256
+                     SPLITK_BLOCK_SIZE = 2 * K_packed
+                     NUM_KSPLIT = 1
+                 else:
+                     BLOCK_SIZE_K = BLOCK_SIZE_K_adj
+             else:
                  NUM_KSPLIT = 1
+                 SPLITK_BLOCK_SIZE = 2 * K_packed
   
              if NUM_KSPLIT > 1:
                  y_pp = torch.empty((NUM_KSPLIT, M, N), dtype=torch.float32, device=A.device)
  ⋯ 3 unchanged lines
   
              y = torch.empty((M, N), dtype=torch.bfloat16, device=A.device)
   
-             config = {
-                 "BLOCK_SIZE_M": BLOCK_SIZE_M,
-                 "BLOCK_SIZE_N": BLOCK_SIZE_N,
-                 "BLOCK_SIZE_K": BLOCK_SIZE_K,
-                 "GROUP_SIZE_M": 8,
-                 "NUM_KSPLIT": NUM_KSPLIT,
-                 "SPLITK_BLOCK_SIZE": SPLITK_BLOCK_SIZE,
-                 "num_warps": 4,
-                 "num_stages": 2,
-                 "waves_per_eu": 0,
-                 "matrix_instr_nonkdim": 32,
-                 "cache_modifier": ".ca",
-             }
+             total_blocks_raw = NUM_KSPLIT * triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N)
+             # Pad to multiple of 8 for correct XCD remapping
+             total_blocks = ((total_blocks_raw + 7) // 8) * 8
   
-             total_blocks = NUM_KSPLIT * triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N)
-  
-             # B_scale strides for shuffled layout (32 rows per N-group)
              bs_stride_n = 32 * SCALE_N_B
              bs_stride_k = 1
   
-             c = (K_packed, scale_n, SCALE_N_B, BSM_q, BSK_q, grid_q,
-                  x_fp4, x_scales, y, y_pp, config, total_blocks,
-                  bs_stride_n, bs_stride_k,
-                  A.stride(0), A.stride(1),
-                  x_fp4.stride(0), x_fp4.stride(1),
-                  x_scales.stride(0), x_scales.stride(1))
-             _cache_triton[key] = c
+             c = (K_packed, SCALE_N_B, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K,
+                  NUM_KSPLIT, SPLITK_BLOCK_SIZE,
+                  y, y_pp, total_blocks, bs_stride_n, bs_stride_k)
+             _cache_fused[key] = c
   
-         (K_packed, scale_n, SCALE_N_B, BSM_q, BSK_q, grid_q,
-          x_fp4, x_scales, y, y_pp, config, total_blocks,
-          bs_stride_n, bs_stride_k,
-          sa0, sa1, sf0, sf1, ss0, ss1) = c
+         (K_packed, SCALE_N_B, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K,
+          NUM_KSPLIT, SPLITK_BLOCK_SIZE,
+          y, y_pp, total_blocks, bs_stride_n, bs_stride_k) = c
   
-         # 1. Raw quant (produces x_fp4 + raw x_scales)
-         _quant_raw_kernel[grid_q](
-             A, x_fp4, x_scales,
-             sa0, sa1, sf0, sf1, ss0, ss1,
-             M, K,
-             BLOCK_M=BSM_q, BLOCK_K=BSK_q,
-             QUANT_BLOCK=32,
-             num_warps=1, waves_per_eu=0, num_stages=1,
-         )
-  
-         # 2. Preshuffle-scales GEMM — reads B_scale_sh directly, no unshuffle needed
          B_q_u8 = B_q.view(torch.uint8) if B_q.dtype != torch.uint8 else B_q
-         B_q_T = B_q_u8.T  # (K//2, N) non-contiguous view
-         # View B_scale_sh as uint8 to avoid Triton float8_e8m0fnu type error
+         B_q_T = B_q_u8.T
          B_scale_u8 = B_scale_sh.view(torch.uint8)
   
-         out_tensor = y if config["NUM_KSPLIT"] == 1 else y_pp
+         out_tensor = y if NUM_KSPLIT == 1 else y_pp
   
-         _ps_kernel[(total_blocks,)](
-             x_fp4,
-             B_q_T,
-             out_tensor,
-             x_scales,
-             B_scale_u8,
-             M, N, K_packed,
-             x_fp4.stride(0), x_fp4.stride(1),
+         _fused_quant_gemm_kernel[(total_blocks,)](
+             A, B_q_T, out_tensor, B_scale_u8,
+             M, N, K,
+             A.stride(0), A.stride(1),
              B_q_T.stride(0), B_q_T.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),
-             x_scales.stride(0), x_scales.stride(1),
+             0 if NUM_KSPLIT == 1 else y_pp.stride(0),
+             y.stride(0) if NUM_KSPLIT == 1 else y_pp.stride(1),
+             y.stride(1) if NUM_KSPLIT == 1 else y_pp.stride(2),
              bs_stride_n, bs_stride_k,
-             **config,
+             BLOCK_SIZE_M=BLOCK_SIZE_M,
+             BLOCK_SIZE_N=BLOCK_SIZE_N,
+             BLOCK_SIZE_K=BLOCK_SIZE_K,
+             GROUP_SIZE_M=8,
+             NUM_KSPLIT=NUM_KSPLIT,
+             SPLITK_BLOCK_SIZE=SPLITK_BLOCK_SIZE,
+             QUANT_BLOCK=32,
+             num_warps=4,
+             num_stages=2,
+             waves_per_eu=0,
          )
   
-         # 3. Reduce if SplitK
-         if config["NUM_KSPLIT"] > 1:
-             ACTUAL_KSPLIT = triton.cdiv(K_packed, (config["SPLITK_BLOCK_SIZE"] // 2))
-             grid_reduce = (
-                 triton.cdiv(M, 16),
-                 triton.cdiv(N, 64),
-             )
+         if NUM_KSPLIT > 1:
+             ACTUAL_KSPLIT = triton.cdiv(K_packed, (SPLITK_BLOCK_SIZE // 2))
+             grid_reduce = (triton.cdiv(M, 16), triton.cdiv(N, 64))
              _reduce_kernel[grid_reduce](
-                 y_pp, y,
-                 M, N,
+                 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,
-                 ACTUAL_KSPLIT,
-                 triton.next_power_of_2(config["NUM_KSPLIT"]),
+                 16, 64, ACTUAL_KSPLIT,
+                 triton.next_power_of_2(NUM_KSPLIT),
              )
   
          return y
   
      else:
-         # --- ASM GEMM path (v35) ---
+         # --- ASM GEMM path ---
          key = (M, K, N)
          c = _cache_asm.get(key)
          if c is None: