submission 503712

- #!POPCORN leaderboard nvfp4_group_gemm
- #!POPCORN gpu B200
-  
  import torch
  from task import input_t, output_t
- from torch.utils.cpp_extension import load_inline
+ from utils import make_match_reference
+ from typing import Tuple, Type, Union, Dict, Optional
   
- # ----------------------------
- # CUDA / C++ extension
- # ----------------------------
+ import cutlass
+ import cutlass.cute as cute
+ import cutlass.utils as utils
+ import cutlass.pipeline as pipeline
+ import cutlass.utils.blackwell_helpers as sm100_utils
+ import cutlass.utils.blockscaled_layout as blockscaled_utils
   
- CUDA_SRC_COMMON = r"""
- #include <cuda.h>
- #include <cudaTypedefs.h>
- #include <cuda_fp16.h>
- #include <stdint.h>
+ from cutlass.cute.nvgpu import cpasync, tcgen05
+ from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
+ from cutlass.cute.runtime import make_ptr
   
- #include <torch/library.h>
- #include <ATen/core/Tensor.h>
   
- constexpr int WARP_SIZE = 32;
- constexpr int MMA_K = 64;  // 32 bytes
+ # ============================================================
+ # Reference (兜底)
+ # ============================================================
+ sf_vec_size = 16
   
- // cache hint (from CUTLASS copy_sm90_desc)
- constexpr uint64_t EVICT_NORMAL = 0x1000000000000000;
- constexpr uint64_t EVICT_FIRST  = 0x12F0000000000000;
- constexpr uint64_t EVICT_LAST   = 0x14F0000000000000;
+ def ceil_div(a, b):
+     return (a + b - 1) // b
   
- __device__ inline constexpr uint64_t desc_encode(uint64_t x) { return (x & 0x3'FFFFULL) >> 4ULL; };
+ def to_blocked(input_matrix):
+     rows, cols = input_matrix.shape
+     n_row_blocks = ceil_div(rows, 128)
+     n_col_blocks = ceil_div(cols, 4)
+     padded_rows = n_row_blocks * 128
+     padded_cols = n_col_blocks * 4
+     if padded_rows != rows or padded_cols != cols:
+         padded = torch.nn.functional.pad(
+             input_matrix, (0, padded_cols - cols, 0, padded_rows - rows),
+             mode="constant", value=0
+         )
+     else:
+         padded = input_matrix
+     blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
+     rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
+     return rearranged.flatten()
   
- // elect one thread in warp
- __device__ uint32_t elect_sync() {
-   uint32_t pred = 0;
-   asm volatile(
-     "{\n\t"
-     ".reg .pred %%px;\n\t"
-     "elect.sync _|%%px, %1;\n\t"
-     "@%%px mov.s32 %0, 1;\n\t"
-     "}"
-     : "+r"(pred)
-     : "r"(0xFFFFFFFF)
-   );
-   return pred;
- }
+ def ref_kernel(data: input_t) -> output_t:
+     if len(data) == 3:
+         abc_tensors, sfasfb_tensors, problem_sizes = data
+     else:
+         abc_tensors, sfasfb_tensors, _, problem_sizes = data
   
- __device__ inline void mbarrier_init(int mbar_addr, int count) {
-   asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" :: "r"(mbar_addr), "r"(count));
- }
+     result_tensors = []
+     for (a_ref, b_ref, c_ref), (sfa_ref, sfb_ref), (m, n, k, l) in zip(
+         abc_tensors, sfasfb_tensors, problem_sizes
+     ):
+         for l_idx in range(l):
+             scale_a = to_blocked(sfa_ref[:, :, l_idx]).cuda()
+             scale_b = to_blocked(sfb_ref[:, :, l_idx]).cuda()
+             res = torch._scaled_mm(
+                 a_ref[:, :, l_idx].view(torch.float4_e2m1fn_x2),
+                 b_ref[:, :, l_idx].transpose(0, 1).view(torch.float4_e2m1fn_x2),
+                 scale_a,
+                 scale_b,
+                 bias=None,
+                 out_dtype=torch.float16,
+             )
+             c_ref[:, :, l_idx] = res
+         result_tensors.append(c_ref)
+     return result_tensors
   
- // parity wait
- __device__ void mbarrier_wait(int mbar_addr, int phase) {
-   uint32_t ticks = 0x989680;
-   asm volatile(
-     "{\n\t"
-     ".reg .pred P1;\n\t"
-     "LAB_WAIT:\n\t"
-     "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1, %2;\n\t"
-     "@P1 bra.uni DONE;\n\t"
-     "bra.uni LAB_WAIT;\n\t"
-     "DONE:\n\t"
-     "}"
-     :: "r"(mbar_addr), "r"(phase), "r"(ticks)
-   );
- }
+ check_implementation = make_match_reference(ref_kernel, rtol=1e-3, atol=1e-3)
   
- // TMA bulk copy (1D)
- __device__ inline
- void tma_gmem2smem(int dst, const void *src, int size, int mbar_addr, uint64_t cache_policy) {
-   asm volatile(
-     "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint "
-     "[%0], [%1], %2, [%3], %4;"
-     :: "r"(dst), "l"(src), "r"(size), "r"(mbar_addr), "l"(cache_policy)
-   );
- }
   
- // TMA tensor 3D copy
- __device__ inline
- void tma_3d_gmem2smem(int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr, uint64_t cache_policy) {
-   asm volatile(
-     "cp.async.bulk.tensor.3d.shared::cta.global.mbarrier::complete_tx::bytes.cta_group::1.L2::cache_hint "
-     "[%0], [%1, {%2, %3, %4}], [%5], %6;"
-     :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(mbar_addr), "l"(cache_policy)
-     : "memory"
-   );
- }
+ # ============================================================
+ # “冠军式”改动：2D cluster multicast + 更激进 stage cap
+ # ============================================================
+ _TMA_CACHE_EVICT_NORMAL = 0x1000000000000000
+ _TMA_CACHE_EVICT_FIRST  = 0x12F0000000000000
+ _TMA_CACHE_EVICT_LAST   = 0x14F0000000000000
   
- // tcgen05 cp scale (nvfp4)
- __device__ inline
- void tcgen05_cp_nvfp4(int taddr, uint64_t s_desc) {
-   asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;" :: "r"(taddr), "l"(s_desc));
- }
+ # (M,N,K) -> (tile_mn, cluster_mn, prefetch_dist, cache_policy, ab_stage_cap, max_active_clusters)
+ config_map: Dict[Tuple[int, int, int], Tuple[Tuple[int, int], Tuple[int, int], int, int, int, int]] = {}
   
- // tcgen05 mma (nvfp4 blockscale)
- __device__ inline
- void tcgen05_mma_nvfp4(
-   uint64_t a_desc,
-   uint64_t b_desc,
-   uint32_t i_desc,
-   int scale_A_tmem,
-   int scale_B_tmem,
-   int enable_input_d
- ) {
-   const int d_tmem = 0;
-   asm volatile(
-     "{\n\t"
-     ".reg .pred p;\n\t"
-     "setp.ne.b32 p, %6, 0;\n\t"
-     "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%4], [%5], p;\n\t"
-     "}"
-     :: "r"(d_tmem), "l"(a_desc), "l"(b_desc), "r"(i_desc),
-        "r"(scale_A_tmem), "r"(scale_B_tmem), "r"(enable_input_d)
-   );
- }
+ def _pick_tile_m(m: int) -> Optional[int]:
+     # 你如果知道你当前 mma 支持哪些 tile_m，就把候选列表换成你确认可用的
+     for tm in (64, 32, 16, 8):
+         if m % tm == 0:
+             return tm
+     return None
   
- // Minimal tcgen05 ld: need 32 regs for 16x256b.x8
- struct SHAPE {
-   static constexpr char _16x256b[] = ".16x256b";
- };
+ def _pick_tile_n(n: int) -> Optional[int]:
+     # 先 128 稳（很多情况下更容易保持驻留 + 更好 pipeline）
+     return 128 if (n % 128 == 0) else None
   
- struct NUM {
-   static constexpr char x8[] = ".x8";
- };
+ def _pick_cluster_shape(m_tiles: int, n_tiles: int) -> Tuple[int, int]:
+     """
+     核心改动：优先 2D cluster -> 同时复用 A/SFA（沿 N）和 B/SFB（沿 M）
+     - A/SFA multicast：mcast_mode=2
+     - B/SFB multicast：mcast_mode=1
+     """
+     # 优先 (2,2)：两边都能复用，通常对 group gemm 更有效
+     if m_tiles >= 2 and n_tiles >= 2:
+         return (2, 2)
+     # 其次 (4,1)：复用 B/SFB（N 很大，B 很重）
+     if m_tiles >= 4:
+         return (4, 1)
+     # 其次 (1,4)：复用 A/SFA（如果你确定 A/SFA 是瓶颈可用）
+     if n_tiles >= 4:
+         return (1, 4)
+     if m_tiles >= 2:
+         return (2, 1)
+     if n_tiles >= 2:
+         return (1, 2)
+     return (1, 1)
   
- template <const char *SHAPE_, const char *NUM_>
- __device__ inline
- void tcgen05_ld_32regs(float *tmp, int row, int col) {
-   asm volatile("tcgen05.ld.sync.aligned%33%34.b32 "
-               "{ %0,  %1,  %2,  %3,  %4,  %5,  %6,  %7, "
-               "  %8,  %9, %10, %11, %12, %13, %14, %15, "
-               " %16, %17, %18, %19, %20, %21, %22, %23, "
-               " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];"
-               : "=f"(tmp[ 0]), "=f"(tmp[ 1]), "=f"(tmp[ 2]), "=f"(tmp[ 3]), "=f"(tmp[ 4]), "=f"(tmp[ 5]), "=f"(tmp[ 6]), "=f"(tmp[ 7]),
-                 "=f"(tmp[ 8]), "=f"(tmp[ 9]), "=f"(tmp[10]), "=f"(tmp[11]), "=f"(tmp[12]), "=f"(tmp[13]), "=f"(tmp[14]), "=f"(tmp[15]),
-                 "=f"(tmp[16]), "=f"(tmp[17]), "=f"(tmp[18]), "=f"(tmp[19]), "=f"(tmp[20]), "=f"(tmp[21]), "=f"(tmp[22]), "=f"(tmp[23]),
-                 "=f"(tmp[24]), "=f"(tmp[25]), "=f"(tmp[26]), "=f"(tmp[27]), "=f"(tmp[28]), "=f"(tmp[29]), "=f"(tmp[30]), "=f"(tmp[31])
-               : "r"((row << 16) | col), "C"(SHAPE_), "C"(NUM_));
- }
+ def _pick_ab_stage_cap(k: int) -> int:
+     # 大 K 更容易把 SMEM 吃爆导致每 SM 只能驻 1 CTA -> 直接 cap=2
+     if k >= 4096:
+         return 2
+     return 3
   
- __device__ inline
- void tcgen05_ld_16x256bx8(float *tmp, int row, int col) {
-   tcgen05_ld_32regs<SHAPE::_16x256b, NUM::x8>(tmp, row, col);
- }
+ def _pick_max_active_clusters(cluster_mn: Tuple[int,int]) -> int:
+     # cluster 越大，每 cluster 吞掉的 CTA 越多；适当提高上限让 scheduler 更积极铺满
+     cm, cn = cluster_mn
+     area = cm * cn
+     if area >= 4:
+         return 256
+     if area == 2:
+         return 224
+     return 192
   
- void check_cu(CUresult err) {
-   if (err == CUDA_SUCCESS) return;
-   const char *error_msg_ptr;
-   if (cuGetErrorString(err, &error_msg_ptr) != CUDA_SUCCESS)
-     error_msg_ptr = "unable to get error string";
-   TORCH_CHECK(false, "CUDA driver error: ", error_msg_ptr);
- }
+ def ensure_config(m: int, n: int, k: int) -> bool:
+     key = (m, n, k)
+     if key in config_map:
+         return True
   
- void init_AB_tmap(
-   CUtensorMap *tmap,
-   const char *ptr,
-   uint64_t global_height, uint64_t global_width,
-   uint32_t shared_height, uint32_t shared_width
- ) {
-   // NOTE: matches champion gemm's fp4 tensor map (16U4)
-   constexpr uint32_t rank = 3;
-   uint64_t globalDim[rank]       = {256, global_height, global_width / 256};
-   uint64_t globalStrides[rank-1] = {global_width / 2, 128};  // bytes
-   uint32_t boxDim[rank]          = {256, shared_height, shared_width / 256};
-   uint32_t elementStrides[rank]  = {1, 1, 1};
+     tm = _pick_tile_m(m)
+     tn = _pick_tile_n(n)
+     if tm is None or tn is None:
+         return False
   
-   auto err = cuTensorMapEncodeTiled(
-     tmap,
-     CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B,
-     rank,
-     (void *)ptr,
-     globalDim,
-     globalStrides,
-     boxDim,
-     elementStrides,
-     CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
-     CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B,
-     CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
-     CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE
-   );
-   check_cu(err);
- }
- """
+     m_tiles = m // tm
+     n_tiles = n // tn
+     cluster = _pick_cluster_shape(m_tiles, n_tiles)
   
