submission 420103

+ import cutlass
+ import cutlass.cute as cute
+ import cutlass.utils as utils
+ import cutlass.pipeline as pipeline
+ from cutlass.cute.nvgpu import cpasync, tcgen05
+ import cutlass.utils.blackwell_helpers as sm100_utils
+ import cutlass.utils.blockscaled_layout as blockscaled_utils
+ from cutlass.cute.runtime import make_ptr
+  
+ import functools
+ from typing import Tuple, List
+  
  import torch
  from task import input_t, output_t
   
- # Scaling factor vector size
- sf_vec_size = 16
+ # Kernel configuration parameters
+ # Size of tma descriptor in bytes
+ bytes_per_tensormap = 128
+ # Number of tensormaps: a, b, sfa, sfb
+ num_tensormaps = 4
+ # Tile sizes for M, N, K dimensions
+ mma_tiler_mnk = (128, 128, 256)  
+ # Shape of the K dimension for the MMA instruction
+ mma_inst_shape_k = 64
+ # FP4 data type for A and B
+ ab_dtype = cutlass.Float4E2M1FN  
+ # FP8 data type for scale factors
+ sf_dtype = cutlass.Float8E4M3FN  
+ # FP16 output type
+ c_dtype = cutlass.Float16  
+ # Scale factor block size (16 elements share one scale)
+ sf_vec_size = 16  
+ # Number of threads per CUDA thread block
+ threads_per_cta = 128  
+ # Stage numbers of shared memory and tmem
+ # OPTIMIZATION: Increased from 1 to 2 for double-buffering of A/B loads
+ # This allows overlapping TMA load of next K-tile with compute of current K-tile
+ num_ab_stage = 4  # Experiment 7: Quad-buffering
+ num_acc_stage = 1
+ # Total number of columns in tmem
+ num_tmem_alloc_cols = 512
   
   
  # Helper function for ceiling division
  ⋯ 1 unchanged lines
      return (a + b - 1) // b
   
   
- # Helper function to convert scale factor tensor to blocked format
- def to_blocked(input_matrix):
-     rows, cols = input_matrix.shape
+ # The CuTe reference implementation for NVFP4 block-scaled GEMM
+ @cute.kernel
+ def kernel(
+     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,
+     tensor_of_abc_ptrs: cute.Tensor,
+     tensor_of_sfasfb_ptrs: cute.Tensor,
+     tensormaps: cute.Tensor,
+     tensor_of_problem_sizes: cute.Tensor,
+     a_smem_layout_staged: cute.ComposedLayout,
+     b_smem_layout_staged: cute.ComposedLayout,
+     sfa_smem_layout_staged: cute.Layout,
+     sfb_smem_layout_staged: cute.Layout,
+     cta_mn_list: List[Tuple[int, int]],
+     num_tma_load_bytes: cutlass.Constexpr[int],
+ ):
+     """
+     GPU device kernel performing the Group GEMM computation.
+     
+     OPTIMIZATION CHANGES (from submission_cute_v1.py):
+     - num_ab_stage = 4  # Experiment 7: Quad-buffering: Double-buffer A/B/SFA/SFB to overlap loads with compute
+     - This allows next K-tile to load while current K-tile computes
+     - Expected improvement: 10-20% from reduced TMA stalls
+     """
+     warp_idx = cute.arch.warp_idx()
+     warp_idx = cute.arch.make_warp_uniform(warp_idx)
+     tidx, _, _ = cute.arch.thread_idx()
   
-     # Please ensure rows and cols are multiples of 128 and 4 respectively
-     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
+     #
+     # Delinearize bidz to coord_x, coord_y and group_idx for each CTA
+     #
+     bidx, bidy, bidz = cute.arch.block_idx()
+     group_idx = 0
+     find = False
+     coord_x = 0
+     coord_y = 0
+     cta_rest = bidz
+     for _, (cta_m, cta_n) in enumerate(cta_mn_list):
+         if cta_rest >= (cta_m * cta_n):
+             group_idx += 1
+             cta_rest -= cta_m * cta_n
+         else:
+             if not find:
+                 coord_y = cta_rest // cta_m
+                 coord_x = cta_rest % cta_m
+                 cta_rest -= cta_m * cta_n
+                 find = True
   