- CUDA_SRC_GEMM = r"""
- template <int K, int BLOCK_M, int BLOCK_N, int BLOCK_K, int NUM_STAGES>
- __global__ __launch_bounds__(BLOCK_M + 2 * WARP_SIZE)
- void kernel_gemm(
-   const __grid_constant__ CUtensorMap A_tmap,
-   const __grid_constant__ CUtensorMap B_tmap,
-   const char *SFA_ptr,
-   const char *SFB_ptr,
-   half *C_ptr,
-   int M, int N
- ) {
-   const int tid = threadIdx.x;
-   const int bid = blockIdx.x;
+     prefetch_dist = 0
+     cache_policy = _TMA_CACHE_EVICT_FIRST
+     ab_stage_cap = _pick_ab_stage_cap(k)
+     max_active = _pick_max_active_clusters(cluster)
   
-   const int lane_id = tid % WARP_SIZE;
-   const int warp_id = tid / WARP_SIZE;
+     config_map[key] = ((tm, tn), cluster, prefetch_dist, cache_policy, ab_stage_cap, max_active)
+     return True
   
-   const int grid_m = M / BLOCK_M;
-   const int grid_n = N / BLOCK_N;
-   const int bid_m = bid / grid_n;
-   const int bid_n = bid % grid_n;
   
-   const int off_m = bid_m * BLOCK_M;
-   const int off_n = bid_n * BLOCK_N;
+ # ============================================================
+ # GroupGemm kernel（保持你之前通过正确性的结构）
+ # 只改：cluster=2D + stage cap=2/3
+ # ============================================================
+ class GroupGemm:
+     def __init__(self, mnk: Tuple[int, int, int]):
+         m, n, k = mnk
+         (mma_tiler_mn, cluster_shape_mn, prefetch_dist, cache_policy, ab_stage_cap, max_active_clusters) = config_map[(m, n, k)]
   
-   constexpr int NUM_WARPS = BLOCK_M / WARP_SIZE + 2; // BLOCK_M=128 => 6 warps
+         self.mma_tiler_mn = mma_tiler_mn
+         self.cluster_shape_mn = cluster_shape_mn
+         self.prefetch_dist_param = prefetch_dist
+         self.cache_policy = cache_policy
+         self.ab_stage_cap = ab_stage_cap
+         self.max_active_clusters = max_active_clusters
   
-   // smem
-   extern __shared__ __align__(1024) char smem_ptr[];
-   const int smem = static_cast<int>(__cvta_generic_to_shared(smem_ptr));
-   constexpr int A_size   = BLOCK_M * BLOCK_K / 2;  // fp4 packed
-   constexpr int B_size   = BLOCK_N * BLOCK_K / 2;
-   constexpr int SFA_size = 128 * BLOCK_K / 16;     // fp8 bytes, always 128 rows
-   constexpr int SFB_size = 128 * BLOCK_K / 16;
-   constexpr int STAGE_SIZE = A_size + B_size + SFA_size + SFB_size;
+         self.acc_dtype = cutlass.Float32
+         self.sf_vec_size = 16
   
-   // mbarriers
-   #pragma nv_diag_suppress static_var_with_dynamic_init
-   __shared__ int64_t mbars[NUM_STAGES * 2 + 1];
-   const int tma_mbar_addr = static_cast<int>(__cvta_generic_to_shared(mbars));
-   const int mma_mbar_addr = tma_mbar_addr + NUM_STAGES * 8;
-   const int mainloop_mbar_addr = mma_mbar_addr + NUM_STAGES * 8;
+         self.epilog_warp_id = (0, 1, 2, 3)
+         self.mma_warp_id = 4
+         self.tma_warp_id = 5
+         self.threads_per_cta = 32 * 6
   
-   // tmem columns for scales
-   constexpr int SFA_tmem = BLOCK_N;
-   constexpr int SFB_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K);
+         self.occupancy = 1
+         self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
   
-   if (warp_id == 0 && elect_sync()) {
-     for (int i = 0; i < NUM_STAGES * 2 + 1; i++)
-       mbarrier_init(tma_mbar_addr + i * 8, 1);
-     asm volatile("fence.mbarrier_init.release.cluster;");
-   } else if (warp_id == 1) {
-     // allocate tmem: BLOCK_N*2 columns
-     asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;"
-                  :: "r"(smem), "r"(BLOCK_N * 2));
-   }
-   __syncthreads();
+         # 先保持 ONE CTA（更稳）；你如果已验证某些 tile 支持 TWO，可在这里加规则
+         self.use_2cta_instrs = False
+         self.cta_group = tcgen05.CtaGroup.ONE
   
-   constexpr int num_iters = K / BLOCK_K;
+         self.epilog_sync_barrier = pipeline.NamedBarrier(
+             barrier_id=1,
+             num_threads=32 * 4,
+         )
+         self.tmem_alloc_barrier = pipeline.NamedBarrier(
+             barrier_id=2,
+             num_threads=32 * 5,
+         )
   
-   // TMA warp
-   if (warp_id == NUM_WARPS - 2 && elect_sync()) {
-     constexpr uint64_t cache_A = EVICT_LAST;
-     constexpr uint64_t cache_B = EVICT_FIRST;
+         self.buffer_align_bytes = 1024
   
-     auto issue_tma = [&](int iter_k, int stage_id) {
-       const int mbar_addr = tma_mbar_addr + stage_id * 8;
-       const int A_smem    = smem + stage_id * STAGE_SIZE;
-       const int B_smem    = A_smem + A_size;
-       const int SFA_smem  = B_smem + B_size;
-       const int SFB_smem  = SFA_smem + SFA_size;
+     @staticmethod
+     def _compute_ab_c_stages_capped(
+         tiled_mma: cute.TiledMma,
+         mma_tiler_mnk: Tuple[int, int, int],
+         a_dtype: Type[cutlass.Numeric],
+         b_dtype: Type[cutlass.Numeric],
+         epi_tile: cute.Tile,
+         c_dtype: Type[cutlass.Numeric],
+         c_layout: utils.LayoutEnum,
+         sf_dtype: Type[cutlass.Numeric],
+         sf_vec_size: int,
+         smem_capacity: int,
+         occupancy: int,
+         ab_stage_cap: int,
+     ) -> Tuple[int, int]:
+         num_c_stage = 2
   
-       const int off_k = iter_k * BLOCK_K;
+         a_smem_1 = sm100_utils.make_smem_layout_a(tiled_mma, mma_tiler_mnk, a_dtype, 1)
+         b_smem_1 = sm100_utils.make_smem_layout_b(tiled_mma, mma_tiler_mnk, b_dtype, 1)
+         sfa_smem_1 = blockscaled_utils.make_smem_layout_sfa(tiled_mma, mma_tiler_mnk, sf_vec_size, 1)
+         sfb_smem_1 = blockscaled_utils.make_smem_layout_sfb(tiled_mma, mma_tiler_mnk, sf_vec_size, 1)
+         c_smem_1 = sm100_utils.make_smem_layout_epi(c_dtype, c_layout, epi_tile, 1)
   
-       // fp4 tiles
-       tma_3d_gmem2smem(A_smem, &A_tmap, 0, off_m, off_k / 256, mbar_addr, cache_A);
-       tma_3d_gmem2smem(B_smem, &B_tmap, 0, off_n, off_k / 256, mbar_addr, cache_B);
+         ab_bytes = (
+             cute.size_in_bytes(a_dtype, a_smem_1)
+             + cute.size_in_bytes(b_dtype, b_smem_1)
+             + cute.size_in_bytes(sf_dtype, sfa_smem_1)
+             + cute.size_in_bytes(sf_dtype, sfb_smem_1)
+         )
+         mbar_bytes = 1024
+         c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_1)
+         c_bytes = c_bytes_per_stage * num_c_stage
   
-       // scale layout assumed: [mn/128, rest_k, 32,4,4] contiguous in memory (512 bytes each block)
-       const int rest_k = K / 16 / 4; // = K/64
-       const char *SFA_src = SFA_ptr + ((off_m / 128) * rest_k + off_k / (16 * 4)) * 512;
-       const char *SFB_src = SFB_ptr + ((off_n / 128) * rest_k + off_k / (16 * 4)) * 512;
+         num_ab_stage = (smem_capacity // occupancy - (mbar_bytes + c_bytes)) // ab_bytes
+         if num_ab_stage < 2:
+             num_ab_stage = 2
+         if num_ab_stage > ab_stage_cap:
+             num_ab_stage = ab_stage_cap
   
-       tma_gmem2smem(SFA_smem, SFA_src, SFA_size, mbar_addr, cache_A);
-       tma_gmem2smem(SFB_smem, SFB_src, SFB_size, mbar_addr, cache_B);
+         num_c_stage += (
+             smem_capacity
+             - occupancy * ab_bytes * num_ab_stage
+             - occupancy * (mbar_bytes + c_bytes)
+         ) // (occupancy * c_bytes_per_stage)
   
-       asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"
-                    :: "r"(mbar_addr), "r"(STAGE_SIZE) : "memory");
-     };
+         return num_ab_stage, num_c_stage
   
-     for (int iter_k = 0; iter_k < NUM_STAGES; iter_k++)
-       issue_tma(iter_k, iter_k);
+     @staticmethod
+     def _compute_grid(
+         c: cute.Tensor,
+         cta_tile_shape_mnk: Tuple[int, int, int],
+         cluster_shape_mn: Tuple[int, int],
+         max_active_clusters: cutlass.Constexpr,
+     ):
+         c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0))
+         gc = cute.zipped_divide(c, tiler=c_shape)
+         num_ctas_mnl = gc[(0, (None, None, None))].shape
+         cluster_shape_mnl = (*cluster_shape_mn, 1)
+         tile_sched_params = utils.PersistentTileSchedulerParams(num_ctas_mnl, cluster_shape_mnl)
+         grid = utils.StaticPersistentTileScheduler.get_grid_shape(tile_sched_params, max_active_clusters)
+         return tile_sched_params, grid
   
-     for (int iter_k = NUM_STAGES; iter_k < num_iters; iter_k++) {
-       const int stage_id = iter_k % NUM_STAGES;
-       const int mma_phase = (iter_k / NUM_STAGES - 1) % 2;
-       mbarrier_wait(mma_mbar_addr + stage_id * 8, mma_phase);
-       issue_tma(iter_k, stage_id);
-     }
-   }
-   // MMA warp
-   else if (warp_id == NUM_WARPS - 1 && elect_sync()) {
-     // fp4 MMA uses MMA_M=128 always
-     constexpr uint32_t i_desc = (1U << 7U)   // atype=E2M1
-                              | (1U << 10U)  // btype=E2M1
-                              | ((uint32_t)BLOCK_N >> 3U << 17U)  // MMA_N
-                              | ((uint32_t)128     >> 7U << 27U); // MMA_M
+     @cute.jit
+     def __call__(
+         self,
+         a_ptr: cute.Pointer,
+         b_ptr: cute.Pointer,
+         sfa_ptr: cute.Pointer,
+         sfb_ptr: cute.Pointer,
+         c_ptr: cute.Pointer,
+         problem_size: cutlass.Constexpr,
+     ):
+         m, n, k, l = problem_size
   
-     for (int iter_k = 0; iter_k < num_iters; iter_k++) {
-       const int stage_id = iter_k % NUM_STAGES;
-       const int tma_phase = (iter_k / NUM_STAGES) % 2;
-       mbarrier_wait(tma_mbar_addr + stage_id * 8, tma_phase);
+         a_tensor = cute.make_tensor(a_ptr, cute.make_layout((m, k, l), stride=(k, 1, m * k)))
+         b_tensor = cute.make_tensor(b_ptr, cute.make_layout((n, k, l), stride=(k, 1, n * k)))
+         c_tensor = cute.make_tensor(c_ptr, cute.make_layout((m, n, l), stride=(n, 1, m * n)))
   
-       const int A_smem   = smem + stage_id * STAGE_SIZE;
-       const int B_smem   = A_smem + A_size;
-       const int SFA_smem = B_smem + B_size;
-       const int SFB_smem = SFA_smem + SFA_size;
+         self.a_dtype = cutlass.Float4E2M1FN
+         self.b_dtype = cutlass.Float4E2M1FN
+         self.sf_dtype = cutlass.Float8E4M3FN
+         self.c_dtype = cutlass.Float16
   
-       auto make_desc_AB = [](int addr) -> uint64_t {
-         const int SBO = 8 * 128;
-         return desc_encode(addr) | (desc_encode(SBO) << 32ULL) | (1ULL << 46ULL) | (2ULL << 61ULL);
-       };
-       auto make_desc_SF = [](int addr) -> uint64_t {
-         const int SBO = 8 * 16;
-         return desc_encode(addr) | (desc_encode(SBO) << 32ULL) | (1ULL << 46ULL);
-       };
+         self.a_major_mode = utils.LayoutEnum.from_tensor(a_tensor).mma_major_mode()
+         self.b_major_mode = utils.LayoutEnum.from_tensor(b_tensor).mma_major_mode()
+         self.c_layout = utils.LayoutEnum.from_tensor(c_tensor)
   
-       constexpr uint64_t SF_desc = make_desc_SF(0);
-       const uint64_t SFA_desc = SF_desc + ((uint64_t)SFA_smem >> 4ULL);
-       const uint64_t SFB_desc = SF_desc + ((uint64_t)SFB_smem >> 4ULL);
+         # SF layout（reordered SF）
+         sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(a_tensor.shape, self.sf_vec_size)
+         sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(b_tensor.shape, self.sf_vec_size)
+         sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)
+         sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout)
   
-       // smem->tmem for scales
-       for (int k = 0; k < BLOCK_K / MMA_K; k++) {
-         uint64_t sfa_desc = SFA_desc + (uint64_t)k * (512ULL >> 4ULL);
-         uint64_t sfb_desc = SFB_desc + (uint64_t)k * (512ULL >> 4ULL);
-         tcgen05_cp_nvfp4(SFA_tmem + k * 4, sfa_desc);
-         tcgen05_cp_nvfp4(SFB_tmem + k * 4, sfb_desc);
-       }
+         self.mma_inst_shape_mn = (self.mma_tiler_mn[0], self.mma_tiler_mn[1])
+         tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
+             self.a_dtype, self.a_major_mode,
+             self.b_major_mode,
+             self.sf_dtype, self.sf_vec_size,
+             self.cta_group,
+             self.mma_inst_shape_mn,
+         )
   
-       // mma over BLOCK_K
-       for (int k1 = 0; k1 < BLOCK_K / 256; k1++)
-         for (int k2 = 0; k2 < 256 / MMA_K; k2++) {
-           uint64_t a_desc = make_desc_AB(A_smem + k1 * BLOCK_M * 128 + k2 * 32);
-           uint64_t b_desc = make_desc_AB(B_smem + k1 * BLOCK_N * 128 + k2 * 32);
+         mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
+         mma_inst_tile_k = 4
+         self.mma_tiler = (
+             self.mma_inst_shape_mn[0],
+             self.mma_inst_shape_mn[1],
+             mma_inst_shape_k * mma_inst_tile_k,
+         )
+         self.cta_tile_shape_mnk = (
+             self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
+             self.mma_tiler[1],
+             self.mma_tiler[2],
+         )
   
-           int k_sf = k1 * 4 + k2;
-           const int scale_A_tmem = SFA_tmem + k_sf * 4 + (bid_m % (128 / BLOCK_M)) * (BLOCK_M / 32);
-           const int scale_B_tmem = SFB_tmem + k_sf * 4 + (bid_n % (128 / BLOCK_N)) * (BLOCK_N / 32);
+         self.cluster_layout_vmnk = cute.tiled_divide(
+             cute.make_layout((*self.cluster_shape_mn, 1)),
+             (tiled_mma.thr_id.shape,),
+         )
   
-           const int enable_input_d = (k1 == 0 && k2 == 0) ? iter_k : 1;
-           tcgen05_mma_nvfp4(a_desc, b_desc, i_desc, scale_A_tmem, scale_B_tmem, enable_input_d);
-         }
+         # 这两个维度的含义保持跟 DualGemm 一致：
+         # - A/SFA multicast 用 shape[2]
+         # - B/SFB multicast 用 shape[1]
+         self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2])
+         self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1])
+         self.is_a_mcast = self.num_mcast_ctas_a > 1
+         self.is_b_mcast = self.num_mcast_ctas_b > 1
   
-       asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
-                    :: "r"(mma_mbar_addr + stage_id * 8) : "memory");
-     }
+         self.epi_tile = sm100_utils.compute_epilogue_tile_shape(
+             self.cta_tile_shape_mnk,
+             self.use_2cta_instrs,
+             self.c_layout,
+             self.c_dtype,
+         )
   
-     asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
-                  :: "r"(mainloop_mbar_addr) : "memory");
-   }
-   // epilogue warps (write row-major C: [M,N])
-   else if (tid < BLOCK_M) {
-     mbarrier_wait(mainloop_mbar_addr, 0);
-     asm volatile("tcgen05.fence::after_thread_sync;");
+         self.num_acc_stage = 1
+         self.num_ab_stage, self.num_c_stage = self._compute_ab_c_stages_capped(
+             tiled_mma, self.mma_tiler,
+             self.a_dtype, self.b_dtype,
+             self.epi_tile,
+             self.c_dtype, self.c_layout,
+             self.sf_dtype, self.sf_vec_size,
+             self.smem_capacity, self.occupancy,
+             self.ab_stage_cap,
+         )
   
-     // C is row-major (N-major in their naming)
-     for (int m = 0; m < 32 / 16; m++) {
-       float tmp[BLOCK_N / 2]; // BLOCK_N=64 => 32 floats
-       tcgen05_ld_16x256bx8(tmp, warp_id * 32 + m * 16, 0);
-       asm volatile("tcgen05.wait::ld.sync.aligned;");
+         self.prefetch_dist = self.num_ab_stage if (self.prefetch_dist_param is None) else self.prefetch_dist_param
   
-       for (int i = 0; i < BLOCK_N / 8; i++) {
-         const int row = off_m + warp_id * 32 + m * 16 + lane_id / 4;
-         const int col = off_n + i * 8 + (lane_id % 4) * 2;
+         self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage)
+         self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage)
+         self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(tiled_mma, self.mma_tiler, self.sf_vec_size, self.num_ab_stage)
+         self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(tiled_mma, self.mma_tiler, self.sf_vec_size, self.num_ab_stage)
+         self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi(self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage)
   
-         reinterpret_cast<half2 *>(C_ptr + (row + 0) * N + col)[0] =
-           __float22half2_rn({tmp[i * 4 + 0], tmp[i * 4 + 1]});
-         reinterpret_cast<half2 *>(C_ptr + (row + 8) * N + col)[0] =
-           __float22half2_rn({tmp[i * 4 + 2], tmp[i * 4 + 3]});
-       }
-     }
+         a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id)
+         b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id)
+         sf_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id)
+         sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, tiled_mma.thr_id)
   
-     asm volatile("bar.sync 1, %0;" :: "r"(BLOCK_M) : "memory");
-     if (warp_id == 0) {
-       asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;"
-                    :: "r"(0), "r"(BLOCK_N * 2));
-     }
-   }
- }
+         a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
+         b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
+         sfa_smem_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0))
+         sfb_smem_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0))
   
- template <int K>
- static inline void launch_one(
-   const at::Tensor& A,
-   const at::Tensor& B,
-   const at::Tensor& SFA_blk,
-   const at::Tensor& SFB_blk,
-         at::Tensor& C
- ) {
-   constexpr int BLOCK_M = 128;
-   constexpr int BLOCK_N = 64;
-   constexpr int BLOCK_K = 256;
-   constexpr int NUM_STAGES = 6;
+         tma_atom_a, mA_mkl = cute.nvgpu.make_tiled_tma_atom_A(
+             a_op, a_tensor, a_smem_layout, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape
+         )
+         tma_atom_b, mB_nkl = cute.nvgpu.make_tiled_tma_atom_B(
+             b_op, b_tensor, b_smem_layout, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape
+         )
+         tma_atom_sfa, mSFA_mkl = cute.nvgpu.make_tiled_tma_atom_A(
+             sf_op, sfa_tensor, sfa_smem_layout, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape,
+             internal_type=cutlass.Int16
+         )
+         tma_atom_sfb, mSFB_nkl = cute.nvgpu.make_tiled_tma_atom_B(
+             sfb_op, sfb_tensor, sfb_smem_layout, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape,
+             internal_type=cutlass.Int16
+         )
   
-   const int M = A.size(0);
-   const int N = B.size(0);
+         epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0))
+         tma_atom_c, _ = cpasync.make_tiled_tma_atom(
+             cpasync.CopyBulkTensorTileS2GOp(),
+             c_tensor,
+             epi_smem_layout,
+             self.epi_tile,
+         )
   
-   auto A_ptr   = reinterpret_cast<const char*>(A.data_ptr());
-   auto B_ptr   = reinterpret_cast<const char*>(B.data_ptr());
-   auto SFA_ptr = reinterpret_cast<const char*>(SFA_blk.data_ptr());
-   auto SFB_ptr = reinterpret_cast<const char*>(SFB_blk.data_ptr());
-   auto C_ptr   = reinterpret_cast<half*>(C.data_ptr());
+         # grid
+         self.tile_sched_params, grid = self._compute_grid(
+             c_tensor, self.cta_tile_shape_mnk, self.cluster_shape_mn, self.max_active_clusters
+         )
   
-   CUtensorMap A_tmap, B_tmap;
-   init_AB_tmap(&A_tmap, A_ptr, M, K, BLOCK_M, BLOCK_K);
-   init_AB_tmap(&B_tmap, B_ptr, N, K, BLOCK_N, BLOCK_K);
+         atom_thr_size = cute.size(tiled_mma.thr_id.shape)
+         a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout)
+         b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout)
+         sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout)
+         sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout)
+         self.num_tma_load_bytes = (a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size) * atom_thr_size
   
-   const int grid = (M / BLOCK_M) * (N / BLOCK_N);
-   const int tb_size = BLOCK_M + 2 * WARP_SIZE;
+         @cute.struct
+         class SharedStorage:
+             ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]
+             acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
+             tmem_dealloc_mbar_ptr: cutlass.Int64
+             tmem_holding_buf: cutlass.Int32
+             sC: cute.struct.Align[
+                 cute.struct.MemRange[self.c_dtype, cute.cosize(self.c_smem_layout_staged.outer)],
+                 self.buffer_align_bytes
+             ]
+             sA: cute.struct.Align[
+                 cute.struct.MemRange[self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)],
+                 self.buffer_align_bytes
+             ]
+             sB: cute.struct.Align[
+                 cute.struct.MemRange[self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)],
+                 self.buffer_align_bytes
+             ]
+             sSFA: cute.struct.Align[
+                 cute.struct.MemRange[self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged)],
+                 self.buffer_align_bytes
+             ]
+             sSFB: cute.struct.Align[
+                 cute.struct.MemRange[self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)],
+                 self.buffer_align_bytes
+             ]
   
-   const int AB_size   = (BLOCK_M + BLOCK_N) * (BLOCK_K / 2);
-   const int SFAB_size = 128 * (BLOCK_K / 16) * 2;
-   const int smem_size = (AB_size + SFAB_size) * NUM_STAGES;
+         self.shared_storage = SharedStorage
   
-   auto ker = kernel_gemm<K, BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES>;
-   if (smem_size > 48'000)
-     cudaFuncSetAttribute(ker, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+         self.kernel(
+             tiled_mma,
+             tma_atom_a, mA_mkl,
+             tma_atom_b, mB_nkl,
+             tma_atom_sfa, mSFA_mkl,
+             tma_atom_sfb, mSFB_nkl,
+             tma_atom_c, c_tensor,
+             self.cluster_layout_vmnk,
+             self.a_smem_layout_staged,
+             self.b_smem_layout_staged,
+             self.sfa_smem_layout_staged,
+             self.sfb_smem_layout_staged,
+             self.c_smem_layout_staged,
+             self.epi_tile,
+             self.tile_sched_params,
+         ).launch(
+             grid=grid,
+             block=[self.threads_per_cta, 1, 1],
+             cluster=(*self.cluster_shape_mn, 1),
+             min_blocks_per_mp=1,
+         )
+         return
   
-   ker<<<grid, tb_size, smem_size>>>(A_tmap, B_tmap, SFA_ptr, SFB_ptr, C_ptr, M, N);
- }
+     @cute.kernel
+     def kernel(
+         self,
+         tiled_mma: cute.TiledMma,
+         tma_atom_a: cute.CopyAtom, mA_mkl: cute.Tensor,
+         tma_atom_b: cute.CopyAtom, mB_nkl: cute.Tensor,
+         tma_atom_sfa: cute.CopyAtom, mSFA_mkl: cute.Tensor,
+         tma_atom_sfb: cute.CopyAtom, mSFB_nkl: cute.Tensor,
+         tma_atom_c: cute.CopyAtom, mC_mnl: cute.Tensor,
+         cluster_layout_vmnk: cute.Layout,
+         a_smem_layout_staged: cute.ComposedLayout,
+         b_smem_layout_staged: cute.ComposedLayout,
+         sfa_smem_layout_staged: cute.Layout,
+         sfb_smem_layout_staged: cute.Layout,
+         c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout],
+         epi_tile: cute.Tile,
+         tile_sched_params: utils.PersistentTileSchedulerParams,
+     ):
+         warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
   
- at::Tensor gemm(
-   const at::Tensor& A,
-   const at::Tensor& B,
-   const at::Tensor& SFA_blk,
-   const at::Tensor& SFB_blk,
-         at::Tensor& C
- ) {
-   const int K = A.size(1) * 2;
+         if warp_idx == self.tma_warp_id:
+             cpasync.prefetch_descriptor(tma_atom_a)
+             cpasync.prefetch_descriptor(tma_atom_b)
+             cpasync.prefetch_descriptor(tma_atom_sfa)
+             cpasync.prefetch_descriptor(tma_atom_sfb)
+             cpasync.prefetch_descriptor(tma_atom_c)
   
-   if (false) {}
-   else if (K == 7168)  launch_one<7168>(A, B, SFA_blk, SFB_blk, C);
-   else if (K == 4096)  launch_one<4096>(A, B, SFA_blk, SFB_blk, C);
-   else if (K == 2048)  launch_one<2048>(A, B, SFA_blk, SFB_blk, C);
-   else if (K == 1536)  launch_one<1536>(A, B, SFA_blk, SFB_blk, C);
-   else if (K == 2304)  launch_one<2304>(A, B, SFA_blk, SFB_blk, C);
-   else if (K == 512)   launch_one<512>(A, B, SFA_blk, SFB_blk, C);
-   else if (K == 256)   launch_one<256>(A, B, SFA_blk, SFB_blk, C);
-   else {
-     // leave to python fallback for correctness
-   }
+         bidx, bidy, bidz = cute.arch.block_idx()
+         mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape)
+         is_leader_cta = mma_tile_coord_v == 0
   
-   return C;
- }
+         cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
+         block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
   
- TORCH_LIBRARY(group_gemm_mod, m) {
-   m.def("gemm(Tensor A, Tensor B, Tensor SFA_blk, Tensor SFB_blk, Tensor(a!) C) -> Tensor");
-   m.impl("gemm", &gemm);
- }
- """
+         tidx, _, _ = cute.arch.thread_idx()
   
- load_inline(
-     name="group_gemm_ext",
-     cpp_sources="",
-     cuda_sources=CUDA_SRC_COMMON + CUDA_SRC_GEMM,
-     functions=None,
-     extra_cuda_cflags=[
-         "-O3",
-         "-gencode=arch=compute_100a,code=sm_100a",
-         "--use_fast_math",
-         "--expt-relaxed-constexpr",
-         "--relocatable-device-code=false",
-         "-lineinfo",
-         "-Xptxas=-v",
-     ],
-     extra_ldflags=["-lcuda"],
-     with_cuda=True,
-     verbose=False,
-     is_python_module=False,
-     no_implicit_headers=True,
- )
+         smem = utils.SmemAllocator()
+         storage = smem.allocate(self.shared_storage)
   
- gemm = torch.ops.group_gemm_mod.gemm
+         ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+         num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
+         ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_tma_producer)
   
- # ----------------------------
- # Python helpers
- # ----------------------------
+         ab_pipeline = pipeline.PipelineTmaUmma.create(
+             barrier_storage=storage.ab_full_mbar_ptr.data_ptr(),
+             num_stages=self.num_ab_stage,
+             producer_group=ab_pipeline_producer_group,
+             consumer_group=ab_pipeline_consumer_group,
+             tx_count=self.num_tma_load_bytes,
+             cta_layout_vmnk=cluster_layout_vmnk,
+             defer_sync=True,
+         )
   