-     # Pad the input matrix if necessary
-     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,
+     #
+     # Construct C Tensor for each CTA
+     #
+     mC_mnl_iter = cute.make_ptr(
+         c_dtype, tensor_of_abc_ptrs[group_idx, 2], cute.AddressSpace.gmem
+     ).align(32)
+     m = tensor_of_problem_sizes[group_idx, 0]
+     n = tensor_of_problem_sizes[group_idx, 1]
+     k = tensor_of_problem_sizes[group_idx, 2]
+     l = tensor_of_problem_sizes[group_idx, 3]
+  
+     mC_mnl_layout = cute.make_layout(
+         (m, n, l),
+         stride=(cute.assume(n, 32), 1, cute.assume(m * n, 32),))
+     mC_mnl = cute.make_tensor(mC_mnl_iter, mC_mnl_layout)
+     # Local partition for global C Tensor
+     # (bM, bN, RestM, RestN, RestL)
+     gC_mnl = cute.local_tile(
+         mC_mnl, cute.slice_(mma_tiler_mnk, (None, None, 0)), (coord_x, coord_y, 0)
+     )
+  
+     #
+     # Define shared storage for kernel
+     #
+     size_tensormap_in_i64 = (
+         num_tensormaps * bytes_per_tensormap // 8
+     )
+     @cute.struct
+     class SharedStorage:
+         tensormap_buffer: cute.struct.MemRange[
+             cutlass.Int64, size_tensormap_in_i64
+         ]
+         ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, num_ab_stage * 2]
+         acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, num_acc_stage * 2]
+         tmem_holding_buf: cutlass.Int32
+     smem = utils.SmemAllocator()
+     storage = smem.allocate(SharedStorage)
+  
+     tensormap_smem_ptr = storage.tensormap_buffer.data_ptr()
+     tensormap_a_smem_ptr = tensormap_smem_ptr
+     tensormap_b_smem_ptr = (
+         tensormap_a_smem_ptr
+         + bytes_per_tensormap // 8
+     )
+     tensormap_sfa_smem_ptr = (
+         tensormap_b_smem_ptr
+         + bytes_per_tensormap // 8
+     )
+     tensormap_sfb_smem_ptr = (
+         tensormap_sfa_smem_ptr
+         + bytes_per_tensormap // 8
+     )
+     # Setup smem tensor for A, B, SFA, SFB
+     # (MMA, MMA_M, MMA_K, STAGE) - now with STAGE dimension due to num_ab_stage=2
+     sA = smem.allocate_tensor(
+         element_type=ab_dtype,
+         layout=a_smem_layout_staged.outer,
+         byte_alignment=128,
+         swizzle=a_smem_layout_staged.inner,
+     )
+     # (MMA, MMA_N, MMA_K, STAGE)
+     sB = smem.allocate_tensor(
+         element_type=ab_dtype,
+         layout=b_smem_layout_staged.outer,
+         byte_alignment=128,
+         swizzle=b_smem_layout_staged.inner,
+     )
+     # (MMA, MMA_M, MMA_K, STAGE)
+     sSFA = smem.allocate_tensor(
+         element_type=sf_dtype,
+         layout=sfa_smem_layout_staged,
+         byte_alignment=128,
+     )
+     # (MMA, MMA_N, MMA_K, STAGE)
+     sSFB = smem.allocate_tensor(
+         element_type=sf_dtype,
+         layout=sfb_smem_layout_staged,
+         byte_alignment=128,
+     )
+  
+     # Initialize mainloop ab_pipeline, acc_pipeline and their states
+     ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+     ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1)
+     ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create(
+         barrier_storage=storage.ab_mbar_ptr.data_ptr(),
+         num_stages=num_ab_stage,
+         producer_group=ab_pipeline_producer_group,
+         consumer_group=ab_pipeline_consumer_group,
+         tx_count=num_tma_load_bytes,
+     ).make_participants()
+     acc_producer, acc_consumer = pipeline.PipelineUmmaAsync.create(
+         barrier_storage=storage.acc_mbar_ptr.data_ptr(),
+         num_stages=num_acc_stage,
+         producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
+         consumer_group=pipeline.CooperativeGroup(
+             pipeline.Agent.Thread,
+             threads_per_cta,
+         ),
+     ).make_participants()
+  
+     #
+     # Local_tile partition global tensors
+     #
+     # (bM, bK, RestM, RestK, RestL)
+     gA_mkl = cute.local_tile(
+         mA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None)
+     )
+     # (bN, bK, RestN, RestK, RestL)
+     gB_nkl = cute.local_tile(
+         mB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None)
+     )
+     # (bM, bK, RestM, RestK, RestL)
+     gSFA_mkl = cute.local_tile(
+         mSFA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None)
+     )
+     # (bN, bK, RestN, RestK, RestL)
+     gSFB_nkl = cute.local_tile(
+         mSFB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None)
+     )
+     #
+     # Partition global tensor for TiledMMA_A/B/C
+     #
+     thr_mma = tiled_mma.get_slice(tidx)
+     # (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
+     tCgA = thr_mma.partition_A(gA_mkl)
+     # (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
+     tCgB = thr_mma.partition_B(gB_nkl)
+     # (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
+     tCgSFA = thr_mma.partition_A(gSFA_mkl)
+     # (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
+     tCgSFB = thr_mma.partition_B(gSFB_nkl)
+     # (MMA, MMA_M, MMA_N, RestM, RestN, RestL)
+     tCgC = thr_mma.partition_C(gC_mnl)
+  
+     # Update tma descriptor with the correct shapes and strides
+     tensormap_manager = utils.TensorMapManager(
+         utils.TensorMapUpdateMode.SMEM,
+         128,
+     )
+     tensormap_a_gmem_ptr = tensormap_manager.get_tensormap_ptr(
+         tensormaps[(bidz, 0, None)].iterator
+     )
+     tensormap_b_gmem_ptr = tensormap_manager.get_tensormap_ptr(
+         tensormaps[(bidz, 1, None)].iterator
+     )
+     tensormap_sfa_gmem_ptr = tensormap_manager.get_tensormap_ptr(
+         tensormaps[(bidz, 2, None)].iterator
+     )
+     tensormap_sfb_gmem_ptr = tensormap_manager.get_tensormap_ptr(
+         tensormaps[(bidz, 3, None)].iterator
+     )
+  
+     mA_mkl_iter = cute.make_ptr(
+         ab_dtype, tensor_of_abc_ptrs[group_idx, 0], cute.AddressSpace.gmem
+     ).align(32)
+     mB_nkl_iter = cute.make_ptr(
+         ab_dtype, tensor_of_abc_ptrs[group_idx, 1], cute.AddressSpace.gmem
+     ).align(32)
+     sfa_mkl_iter = cute.make_ptr(
+         sf_dtype, tensor_of_sfasfb_ptrs[group_idx, 0], cute.AddressSpace.gmem
+     ).align(32)
+     sfb_nkl_iter = cute.make_ptr(
+         sf_dtype, tensor_of_sfasfb_ptrs[group_idx, 1], cute.AddressSpace.gmem
+     ).align(32)
+     mA_mkl_layout = cute.make_layout(
+         (m, k, l), stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32),))
+     mB_nkl_layout = cute.make_layout(
+         (n, k, l), stride=(cute.assume(k, 32), 1, cute.assume(n * k, 32),))
+  
+     # SFA, SFB follows specialized layout defined in the following link:
+     # https://docs.nvidia.com/cuda/cublas/index.html?highlight=fp4#d-block-scaling-factors-layout
+     atom_shape = ((32, 4), (sf_vec_size, 4))
+     atom_stride = ((16, 4), (0, 1))
+     sfa_layout = cute.tile_to_shape(
+         cute.make_layout(atom_shape, stride=atom_stride),
+         mA_mkl_layout.shape,
+         (2, 1, 3),
+     )
+     sfb_layout = cute.tile_to_shape(
+         cute.make_layout(atom_shape, stride=atom_stride),
+         mB_nkl_layout.shape,
+         (2, 1, 3),
+     )
+     real_tensor_a = cute.make_tensor(mA_mkl_iter, mA_mkl_layout)
+     real_tensor_b = cute.make_tensor(mB_nkl_iter, mB_nkl_layout)
+     real_tensor_sfa = cute.make_tensor(sfa_mkl_iter, sfa_layout)
+     real_tensor_sfb = cute.make_tensor(sfb_nkl_iter, sfb_layout)
+  
+     # Let warp 0 initialize tensormap
+     if warp_idx == 0:
+         tensormap_manager.init_tensormap_from_atom(
+             tma_atom_a, tensormap_a_smem_ptr, 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)
+         tensormap_manager.init_tensormap_from_atom(
+             tma_atom_b, tensormap_b_smem_ptr, 0
+         )
+         tensormap_manager.init_tensormap_from_atom(
+             tma_atom_sfa, tensormap_sfa_smem_ptr, 0
+         )
+         tensormap_manager.init_tensormap_from_atom(
+             tma_atom_sfb, tensormap_sfb_smem_ptr, 0
+         )
+         tensormap_manager.update_tensormap(
+             (
+                 real_tensor_a,
+                 real_tensor_b,
+                 real_tensor_sfa,
+                 real_tensor_sfb,
+             ),
+             (tma_atom_a, tma_atom_b, tma_atom_sfa, tma_atom_sfb),
+             (
+                 tensormap_a_gmem_ptr,
+                 tensormap_b_gmem_ptr,
+                 tensormap_sfa_gmem_ptr,
+                 tensormap_sfb_gmem_ptr,
+             ),
+             0,  # tma warp id
+             (
+                 tensormap_a_smem_ptr,
+                 tensormap_b_smem_ptr,
+                 tensormap_sfa_smem_ptr,
+                 tensormap_sfb_smem_ptr,
+             ),
+         )
   