- def ceil_div(a: int, b: int) -> int:
-     return (a + b - 1) // b
+         acc_pipeline = pipeline.PipelineUmmaAsync.create(
+             barrier_storage=storage.acc_full_mbar_ptr.data_ptr(),
+             num_stages=self.num_acc_stage,
+             producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+             consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 4),
+             cta_layout_vmnk=cluster_layout_vmnk,
+             defer_sync=True,
+         )
   
- @torch.no_grad()
- def reorder_sf_blocked_fast(sf_mn_k16: torch.Tensor, mn: int, k: int, mn_pad: int) -> torch.Tensor:
-     """
-     FAST reshape/permute version.
-     input:  [mn, k//16]
-     output: [mn_pad//128, (k//16)//4, 32, 4, 4]
-     Layout mapping (matches your original scatter):
-       mm   = i//128
-       mm32 = i%32
-       mm4  = (i%128)//32
-       kk   = j//4
-       kk4  = j%4
-     => out[mm, kk, mm32, mm4, kk4] = in[i, j]
-     """
-     assert sf_mn_k16.dim() == 2
-     assert sf_mn_k16.shape[0] == mn
-     sf_k = k // 16
-     assert sf_mn_k16.shape[1] == sf_k
-     assert (mn_pad % 128) == 0
-     assert (sf_k % 4) == 0
+         tmem = utils.TmemAllocator(
+             storage.tmem_holding_buf,
+             barrier_for_retrieve=self.tmem_alloc_barrier,
+             allocator_warp_id=self.epilog_warp_id[0],
+             is_two_cta=False,
+             two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
+         )
   
-     device = sf_mn_k16.device
-     dtype = sf_mn_k16.dtype
+         pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
   
-     # pad rows to mn_pad (cheap, contiguous)
-     if mn_pad == mn:
-         sf_pad = sf_mn_k16
-     else:
-         sf_pad = torch.zeros((mn_pad, sf_k), device=device, dtype=dtype)
-         sf_pad[:mn, :] = sf_mn_k16
+         sC = storage.sC.get_tensor(c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner)
+         sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner)
+         sB = storage.sB.get_tensor(b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner)
+         sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged)
+         sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged)
   
-     rest_m = mn_pad // 128
-     rest_k = sf_k // 4
+         # ★关键：2D multicast
+         a_full_mcast_mask = None
+         b_full_mcast_mask = None
+         sfa_full_mcast_mask = None
+         sfb_full_mcast_mask = None
+         if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast):
+             a_full_mcast_mask = cpasync.create_tma_multicast_mask(
+                 cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
+             )
+             b_full_mcast_mask = cpasync.create_tma_multicast_mask(
+                 cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1
+             )
+             sfa_full_mcast_mask = cpasync.create_tma_multicast_mask(
+                 cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
+             )
+             sfb_full_mcast_mask = cpasync.create_tma_multicast_mask(
+                 cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1
+             )
   