-     return rearranged.flatten()
+         tensormap_manager.fence_tensormap_update(tensormap_a_gmem_ptr)
+         tensormap_manager.fence_tensormap_update(tensormap_b_gmem_ptr)
+         tensormap_manager.fence_tensormap_update(tensormap_sfa_gmem_ptr)
+         tensormap_manager.fence_tensormap_update(tensormap_sfb_gmem_ptr)
   
+     cute.arch.barrier()
   
- def custom_kernel(data: input_t) -> output_t:
+     #
+     # Partition global/shared tensor for TMA load A/B/SFA/SFB
+     #
+     # TMA Partition_S/D for A
+     # ((atom_v, rest_v), STAGE)
+     # ((atom_v, rest_v), RestM, RestK, RestL)
+     tAsA, tAgA = cpasync.tma_partition(
+         tma_atom_a,
+         0,
+         cute.make_layout(1),
+         cute.group_modes(sA, 0, 3),
+         cute.group_modes(tCgA, 0, 3),
+     )
+     # TMA Partition_S/D for B
+     # ((atom_v, rest_v), STAGE)
+     # ((atom_v, rest_v), RestN, RestK, RestL)
+     tBsB, tBgB = cpasync.tma_partition(
+         tma_atom_b,
+         0,
+         cute.make_layout(1),
+         cute.group_modes(sB, 0, 3),
+         cute.group_modes(tCgB, 0, 3),
+     )
+     #  TMA Partition_S/D for SFA
+     # ((atom_v, rest_v), STAGE)
+     # ((atom_v, rest_v), RestM, RestK, RestL)
+     tAsSFA, tAgSFA = cpasync.tma_partition(
+         tma_atom_sfa,
+         0,
+         cute.make_layout(1),
+         cute.group_modes(sSFA, 0, 3),
+         cute.group_modes(tCgSFA, 0, 3),
+     )
+     tAsSFA = cute.filter_zeros(tAsSFA)
+     tAgSFA = cute.filter_zeros(tAgSFA)
+     # TMA Partition_S/D for SFB
+     # ((atom_v, rest_v), STAGE)
+     # ((atom_v, rest_v), RestN, RestK, RestL)
+     tBsSFB, tBgSFB = cpasync.tma_partition(
+         tma_atom_sfb,
+         0,
+         cute.make_layout(1),
+         cute.group_modes(sSFB, 0, 3),
+         cute.group_modes(tCgSFB, 0, 3),
+     )
+     tBsSFB = cute.filter_zeros(tBsSFB)
+     tBgSFB = cute.filter_zeros(tBgSFB)
+  
+     #
+     # Partition shared/tensor memory tensor for TiledMMA_A/B/C
+     #
+     # (MMA, MMA_M, MMA_K, STAGE)
+     tCrA = tiled_mma.make_fragment_A(sA)
+     # (MMA, MMA_N, MMA_K, STAGE)
+     tCrB = tiled_mma.make_fragment_B(sB)
+     # (MMA, MMA_M, MMA_N)
+     acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2])
+     # (MMA, MMA_M, MMA_N)
+     tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape)
+     #
+     # Alloc tensor memory buffer
+     #
+     tmem_alloc_barrier = pipeline.NamedBarrier(
+         barrier_id=1,
+         num_threads=threads_per_cta,
+     )
+     tmem = utils.TmemAllocator(
+         storage.tmem_holding_buf,
+         barrier_for_retrieve=tmem_alloc_barrier,
+     )
+     tmem.allocate(num_tmem_alloc_cols)
+     tmem.wait_for_alloc()
+     acc_tmem_ptr = tmem.retrieve_ptr(cutlass.Float32)
+     tCtAcc = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
+  
+     #
+     # Make SFA/SFB tmem tensor
+     #
+     # Get SFA tmem ptr
+     sfa_tmem_ptr = cute.recast_ptr(
+         acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc),
+         dtype=sf_dtype,
+     )
+     # (MMA, MMA_M, MMA_K)
+     tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
+         tiled_mma,
+         mma_tiler_mnk,
+         sf_vec_size,
+         cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)),
+     )
+     tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout)
+     # Get SFB tmem ptr
+     sfb_tmem_ptr = cute.recast_ptr(
+         acc_tmem_ptr
+         + tcgen05.find_tmem_tensor_col_offset(tCtAcc)
+         + tcgen05.find_tmem_tensor_col_offset(tCtSFA),
+         dtype=sf_dtype,
+     )
+     # (MMA, MMA_N, MMA_K)
+     tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
+         tiled_mma,
+         mma_tiler_mnk,
+         sf_vec_size,
+         cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)),
+     )
+     tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout)
+  
+     #
+     # Partition for S2T copy of SFA/SFB
+     #
+     # Make S2T CopyAtom
+     copy_atom_s2t = cute.make_copy_atom(
+         tcgen05.Cp4x32x128bOp(tcgen05.CtaGroup.ONE),
+         sf_dtype,
+     )
+     # (MMA, MMA_MN, MMA_K, STAGE)
+     tCsSFA_compact = cute.filter_zeros(sSFA)
+     tCtSFA_compact = cute.filter_zeros(tCtSFA)
+     tiled_copy_s2t_sfa = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFA_compact)
+     thr_copy_s2t_sfa = tiled_copy_s2t_sfa.get_slice(0)
+     # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
+     tCsSFA_compact_s2t_ = thr_copy_s2t_sfa.partition_S(tCsSFA_compact)
+     # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
+     tCsSFA_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
+         tiled_copy_s2t_sfa, tCsSFA_compact_s2t_
+     )
+     # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
+     tCtSFA_compact_s2t = thr_copy_s2t_sfa.partition_D(tCtSFA_compact)
+  
+     # (MMA, MMA_MN, MMA_K, STAGE)
+     tCsSFB_compact = cute.filter_zeros(sSFB)
+     # (MMA, MMA_MN, MMA_K)
+     tCtSFB_compact = cute.filter_zeros(tCtSFB)
+     tiled_copy_s2t_sfb = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFB_compact)
+     thr_copy_s2t_sfb = tiled_copy_s2t_sfb.get_slice(0)
+     # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
+     tCsSFB_compact_s2t_ = thr_copy_s2t_sfb.partition_S(tCsSFB_compact)
+     # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
+     tCsSFB_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
+         tiled_copy_s2t_sfb, tCsSFB_compact_s2t_
+     )
+     # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
+     tCtSFB_compact_s2t = thr_copy_s2t_sfb.partition_D(tCtSFB_compact)
+  
+     # Number of K loops
+     k_tile_cnt = cute.ceil_div(real_tensor_a.shape[1], mma_tiler_mnk[2])
+  
+     #
+     # Slice to per mma tile index
+     #
+     mma_tile_coord_mnl = (coord_x, coord_y, 0)
+     # ((atom_v, rest_v), RestK)
+     tAgA = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
+     # ((atom_v, rest_v), RestK)
+     tBgB = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
+     # ((atom_v, rest_v), RestK)
+     tAgSFA = tAgSFA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
+     # ((atom_v, rest_v), RestK)
+     tBgSFB = tBgSFB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
+  
+     #
+     # Main loop - now with double-buffered pipelining
+     # OPTIMIZATION: With num_ab_stage=2, we can overlap load of K-tile (i+1)
+     # with compute of K-tile (i), reducing stalls from TMA latency
+     #
+     if warp_idx == 0:
+         # Wait for accumulator buffer empty
+         acc_empty = acc_producer.acquire_and_advance()
+         # Set ACCUMULATE field to False for the first k_tile iteration
+         tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
+         
+         # Preload first K-tile
+         ab_empty = ab_producer.acquire_and_advance()
+         cute.copy(
+             tma_atom_a,
+             tAgA[(None, 0)],
+             tAsA[(None, ab_empty.index)],
+             tma_bar_ptr=ab_empty.barrier,
+             tma_desc_ptr=tensormap_manager.get_tensormap_ptr(
+                 tensormap_a_gmem_ptr,
+                 cute.AddressSpace.generic,
+             ),
+         )
+         cute.copy(
+             tma_atom_b,
+             tBgB[(None, 0)],
+             tBsB[(None, ab_empty.index)],
+             tma_bar_ptr=ab_empty.barrier,
+             tma_desc_ptr=tensormap_manager.get_tensormap_ptr(
+                 tensormap_b_gmem_ptr,
+                 cute.AddressSpace.generic,
+             ),
+         )
+         cute.copy(
+             tma_atom_sfa,
+             tAgSFA[(None, 0)],
+             tAsSFA[(None, ab_empty.index)],
+             tma_bar_ptr=ab_empty.barrier,
+             tma_desc_ptr=tensormap_manager.get_tensormap_ptr(
+                 tensormap_sfa_gmem_ptr,
+                 cute.AddressSpace.generic,
+             ),
+         )
+         cute.copy(
+             tma_atom_sfb,
+             tBgSFB[(None, 0)],
+             tBsSFB[(None, ab_empty.index)],
+             tma_bar_ptr=ab_empty.barrier,
+             tma_desc_ptr=tensormap_manager.get_tensormap_ptr(
+                 tensormap_sfb_gmem_ptr,
+                 cute.AddressSpace.generic,
+             ),
+         )
+         
+         # Execute k_tile loop with overlapped loads
+         for k_tile in range(k_tile_cnt):
+             # Wait for current AB buffer full
+             ab_full = ab_consumer.wait_and_advance()
+  
+             # If not the last tile, preload next K-tile while computing current
+             if k_tile < k_tile_cnt - 1:
+                 ab_empty = ab_producer.acquire_and_advance()
+                 cute.copy(
+                     tma_atom_a,
+                     tAgA[(None, k_tile + 1)],
+                     tAsA[(None, ab_empty.index)],
+                     tma_bar_ptr=ab_empty.barrier,
+                     tma_desc_ptr=tensormap_manager.get_tensormap_ptr(
+                         tensormap_a_gmem_ptr,
+                         cute.AddressSpace.generic,
+                     ),
+                 )
+                 cute.copy(
+                     tma_atom_b,
+                     tBgB[(None, k_tile + 1)],
+                     tBsB[(None, ab_empty.index)],
+                     tma_bar_ptr=ab_empty.barrier,
+                     tma_desc_ptr=tensormap_manager.get_tensormap_ptr(
+                         tensormap_b_gmem_ptr,
+                         cute.AddressSpace.generic,
+                     ),
+                 )
+                 cute.copy(
+                     tma_atom_sfa,
+                     tAgSFA[(None, k_tile + 1)],
+                     tAsSFA[(None, ab_empty.index)],
+                     tma_bar_ptr=ab_empty.barrier,
+                     tma_desc_ptr=tensormap_manager.get_tensormap_ptr(
+                         tensormap_sfa_gmem_ptr,
+                         cute.AddressSpace.generic,
+                     ),
+                 )
+                 cute.copy(
+                     tma_atom_sfb,
+                     tBgSFB[(None, k_tile + 1)],
+                     tBsSFB[(None, ab_empty.index)],
+                     tma_bar_ptr=ab_empty.barrier,
+                     tma_desc_ptr=tensormap_manager.get_tensormap_ptr(
+                         tensormap_sfb_gmem_ptr,
+                         cute.AddressSpace.generic,
+                     ),
+                 )
+  
+             #  Copy SFA/SFB from shared memory to TMEM
+             s2t_stage_coord = (None, None, None, None, ab_full.index)
+             tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord]
+             tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord]
+             cute.copy(
+                 tiled_copy_s2t_sfa,
+                 tCsSFA_compact_s2t_staged,
+                 tCtSFA_compact_s2t,
+             )
+             cute.copy(
+                 tiled_copy_s2t_sfb,
+                 tCsSFB_compact_s2t_staged,
+                 tCtSFB_compact_s2t,
+             )
+  
+             # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB
+             num_kblocks = cute.size(tCrA, mode=[2])
+             for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):
+                 kblock_coord = (
+                     None,
+                     None,
+                     kblock_idx,
+                     ab_full.index,
+                 )
+  
+                 # Set SFA/SFB tensor to tiled_mma
+                 sf_kblock_coord = (None, None, kblock_idx)
+                 tiled_mma.set(
+                     tcgen05.Field.SFA,
+                     tCtSFA[sf_kblock_coord].iterator,
+                 )
+                 tiled_mma.set(
+                     tcgen05.Field.SFB,
+                     tCtSFB[sf_kblock_coord].iterator,
+                 )
+  
+                 cute.gemm(
+                     tiled_mma,
+                     tCtAcc,
+                     tCrA[kblock_coord],
+                     tCrB[kblock_coord],
+                     tCtAcc,
+                 )
+                 # Enable accumulate on tCtAcc after first kblock
+                 tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+  
+             # Async arrive AB buffer empty
+             ab_full.release()
+         acc_empty.commit()
+  
+     #
+     # Epilogue
+     # Partition for epilogue
+     #
+     op = tcgen05.Ld32x32bOp(tcgen05.Repetition.x128, tcgen05.Pack.NONE)
+     copy_atom_t2r = cute.make_copy_atom(op, cutlass.Float32)
+     tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tCtAcc[None,0,0])
+     thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
+     # (TmemCpy, NumTmemCpy)
+     tDtAcc = thr_copy_t2r.partition_S(tCtAcc[None,0,0])
+     # (TmemCpy, NumTmemCpy)
+     tDgC = thr_copy_t2r.partition_D(tCgC[None,0,0])
+  
+     # (TmemCpy, NumTmemCpy)
+     tDrAcc = cute.make_rmem_tensor(tDgC.shape, cutlass.Float32)
+     # (TmemCpy, NumTmemCpy)
+     tDrC = cute.make_rmem_tensor(tDgC.shape, c_dtype)
+  
+     # Release TMEM allocation lock
+     tmem.relinquish_alloc_permit()
+     # Wait for accumulator buffer full
+     acc_full = acc_consumer.wait_and_advance()
+  
+     # Copy accumulator to register
+     cute.copy(tiled_copy_t2r, tDtAcc, tDrAcc)
+     acc_vec = tDrAcc.load()
+     tDrC.store(acc_vec.to(c_dtype))
+  
+     # STG Atom, just to ensure functionality
+     # For performance optimization, better to use Tma store operation to
+     # reduce address calculation and predicate calulation instructions
+     simt_atom = cute.make_copy_atom(
+         cute.nvgpu.CopyUniversalOp(), c_dtype, num_bits_per_copy=16
+     )
+     thread_layout = cute.make_layout(
+         (1, threads_per_cta), stride=(threads_per_cta, 1))
+     value_layout = cute.make_layout((1, 1))
+     tiled_copy_r2g = cute.make_tiled_copy_tv(
+         simt_atom, thread_layout, value_layout
+     )
+     thr_copy_r2g = tiled_copy_r2g.get_slice(tidx)
+     cC = cute.make_identity_tensor(gC_mnl.shape)
+     # ((atom_v, rest_v), NumGmemCpy)
+     tDcC = thr_copy_r2g.partition_D(cC)
+  
+     # ((atom_v, rest_v), NumGmemCpy)
+     tDpC = cute.make_rmem_tensor(tDrC.shape, cutlass.Boolean)
+     residue_m = mC_mnl.shape[0] - cutlass.Int32(coord_x) * mma_tiler_mnk[0]
+     residue_n = mC_mnl.shape[1] - cutlass.Int32(coord_y) * mma_tiler_mnk[1]
+     for i in range(cute.size(tDrC.shape)):
+         # Swap residue_m and residue_n to match the order of tDcC
+         tDpC[i] = cute.elem_less(tDcC[i], (residue_n, residue_m))
+     cute.copy(simt_atom, cute.flatten(tDrC), cute.flatten(tDgC), pred=cute.flatten(tDpC))
+  
+     acc_full.release()
+     # Deallocate TMEM
+     cute.arch.barrier()
+     tmem.free(acc_tmem_ptr)
+     pass
+  
+  
+ # Host-side JIT function to prepare tensors and launch GPU kernel.
+ @cute.jit
+ def my_kernel(
+     ptr_of_tensor_of_problem_sizes: cute.Pointer,
+     ptr_of_tensor_of_abc_ptrs: cute.Pointer,
+     ptr_of_tensor_of_sfasfb_ptrs: cute.Pointer,
+     ptr_of_tensor_of_tensormap: cute.Pointer,
+     total_num_clusters: cutlass.Int32,
+     problem_sizes: List[
+         Tuple[int, int, int, int]
+     ],  # Problem sizes for each group
+     num_groups: cutlass.Int32,
+ ):
+  
+     tensor_of_abc_ptrs = cute.make_tensor(
+         ptr_of_tensor_of_abc_ptrs, cute.make_layout((num_groups, 3), stride=(3, 1))
+     )
+     tensor_of_sfasfb_ptrs = cute.make_tensor(
+         ptr_of_tensor_of_sfasfb_ptrs, cute.make_layout((num_groups, 2), stride=(2, 1))
+     )
+     tensor_of_problem_sizes = cute.make_tensor(
+         ptr_of_tensor_of_problem_sizes, cute.make_layout((num_groups, 4), stride=(4, 1))
+     )
+     tensor_of_tensormap = cute.make_tensor(
+         ptr_of_tensor_of_tensormap, cute.make_layout((total_num_clusters, 4, 16), stride=(64, 16, 1))
+     )
+  
+     # Use fake shape for initial Tma descriptor and atom setup
+     # The real Tma desc and atom will be updated during kernel execution.
+     min_a_shape = (cutlass.Int32(64), cutlass.Int32(64), cutlass.Int32(64), cutlass.Int32(1))
+     min_b_shape = (cutlass.Int32(64), cutlass.Int32(64), cutlass.Int32(64), cutlass.Int32(1))
+     initial_a = cute.make_tensor(
+         cute.make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16,),
+         cute.make_layout(
+             (min_a_shape[0], cute.assume(min_a_shape[2], 32), min_a_shape[3]),
+             stride=(
+                 cute.assume(min_a_shape[2], 32),
+                 1,
+                 cute.assume(min_a_shape[0] * min_a_shape[2], 32),
+             ),
+         ),
+     )
+     initial_b = cute.make_tensor(
+         cute.make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16,),
+         cute.make_layout(
+             (min_b_shape[1], cute.assume(min_b_shape[2], 32), min_b_shape[3]),
+             stride=(
+                 cute.assume(min_b_shape[2], 32),
+                 1,
+                 cute.assume(min_b_shape[1] * min_b_shape[2], 32),
+             ),
+         ),
+     )
+  
+     # Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout
+     # ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL)
+     sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(
+         initial_a.shape, sf_vec_size
+     )
+     # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL)
+     sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(
+         initial_b.shape, sf_vec_size
+     )
+     # Create initial SFA and SFB tensors with fake shape and null pointer.
+     initial_sfa = cute.make_tensor(
+         cute.make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=16,), sfa_layout)
+     initial_sfb = cute.make_tensor(
+         cute.make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=16,), sfb_layout)
+  
+     # Select MMA operation
+     mma_op = tcgen05.MmaMXF4NVF4Op(
+         sf_dtype,
+         (mma_tiler_mnk[0], mma_tiler_mnk[1], mma_inst_shape_k),
+         tcgen05.CtaGroup.ONE,
+         tcgen05.OperandSource.SMEM,
+     )
+     tiled_mma = cute.make_tiled_mma(mma_op)
+  
+     cluster_layout_vmnk = cute.tiled_divide(
+         cute.make_layout((1, 1, 1)),
+         (tiled_mma.thr_id.shape,),
+     )
+  
+     # Compute A/B/SFA/SFB/C shared memory layout
+     # OPTIMIZATION: num_ab_stage=2 increases SMEM for staging but enables pipelining
+     a_smem_layout_staged = sm100_utils.make_smem_layout_a(
+         tiled_mma,
+         mma_tiler_mnk,
+         ab_dtype,
+         num_ab_stage,
+     )
+     b_smem_layout_staged = sm100_utils.make_smem_layout_b(
+         tiled_mma,
+         mma_tiler_mnk,
+         ab_dtype,
+         num_ab_stage,
+     )
+     sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
+         tiled_mma,
+         mma_tiler_mnk,
+         sf_vec_size,
+         num_ab_stage,
+     )
+     sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(
+         tiled_mma,
+         mma_tiler_mnk,
+         sf_vec_size,
+         num_ab_stage,
+     )
+     atom_thr_size = cute.size(tiled_mma.thr_id.shape)
+  
+     # Setup TMA for A
+     a_smem_layout = cute.slice_(a_smem_layout_staged, (None, None, None, 0))
+     tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
+         cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
+         initial_a,
+         a_smem_layout,
+         mma_tiler_mnk,
+         tiled_mma,
+         cluster_layout_vmnk.shape,
+     )
+     # Setup TMA for B
+     b_smem_layout = cute.slice_(b_smem_layout_staged, (None, None, None, 0))
+     tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
+         cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
+         initial_b,
+         b_smem_layout,
+         mma_tiler_mnk,
+         tiled_mma,
+         cluster_layout_vmnk.shape,
+     )
+     # Setup TMA for SFA
+     sfa_smem_layout = cute.slice_(
+         sfa_smem_layout_staged, (None, None, None, 0)
+     )
+     tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A(
+         cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
+         initial_sfa,
+         sfa_smem_layout,
+         mma_tiler_mnk,
+         tiled_mma,
+         cluster_layout_vmnk.shape,
+         internal_type=cutlass.Int16,
+     )
+     # Setup TMA for SFB
+     sfb_smem_layout = cute.slice_(
+         sfb_smem_layout_staged, (None, None, None, 0)
+     )
+     tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B(
+         cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
+         initial_sfb,
+         sfb_smem_layout,
+         mma_tiler_mnk,
+         tiled_mma,
+         cluster_layout_vmnk.shape,
+         internal_type=cutlass.Int16,
+     )
+  
+     # Compute TMA load bytes
+     a_copy_size = cute.size_in_bytes(ab_dtype, a_smem_layout)
+     b_copy_size = cute.size_in_bytes(ab_dtype, b_smem_layout)
+     sfa_copy_size = cute.size_in_bytes(sf_dtype, sfa_smem_layout)
+     sfb_copy_size = cute.size_in_bytes(sf_dtype, sfb_smem_layout)
+     num_tma_load_bytes = (
+         a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size
+     ) * atom_thr_size
+  
+     # Store CTA shape information for each Group in a List
+     cta_mn_list = []
+     for group_idx, (m, n, k, l) in enumerate(problem_sizes):
+         x, y = cute.ceil_div(problem_sizes[group_idx][:2], mma_tiler_mnk[0:2])
+         cta_mn_list.append((x, y))
+  
+     # Compute grid size
+     grid = (1, 1, total_num_clusters)
+  
+     # Launch the kernel
+     kernel(
+         # MMA (Matrix Multiply-Accumulate) configuration
+         tiled_mma,                  # Tiled MMA object defining NVFP4 GEMM compute pattern
+         
+         # TMA (Tensor Memory Accelerator) atoms and tensors for input matrix A
+         tma_atom_a,                 # TMA copy atom defining how to load A from global memory
+         tma_tensor_a,               # Tensor descriptor for A (created from smallest A tensor)
+         
+         # TMA atoms and tensors for input matrix B
+         tma_atom_b,                 # TMA copy atom defining how to load B from global memory
+         tma_tensor_b,               # Tensor descriptor for B (created from smallest B tensor)
+         
+         # TMA atoms and tensors for scale factor A
+         tma_atom_sfa,               # TMA copy atom for loading scale factors for A
+         tma_tensor_sfa,             # Tensor descriptor for SFA (block scale factors for A)
+         
+         # TMA atoms and tensors for scale factor B
+         tma_atom_sfb,               # TMA copy atom for loading scale factors for B
+         tma_tensor_sfb,             # Tensor descriptor for SFB (block scale factors for B)
+         
+         # Runtime tensor metadata for dynamic group access
+         tensor_of_abc_ptrs,         # Device tensor containing pointers to A, B, C for all groups
+         tensor_of_sfasfb_ptrs,      # Device tensor containing pointers to SFA, SFB for all groups
+         tensor_of_tensormap,        # Pre-allocated buffer for tensormap descriptors per CTA
+         tensor_of_problem_sizes,    # Device tensor containing (m, n, k, l) for each group
+         
+         # Shared memory layouts with staging for pipelined execution
+         a_smem_layout_staged,       # Staged shared memory layout for A (includes stage dimension)
+         b_smem_layout_staged,       # Staged shared memory layout for B (includes stage dimension)
+         sfa_smem_layout_staged,     # Staged shared memory layout for SFA (includes stage dimension)
+         sfb_smem_layout_staged,     # Staged shared memory layout for SFB (includes stage dimension)
+         
+         # CTA grid configuration per group
+         cta_mn_list,                # List of (M_tiles, N_tiles) for each group
+         
+         # Pipeline synchronization parameter
+         num_tma_load_bytes,         # Total bytes to load per TMA transaction (for barrier setup)
+     ).launch(
+         grid=grid,
+         block=[threads_per_cta, 1, 1],
+         cluster=(1, 1, 1),
+     )
+     return
+  
+  
+ # Global cache for compiled kernels (keyed by group size)
+ _compiled_kernel_cache = {}
+ # This function is used to compile the kernel once and cache it and then allow users to 
+ # run the kernel multiple times to get more accurate timing results.
+ def compile_kernel(problem_sizes):
      """
-     PyTorch implementation of NVFP4 block-scaled group GEMM.
-     Using torch._scaled_mm for each group as baseline.
+     Compile the kernel once and cache it using problem_sizes as the key.
+     This should be called before any timing measurements.
   
+     Returns:
+         The compiled kernel function
+     """
+     global _compiled_kernel_cache
+     
+     # Convert problem_sizes list to a hashable tuple for use as dictionary key
+     cache_key = f"{len(problem_sizes)}"
+  
+     # Check if we already have a compiled kernel for these problem sizes
+     if cache_key in _compiled_kernel_cache:
+         return _compiled_kernel_cache[cache_key]
+  
+     cute_ptr_of_tensor_of_problem_sizes = make_ptr(
+         cutlass.Int32, 0, cute.AddressSpace.gmem, assumed_align=16,
+     )
+     cute_ptr_of_tensor_of_abc_ptrs = make_ptr(
+         cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16,
+     )
+     cute_ptr_of_tensor_of_sfasfb_ptrs = make_ptr(
+         cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16,
+     )
+     # Fake cluster numbers for compile only.
+     total_num_clusters = cutlass.Int32(1)
+     num_groups = cutlass.Int32(len(problem_sizes))
+     # Each cluster needs its own set of tensormaps (one for A, B, SFA, SFB)
+     # Shape: (total_num_clusters, num_tensormaps=4, bytes_per_tensormap/8=16)
+     cute_ptr_of_tensor_of_tensormap = make_ptr(
+         cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16,
+     )
+     compiled_func = cute.compile(
+         my_kernel,
+         cute_ptr_of_tensor_of_problem_sizes,
+         cute_ptr_of_tensor_of_abc_ptrs,
+         cute_ptr_of_tensor_of_sfasfb_ptrs,
+         cute_ptr_of_tensor_of_tensormap,
+         total_num_clusters,
+         problem_sizes,
+         num_groups
+     )
+     # Store compiled kernel in cache with problem_sizes as key
+     _compiled_kernel_cache[cache_key] = compiled_func
+     return compiled_func
+  
+  
+ def custom_kernel(data: input_t) -> output_t:
+     """
+     Execute the block-scaled group GEMM kernel.
+     
+     This is the main entry point called by the evaluation framework.
+     It converts PyTorch tensors to CuTe tensors, launches the kernel,
+     and returns the result.
+     
      Args:
-         data: list of tuples (abc_tensors, sfasfb_tensors, sfasfb_reordered_tensors, problem_sizes)
+         data: Tuple of (abc_tensors, sfasfb_tensors, problem_sizes) where:
+             abc_tensors: list of tuples (a, b, c) where 
+                 a is torch.Tensor[float4e2m1fn_x2] of shape [m, k // 2, l]
+                 b is torch.Tensor[float4e2m1fn_x2] of shape [n, k // 2, l]
+                 c is torch.Tensor[float16] of shape [m, n, l]
+             sfasfb_tensors: list of tuples (sfa, sfb) where 
+                 sfa is torch.Tensor[float8_e4m3fnuz] of shape [m, k // 16, l]
+                 sfb is torch.Tensor[float8_e4m3fnuz] of shape [n, k // 16, l]
+             problem_sizes: list of tuples (m, n, k, l)
+             each group has its own a, b, c, sfa, sfb with different m, n, k, l problem sizes
+             l should always be 1 for each group.
+             list size is the number of groups.
+     
      Returns:
-         list of output tensors [c_0, c_1, ..., c_G-1]
+         list of c tensors where c is torch.Tensor[float16] of shape [m, n, l] for each group
      """