-     # reshape: [rest_m, 128, rest_k, 4]
-     # then split 128 -> [4, 32] with order (mm4, mm32)
-     # then permute to [rest_m, rest_k, 32, 4, 4]
-     out = (
-         sf_pad
-         .view(rest_m, 128, rest_k, 4)
-         .view(rest_m, 4, 32, rest_k, 4)
-         .permute(0, 3, 2, 1, 4)
-         .contiguous()
-     )
-     return out
+         # tile tensors
+         gA_mkl = cute.local_tile(mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None))
+         gB_nkl = cute.local_tile(mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None))
+         gSFA_mkl = cute.local_tile(mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None))
+         gSFB_nkl = cute.local_tile(mSFB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None))
+         gC_mnl = cute.local_tile(mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None))
   
- @torch.no_grad()
- def pad_fp4_rows(a_fp4: torch.Tensor, m: int, mn_pad: int) -> torch.Tensor:
-     """
-     Pad fp4 packed tensor [M, K//2, L] on rows to mn_pad with zeros (bitwise),
-     returning float4_e2m1fn_x2 tensor.
-     """
-     a_u8 = a_fp4.contiguous().view(torch.uint8)
-     out_u8 = torch.zeros((mn_pad, a_u8.shape[1], a_u8.shape[2]), device=a_u8.device, dtype=torch.uint8)
-     out_u8[:m, :, :] = a_u8
-     return out_u8.view(torch.float4_e2m1fn_x2)
+         k_tile_cnt = cute.size(gA_mkl, mode=[3])
   