-     abc_tensors, sfasfb_tensors, sfasfb_reordered_tensors, problem_sizes = data
-     result_tensors = []
+     abc_tensors, _, sfasfb_reordered_tensors, problem_sizes = data
   
-     for i, ((a, b, c), (sfa, sfb), (m, n, k, l)) in enumerate(
-         zip(abc_tensors, sfasfb_tensors, problem_sizes)
-     ):
-         # Process each group's matrices
-         for l_idx in range(l):
-             # Convert the scale factor tensor to blocked format
-             scale_a = to_blocked(sfa[:, :, l_idx])
-             scale_b = to_blocked(sfb[:, :, l_idx])
+     compiled_func = compile_kernel(problem_sizes)
   
-             # Compute GEMM using torch._scaled_mm
-             # (m, k) @ (n, k).T -> (m, n)
-             res = torch._scaled_mm(
-                 a[:, :, l_idx].view(torch.float4_e2m1fn_x2),
-                 b[:, :, l_idx].transpose(0, 1).view(torch.float4_e2m1fn_x2),
-                 scale_a.cuda(),
-                 scale_b.cuda(),
-                 bias=None,
-                 out_dtype=torch.float16,
-             )
-             c[:, :, l_idx] = res
+     # Extract raw data pointers from all input tensors for each group
+     # These will be passed to the GPU kernel to access the actual tensor data
+     abc_ptrs = []
+     sfasfb_ptrs = []
+     for i, ((a, b, c), (sfa_reordered, sfb_reordered), (m, n, k, l)) in enumerate(zip(abc_tensors, sfasfb_reordered_tensors, problem_sizes)):
+         # Store pointers to A, B, and C matrices for this group
+         abc_ptrs.append((a.data_ptr(), b.data_ptr(), c.data_ptr()))
+         # Store pointers to scale factor tensors for this group
+         sfasfb_ptrs.append((sfa_reordered.data_ptr(), sfb_reordered.data_ptr()))
   
-         result_tensors.append(c)
+     # Create torch tensor to store problem sizes for all groups
+     # Shape: (num_groups, 4) where each row contains (m, n, k, l) for that group
+     # Layout: (num_groups, 4):(4, 1) means row-major storage
+     tensor_of_problem_sizes = torch.tensor(
+         problem_sizes, dtype=torch.int32, device="cuda"
+     )
   
-     return result_tensors
+     # Create torch tensors to store data pointers for all groups
+     # These allow the GPU kernel to dynamically access different tensors per group
+     # tensor_of_abc_ptrs: Shape (num_groups, 3) containing (a_ptr, b_ptr, c_ptr) per group
+     # tensor_of_sfasfb_ptrs: Shape (num_groups, 2) containing (sfa_ptr, sfb_ptr) per group
+     tensor_of_abc_ptrs = torch.tensor(abc_ptrs, dtype=torch.int64, device="cuda")
+     tensor_of_sfasfb_ptrs = torch.tensor(sfasfb_ptrs, dtype=torch.int64, device="cuda")
+  
+     # Compute the tile shape for each CUDA Thread Block (CTA)
+     # cta_tile_shape_mn: [M_tile, N_tile] = [128, 128] for this kernel
+     cta_tile_shape_mn = [128, mma_tiler_mnk[1]]
+     # cluster_tile_shape_mn: Total tile shape per cluster (same as CTA since cluster is 1x1)
+     cluster_tile_shape_mn = tuple(
+         x * y for x, y in zip(cta_tile_shape_mn, (1, 1))
+     )
+     
+     # Compute total number of cluster tiles needed across all groups
+     # Each group's (m, n) dimensions are divided into tiles of size cluster_tile_shape_mn
+     # This determines the total grid size (bidz dimension) for kernel launch
+     total_num_clusters = 0
+     num_groups = len(problem_sizes)
+     for m, n, _, _ in problem_sizes:
+         # Calculate number of tiles needed in M and N dimensions for this group
+         num_clusters_mn = tuple(
+             (x + y - 1) // y for x, y in zip((m, n), cluster_tile_shape_mn)
+         )
+         # Multiply M_tiles * N_tiles to get total tiles for this group
+         total_num_clusters += functools.reduce(lambda x, y: x * y, num_clusters_mn)
+  
+     # Allocate device memory for tensormap descriptors
+     # Each cluster needs its own set of tensormaps (one for A, B, SFA, SFB)
+     # Shape: (total_num_clusters, num_tensormaps=4, bytes_per_tensormap/8=16)
+     # Tensormaps are hardware descriptors used by TMA for efficient memory transfers
+     tensormap_shape = (
+         total_num_clusters,
+         num_tensormaps,
+         bytes_per_tensormap // 8,
+     )
+     tensor_of_tensormap = torch.empty(tensormap_shape, dtype=torch.int64, device="cuda")
+  
+     # Create CuTe pointers to the metadata tensors that will be passed to the kernel
+     # These allow the GPU kernel to read problem sizes and tensor pointers
+     cute_ptr_of_tensor_of_abc_ptrs = make_ptr(
+         cutlass.Int64,
+         tensor_of_abc_ptrs.data_ptr(),
+         cute.AddressSpace.gmem,
+         assumed_align=16,
+     )
+     cute_ptr_of_tensor_of_sfasfb_ptrs = make_ptr(
+         cutlass.Int64,
+         tensor_of_sfasfb_ptrs.data_ptr(),
+         cute.AddressSpace.gmem,
+         assumed_align=16,
+     )
+     cute_ptr_of_tensor_of_problem_sizes = make_ptr(
+         cutlass.Int32,
+         tensor_of_problem_sizes.data_ptr(),
+         cute.AddressSpace.gmem,
+         assumed_align=16,
+     )
+     cute_ptr_of_tensor_of_tensormap = make_ptr(
+         cutlass.Int64,
+         tensor_of_tensormap.data_ptr(),
+         cute.AddressSpace.gmem,
+         assumed_align=16,
+     )
+  
+     # Launch the JIT-compiled GPU kernel with all prepared data
+     # The kernel will perform block-scaled group GEMM: C = A * SFA * B * SFB for all groups
+     compiled_func(
+         cute_ptr_of_tensor_of_problem_sizes, # Pointer to problem sizes array
+         cute_ptr_of_tensor_of_abc_ptrs,      # Pointer to ABC tensor pointers array
+         cute_ptr_of_tensor_of_sfasfb_ptrs,   # Pointer to scale factor pointers array
+         cute_ptr_of_tensor_of_tensormap,     # Pointer to tensormap buffer
+         total_num_clusters,                  # Total number of CTAs to launch
+         problem_sizes,                       # Problem sizes list (for host-side processing)
+         num_groups,                          # Number of groups in this batch
+     )
+  
+     res = []
+     for i in range(num_groups):
+         res.append(abc_tensors[i][2])
+     return res