- @torch.no_grad()
- def torch_scaled_mm_fallback(a_fp4, b_fp4, sfa, sfb, c_out, m, n, k, l):
-     sf_vec_size = 16
+         thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
+         tCgA = thr_mma.partition_A(gA_mkl)
+         tCgB = thr_mma.partition_B(gB_nkl)
+         tCgSFA = thr_mma.partition_A(gSFA_mkl)
+         tCgSFB = thr_mma.partition_B(gSFB_nkl)
+         tCgC = thr_mma.partition_C(gC_mnl)
   
-     def to_block(input_matrix):
-         rows, cols = input_matrix.shape
-         n_row_blocks = ceil_div(rows, 128)
-         n_col_blocks = ceil_div(cols, 4)
-         padded_rows = n_row_blocks * 128
-         padded_cols = n_col_blocks * 4
-         if padded_rows != rows or padded_cols != cols:
-             padded = torch.nn.functional.pad(
-                 input_matrix, (0, padded_cols - cols, 0, padded_rows - rows),
-                 mode="constant", value=0
-             )
-         else:
-             padded = input_matrix
-         blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
-         rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
-         return rearranged.flatten()
+         a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape)
+         b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape)
   
-     for l_idx in range(l):
-         scale_a = to_block(sfa[:, :, l_idx]).cuda()
-         scale_b = to_block(sfb[:, :, l_idx]).cuda()
-         res = torch._scaled_mm(
-             a_fp4[:, :, l_idx].view(torch.float4_e2m1fn_x2),
-             b_fp4[:, :, l_idx].transpose(0, 1).view(torch.float4_e2m1fn_x2),
-             scale_a, scale_b,
-             bias=None,
-             out_dtype=torch.float16,
+         tAsA, tAgA = cpasync.tma_partition(
+             tma_atom_a, block_in_cluster_coord_vmnk[2],
+             a_cta_layout,
+             cute.group_modes(sA, 0, 3),
+             cute.group_modes(tCgA, 0, 3),
          )
-         c_out[:, :, l_idx].copy_(res)
+         tBsB, tBgB = cpasync.tma_partition(
+             tma_atom_b, block_in_cluster_coord_vmnk[1],
+             b_cta_layout,
+             cute.group_modes(sB, 0, 3),
+             cute.group_modes(tCgB, 0, 3),
+         )
+         tAsSFA, tAgSFA = cpasync.tma_partition(
+             tma_atom_sfa, block_in_cluster_coord_vmnk[2],
+             a_cta_layout,
+             cute.group_modes(sSFA, 0, 3),
+             cute.group_modes(tCgSFA, 0, 3),
+         )
+         tBsSFB, tBgSFB = cpasync.tma_partition(
+             tma_atom_sfb, block_in_cluster_coord_vmnk[1],
+             b_cta_layout,
+             cute.group_modes(sSFB, 0, 3),
+             cute.group_modes(tCgSFB, 0, 3),
+         )
   
- def _is_case1(problem_sizes):
-     # case1: g=8, all (n=4096,k=7168,l=1), m varies
-     if len(problem_sizes) != 8:
-         return False
-     for (m, n, k, l) in problem_sizes:
-         if not (n == 4096 and k == 7168 and l == 1):
-             return False
-     return True
+         tAsSFA = cute.filter_zeros(tAsSFA); tAgSFA = cute.filter_zeros(tAgSFA)
+         tBsSFB = cute.filter_zeros(tBsSFB); tBgSFB = cute.filter_zeros(tBgSFB)
   
- def custom_kernel(data: input_t) -> output_t:
-     """
-     Optimized for case1 preparation cost:
-     - reorder_sf_blocked_fast: view/permute instead of meshgrid/scatter
-     - keep rest logic same
-     """
-     if len(data) == 4:
-         abc_tensors, sfasfb_tensors, _maybe_reordered, problem_sizes = data
-     else:
-         abc_tensors, sfasfb_tensors, problem_sizes = data
-         _maybe_reordered = None
+         tCrA = tiled_mma.make_fragment_A(sA)
+         tCrB = tiled_mma.make_fragment_B(sB)
+         acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
+         tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage))
   
-     result_tensors = []
+         pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
   
-     # (Optional) if future you finds _maybe_reordered contains ready-to-use blocked SF,
-     # you can plug it here. For now, we just ignore it safely.
+         # Producer warp
+         if warp_idx == self.tma_warp_id:
+             ab_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_ab_stage)
+             cur_tile_coord = (bidx, bidz, 0)
+             mma_tile_coord_mnl = (
+                 cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),
+                 cur_tile_coord[1],
+                 cur_tile_coord[2],
+             )
   
-     supported_k = {7168, 4096, 2048, 1536, 2304, 512, 256}
+             tAgA_slice = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
+             tBgB_slice = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
+             tAgSFA_slice = tAgSFA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
+             tBgSFB_slice = tBgSFB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
   
-     for (a, b, c), (sfa, sfb), (m, n, k, l) in zip(abc_tensors, sfasfb_tensors, problem_sizes):
-         assert l == 1, "This kernel assumes L==1 (matches provided benchmark shapes)."
+             for _ in cutlass.range(0, k_tile_cnt, 1, unroll=1):
+                 ab_pipeline.producer_acquire(ab_producer_state)
   
-         if sfa.device.type != "cuda":
-             sfa = sfa.cuda(non_blocking=True)
-         if sfb.device.type != "cuda":
-             sfb = sfb.cuda(non_blocking=True)
+                 cute.copy(
+                     tma_atom_a,
+                     tAgA_slice[(None, ab_producer_state.count)],
+                     tAsA[(None, ab_producer_state.index)],
+                     tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
+                     mcast_mask=a_full_mcast_mask,
+                     cache_policy=cutlass.Int64(cutlass.Int64(self.cache_policy).ir_value()),
+                 )
+                 cute.copy(
+                     tma_atom_b,
+                     tBgB_slice[(None, ab_producer_state.count)],
+                     tBsB[(None, ab_producer_state.index)],
+                     tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
+                     mcast_mask=b_full_mcast_mask,
+                     cache_policy=cutlass.Int64(cutlass.Int64(self.cache_policy).ir_value()),
+                 )
+                 cute.copy(
+                     tma_atom_sfa,
+                     tAgSFA_slice[(None, ab_producer_state.count)],
+                     tAsSFA[(None, ab_producer_state.index)],
+                     tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
+                     mcast_mask=sfa_full_mcast_mask,
+                     cache_policy=cutlass.Int64(cutlass.Int64(self.cache_policy).ir_value()),
+                 )
+                 cute.copy(
+                     tma_atom_sfb,
+                     tBgSFB_slice[(None, ab_producer_state.count)],
+                     tBsSFB[(None, ab_producer_state.index)],
+                     tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
+                     mcast_mask=sfb_full_mcast_mask,
+                     cache_policy=cutlass.Int64(cutlass.Int64(self.cache_policy).ir_value()),
+                 )
   
-         mn_pad = ceil_div(m, 128) * 128
+                 ab_producer_state.advance()
   
-         a_pad = pad_fp4_rows(a, m, mn_pad)
-         b_use = b.contiguous()
+         # MMA warp（你原来怎么写 gemm / s2t / epilogue，就继续用你原来的——
+         # 这里不展开重写，避免我替你“换掉正确性已过”的 mainloop）
+         #
+         # 你只需要确保：当 cluster_shape_mn=(2,2)/(4,1) 时
+         # b_full_mcast_mask / sfb_full_mcast_mask 真正生效（is_b_mcast = True）
+         #
+         # ↓↓↓ 如果你希望我把你当前通过正确性的 group-gemm mainloop 原封不动塞进来，
+         # 你把那段 mainloop+epilogue 贴出来我就帮你合并成一份完整最终版。
+         pass
   
-         sfa_m = sfa[:, :, 0].contiguous()
-         sfb_n = sfb[:, :, 0].contiguous()
   
-         # >>> the key change:
-         sfa_blk = reorder_sf_blocked_fast(sfa_m, m, k, mn_pad)
-         sfb_blk = reorder_sf_blocked_fast(sfb_n, n, k, n)  # N already multiple of 128 for your target shapes
+ # ============================================================
+ # compile cache + dispatcher
+ # ============================================================
+ _compiled_kernel_cache: Dict[Tuple[int,int,int,int], Optional[object]] = {}
   
-         if mn_pad == m:
-             c_out = c
-         else:
-             c_out = torch.empty((mn_pad, n, l), device="cuda", dtype=torch.float16)
+ def compile_kernel(m: int, n: int, k: int, l: int):
+     key = (m, n, k, l)
+     if key in _compiled_kernel_cache:
+         return _compiled_kernel_cache[key]
   
-         if k in supported_k:
-             gemm(a_pad, b_use, sfa_blk, sfb_blk, c_out)
-             if mn_pad != m:
-                 c.copy_(c_out[:m, :, :])
-         else:
-             torch_scaled_mm_fallback(a, b, sfa, sfb, c, m, n, k, l)
+     a_ptr = make_ptr(cutlass.Float4E2M1FN, 0, cute.AddressSpace.gmem, assumed_align=16)
+     b_ptr = make_ptr(cutlass.Float4E2M1FN, 0, cute.AddressSpace.gmem, assumed_align=16)
+     c_ptr = make_ptr(cutlass.Float16,       0, cute.AddressSpace.gmem, assumed_align=16)
+     sfa_ptr = make_ptr(cutlass.Float8E4M3FN, 0, cute.AddressSpace.gmem, assumed_align=32)
  ⋯ diff truncated