submission 402102

  ⋯ 7 unchanged lines
  from cutlass.cute.runtime import make_ptr
   
  import functools
+ import dataclasses
  from typing import Tuple, List
   
  import torch
  from task import input_t, output_t
   
  # Kernel configuration parameters
+ # Size of tma descriptor in bytes
  bytes_per_tensormap = 128
+ # Number of tensormaps: a, b, sfa, sfb
  num_tensormaps = 4
- mma_tiler_mnk = (128, 128, 256)
+ # 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
- ab_dtype = cutlass.Float4E2M1FN
- sf_dtype = cutlass.Float8E4M3FN
- c_dtype = cutlass.Float16
- sf_vec_size = 16
+ # 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
  num_acc_stage = 1
- num_ab_stage = 4
+ num_ab_stage = 1
+ # Total number of columns in tmem
  num_tmem_alloc_cols = 512
   
- # Warp specialization configuration (6 warps total = 192 threads)
- EPILOG_WARP_IDS = (0, 1, 2, 3)
- MMA_WARP_ID = 4
- TMA_WARP_ID = 5
- THREADS_PER_CTA = 192  # 6 warps * 32 threads
   
-  
+ # Helper function for ceiling division
  def ceil_div(a, b):
      return (a + b - 1) // b
   
   
+ # The CuTe reference implementation for NVFP4 block-scaled GEMM
  @cute.kernel
  def kernel(
      tiled_mma: cute.TiledMma,
  ⋯ 17 unchanged lines
      num_tma_load_bytes: cutlass.Constexpr[int],
  ):
      """
-     Warp-specialized GPU kernel for NVFP4 Group GEMM.
-     - TMA warp (5): Producer for TMA loads
-     - MMA warp (4): Consumer of AB data, producer of accumulator
-     - Epilogue warps (0-3): Consumer of accumulator, store to global memory
+     GPU device kernel performing the Group GEMM computation.
      """
      warp_idx = cute.arch.warp_idx()
      warp_idx = cute.arch.make_warp_uniform(warp_idx)
      tidx, _, _ = cute.arch.thread_idx()
   
-     # Named barriers for synchronization
-     # tmem_alloc_barrier: 160 threads (warps 0-4, excludes TMA warp)
-     tmem_alloc_barrier = pipeline.NamedBarrier(
-         barrier_id=1,
-         num_threads=32 * 5,  # MMA + epilogue warps = 160 threads
-     )
-     # tensormap_ab_init_barrier: 64 threads (MMA + TMA warps)
-     tensormap_ab_init_barrier = pipeline.NamedBarrier(
-         barrier_id=2,
-         num_threads=64,  # MMA + TMA = 2 warps
-     )
-  
      #
      # Delinearize bidz to coord_x, coord_y and group_idx for each CTA
      #
  ⋯ 29 unchanged lines
          (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)
      )
  ⋯ 1 unchanged lines
      #
      # Define shared storage for kernel
      #
-     size_tensormap_in_i64 = num_tensormaps * bytes_per_tensormap // 8
-  
+     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]
+         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
-     tmem_holding_buf = storage.tmem_holding_buf
-  
-     # Setup smem tensors for A, B, SFA, SFB
+     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)
      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 pipelines
-     # AB pipeline: TMA warp is producer, MMA warp is consumer
-     ab_pipeline = pipeline.PipelineTmaUmma.create(
+     # 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=pipeline.CooperativeGroup(pipeline.Agent.Thread),
-         consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1),
+         producer_group=ab_pipeline_producer_group,
+         consumer_group=ab_pipeline_consumer_group,
          tx_count=num_tma_load_bytes,
-     )
-  
-     # ACC pipeline: MMA warp is producer, epilogue warps are consumers
-     acc_pipeline = pipeline.PipelineUmmaAsync.create(
+     ).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, 32 * len(EPILOG_WARP_IDS)),
-     )
+         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
  ⋯ 31 unchanged lines
      mB_nkl_layout = cute.make_layout(
          (n, k, l), stride=(cute.assume(k, 32), 1, cute.assume(n * k, 32),))
   
-     # SFA, SFB layout
+     # 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(
  ⋯ 11 unchanged lines
      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
+         )
+         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,
+             ),
+         )
+  
+         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()
+  
      #
-     # TMA Partition for A/B/SFA/SFB
+     # 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,
  ⋯ 1 unchanged lines
          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,
  ⋯ 1 unchanged lines
          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,
  ⋯ 3 unchanged lines
      )
      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,
  ⋯ 7 unchanged lines
      #
      # 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])
   
  ⋯ 1 unchanged lines
      # 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])]
   
      #
-     # ==================== TMA WARP (PRODUCER) ====================
+     # Main loop
      #
-     # Get raw mbarrier pointers from pipeline's sync objects
-     ab_full_mbar_ptr = ab_pipeline.sync_object_full.mbarrier_base
-     ab_empty_mbar_ptr = ab_pipeline.sync_object_empty.mbarrier_base
-     acc_full_mbar_ptr = acc_pipeline.sync_object_full.mbarrier_base
-     acc_empty_mbar_ptr = acc_pipeline.sync_object_empty.mbarrier_base
-     
-     if warp_idx == TMA_WARP_ID:
-         # Wait for MMA warp to initialize tensormaps
-         tensormap_ab_init_barrier.arrive_and_wait()
+     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)
+         # Execute k_tile loop
+         for k_tile in range(k_tile_cnt):
+             # Wait for AB buffer empty
+             ab_empty = ab_producer.acquire_and_advance()
   
-         # Raw mbarrier state tracking
-         tma_wr_k_tile = cutlass.Int32(0)
-         smem_wr_buffer = tma_wr_k_tile % num_ab_stage
-         tma_wr_phase = tma_wr_k_tile // num_ab_stage % 2
-  
-         # Peek (try_wait) AB buffer empty for first iteration
-         peek_ab_empty_status = cute.arch.mbarrier_conditional_try_wait(
-             tma_wr_k_tile < k_tile_cnt,
-             ab_empty_mbar_ptr + smem_wr_buffer,
-             tma_wr_phase ^ 1,  # empty phase is inverted
-         )
-  
-         for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):
-             # Calculate next iteration values
-             tma_wr_k_tile_next = tma_wr_k_tile + 1
-             smem_wr_buffer_next = tma_wr_k_tile_next % num_ab_stage
-             tma_wr_phase_next = tma_wr_k_tile_next // num_ab_stage % 2
-             
-             # Wait for AB buffer empty (conditionally)
-             if peek_ab_empty_status == 0:
-                 cute.arch.mbarrier_wait(
-                     ab_empty_mbar_ptr + smem_wr_buffer, tma_wr_phase ^ 1
-                 )
-             
-             # Arrive and expect tx bytes on full barrier
-             with cute.arch.elect_one():
-                 cute.arch.mbarrier_arrive_and_expect_tx(
-                     ab_full_mbar_ptr + smem_wr_buffer, num_tma_load_bytes
-                 )
-  
-             # TMA load A/B/SFA/SFB to shared memory
+             #  TMA load A/B/SFA/SFB to shared memory
              cute.copy(
                  tma_atom_a,
-                 tAgA[(None, tma_wr_k_tile)],
-                 tAsA[(None, smem_wr_buffer)],
-                 tma_bar_ptr=ab_full_mbar_ptr + smem_wr_buffer,
+                 tAgA[(None, k_tile)],
+                 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,
  ⋯ 1 unchanged lines
              )
              cute.copy(
                  tma_atom_b,
-                 tBgB[(None, tma_wr_k_tile)],
-                 tBsB[(None, smem_wr_buffer)],
-                 tma_bar_ptr=ab_full_mbar_ptr + smem_wr_buffer,
+                 tBgB[(None, k_tile)],
+                 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,
  ⋯ 1 unchanged lines
              )
              cute.copy(
                  tma_atom_sfa,
-                 tAgSFA[(None, tma_wr_k_tile)],
-                 tAsSFA[(None, smem_wr_buffer)],
-                 tma_bar_ptr=ab_full_mbar_ptr + smem_wr_buffer,
+                 tAgSFA[(None, k_tile)],
+                 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,
  ⋯ 1 unchanged lines
              )
              cute.copy(
                  tma_atom_sfb,
-                 tBgSFB[(None, tma_wr_k_tile)],
-                 tBsSFB[(None, smem_wr_buffer)],
-                 tma_bar_ptr=ab_full_mbar_ptr + smem_wr_buffer,
+                 tBgSFB[(None, k_tile)],
+                 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,
                  ),
              )
   
-             # Peek (try_wait) AB buffer empty for next iteration
-             peek_ab_empty_status = cute.arch.mbarrier_conditional_try_wait(
-                 tma_wr_k_tile_next < k_tile_cnt,
-                 ab_empty_mbar_ptr + smem_wr_buffer_next,
-                 tma_wr_phase_next ^ 1,
-             )
-             
-             tma_wr_k_tile = tma_wr_k_tile_next
-             smem_wr_buffer = smem_wr_buffer_next
-             tma_wr_phase = tma_wr_phase_next
+             # Wait for AB buffer full
+             ab_full = ab_consumer.wait_and_advance()
   
-         # TMA warp done - no explicit tail needed
-  
-     #
-     # ==================== MMA WARP (CONSUMER/PRODUCER) ====================
-     #
-     if warp_idx == MMA_WARP_ID:
-         # Initialize tensormaps for A, B, SFA and SFB
-         tensormap_manager.init_tensormap_from_atom(
-             tma_atom_a, tensormap_a_smem_ptr, MMA_WARP_ID
-         )
-         tensormap_manager.init_tensormap_from_atom(
-             tma_atom_b, tensormap_b_smem_ptr, MMA_WARP_ID
-         )
-         tensormap_manager.init_tensormap_from_atom(
-             tma_atom_sfa, tensormap_sfa_smem_ptr, MMA_WARP_ID
-         )
-         tensormap_manager.init_tensormap_from_atom(
-             tma_atom_sfb, tensormap_sfb_smem_ptr, MMA_WARP_ID
-         )
-  
-         # Update tensormaps with actual tensor information
-         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),
-             MMA_WARP_ID,
-             (tensormap_a_smem_ptr, tensormap_b_smem_ptr,
-              tensormap_sfa_smem_ptr, tensormap_sfb_smem_ptr),
-         )
-  
-         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)
-  
-         # Prefetch TMA descriptors
-         cpasync.prefetch_descriptor(tma_atom_a)
-         cpasync.prefetch_descriptor(tma_atom_b)
-         cpasync.prefetch_descriptor(tma_atom_sfa)
-         cpasync.prefetch_descriptor(tma_atom_sfb)
-  
-         # Signal TMA warp that tensormaps are ready
-         tensormap_ab_init_barrier.arrive_and_wait()
-  
-         # Sync for TMEM allocation
-         tmem_alloc_barrier.arrive_and_wait()
-  
-         # Retrieve TMEM pointer
-         acc_tmem_ptr = cute.arch.retrieve_tmem_ptr(
-             cutlass.Float32,
-             alignment=16,
-             ptr_to_buffer_holding_addr=tmem_holding_buf,
-         )
-         tCtAcc = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
-  
-         # Make SFA/SFB tmem tensors
-         sfa_tmem_ptr = cute.recast_ptr(
-             acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc),
-             dtype=sf_dtype,
-         )
-         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)
-  
-         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,
-         )
-         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)
-  
-         # S2T copy setup for SFA/SFB
-         copy_atom_s2t = cute.make_copy_atom(
-             tcgen05.Cp4x32x128bOp(tcgen05.CtaGroup.ONE),
-             sf_dtype,
-         )
-  
-         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)
-         tCsSFA_compact_s2t_ = thr_copy_s2t_sfa.partition_S(tCsSFA_compact)
-         tCsSFA_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
-             tiled_copy_s2t_sfa, tCsSFA_compact_s2t_
-         )
-         tCtSFA_compact_s2t = thr_copy_s2t_sfa.partition_D(tCtSFA_compact)
-  
-         tCsSFB_compact = cute.filter_zeros(sSFB)
-         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)
-         tCsSFB_compact_s2t_ = thr_copy_s2t_sfb.partition_S(tCsSFB_compact)
-         tCsSFB_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
-             tiled_copy_s2t_sfb, tCsSFB_compact_s2t_
-         )
-         tCtSFB_compact_s2t = thr_copy_s2t_sfb.partition_D(tCtSFB_compact)
-  
-         # ACC state tracking (single stage, so always use index 0)
-         acc_stage_idx = cutlass.Int32(0)
-  
-         # Raw mbarrier state tracking for AB consumer
-         mma_rd_k_tile = cutlass.Int32(0)
-         smem_rd_buffer = mma_rd_k_tile % num_ab_stage
-         mma_rd_phase = mma_rd_k_tile // num_ab_stage % 2
-  
-         # Peek (try_wait) AB buffer full for first iteration
-         peek_ab_full_status = cute.arch.mbarrier_conditional_try_wait(
-             mma_rd_k_tile < k_tile_cnt,
-             ab_full_mbar_ptr + smem_rd_buffer,
-             mma_rd_phase,  # full phase (no inversion)
-         )
-  
-         # Wait for accumulator buffer empty using raw mbarrier
-         # Phase 0 ^ 1 = 1 for producer waiting on empty (consumer signals)
-         cute.arch.mbarrier_wait(acc_empty_mbar_ptr + acc_stage_idx, cutlass.Int32(1))
-  
-         # Set ACCUMULATE field to False for the first k_tile iteration
-         tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
-  
-         # MMA mainloop
-         for k_tile in range(k_tile_cnt):
-             # Calculate next iteration values
-             mma_rd_k_tile_next = cutlass.Int32(k_tile + 1)
-             smem_rd_buffer_next = mma_rd_k_tile_next % num_ab_stage
-             mma_rd_phase_next = mma_rd_k_tile_next // num_ab_stage % 2
-             
-             # Conditionally wait for AB buffer full
-             if peek_ab_full_status == 0:
-                 cute.arch.mbarrier_wait(
-                     ab_full_mbar_ptr + smem_rd_buffer, mma_rd_phase
-                 )
-  
-             # Copy SFA/SFB from shared memory to TMEM
-             s2t_stage_coord = (None, None, None, None, smem_rd_buffer)
+             #  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(
  ⋯ 10 unchanged lines
              # 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, smem_rd_buffer)
+                 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)
+                 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,
  ⋯ 5 unchanged lines
                  # Enable accumulate on tCtAcc after first kblock
                  tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
   
-             # Release AB buffer for TMA warp using tcgen05.commit (async after MMA)
-             with cute.arch.elect_one():
-                 tcgen05.commit(ab_empty_mbar_ptr + smem_rd_buffer)
+             # Async arrive AB buffer empty
+             ab_full.release()
+         acc_empty.commit()
   
-             # Peek for next iteration
-             peek_ab_full_status = cute.arch.mbarrier_conditional_try_wait(
-                 mma_rd_k_tile_next < k_tile_cnt,
-                 ab_full_mbar_ptr + smem_rd_buffer_next,
-                 mma_rd_phase_next,
-             )
-             
-             mma_rd_k_tile = mma_rd_k_tile_next
-             smem_rd_buffer = smem_rd_buffer_next
-             mma_rd_phase = mma_rd_phase_next
-  
-         # Commit accumulator for epilogue warps using tcgen05.commit (async after MMA)
-         with cute.arch.elect_one():
-             tcgen05.commit(acc_full_mbar_ptr + acc_stage_idx)
-  
      #
-     # ==================== EPILOGUE WARPS (CONSUMER) ====================
+     # Epilogue
+     # Partition for epilogue
      #
-     if warp_idx < MMA_WARP_ID:
-         # Allocate TMEM (only first epilogue warp does allocation)
-         if warp_idx == EPILOG_WARP_IDS[0]:
-             cute.arch.alloc_tmem(
-                 num_tmem_alloc_cols,
-                 tmem_holding_buf,
-                 is_two_cta=False,
-             )
+     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])
   
-         # Sync for TMEM allocation
-         tmem_alloc_barrier.arrive_and_wait()
+     # (TmemCpy, NumTmemCpy)
+     tDrAcc = cute.make_rmem_tensor(tDgC.shape, cutlass.Float32)
+     # (TmemCpy, NumTmemCpy)
+     tDrC = cute.make_rmem_tensor(tDgC.shape, c_dtype)
   
-         # Retrieve TMEM pointer
-         acc_tmem_ptr = cute.arch.retrieve_tmem_ptr(
-             cutlass.Float32,
-             alignment=16,
-             ptr_to_buffer_holding_addr=tmem_holding_buf,
-         )
-         tCtAcc = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
+     # Release TMEM allocation lock
+     tmem.relinquish_alloc_permit()
+     # Wait for accumulator buffer full
+     acc_full = acc_consumer.wait_and_advance()
   
-         # 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)
-         tDtAcc = thr_copy_t2r.partition_S(tCtAcc[None, 0, 0])
-         tDgC = thr_copy_t2r.partition_D(tCgC[None, 0, 0])
+     # Copy accumulator to register
+     cute.copy(tiled_copy_t2r, tDtAcc, tDrAcc)
+     acc_vec = tDrAcc.load()
+     tDrC.store(acc_vec.to(c_dtype))
   
-         tDrAcc = cute.make_rmem_tensor(tDgC.shape, cutlass.Float32)
-         tDrC = cute.make_rmem_tensor(tDgC.shape, 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)
   
-         # Wait for accumulator buffer full using raw mbarrier
-         # Phase 0 for consumer waiting on full (producer signals)
-         cute.arch.mbarrier_wait(acc_full_mbar_ptr, cutlass.Int32(0))
+     # ((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))
   
-         # Copy accumulator to register
-         cute.copy(tiled_copy_t2r, tDtAcc, tDrAcc)
-         acc_vec = tDrAcc.load()
-         tDrC.store(acc_vec.to(c_dtype))
-  
-         # STG Atom for global memory store
-         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)
-         tDcC = thr_copy_r2g.partition_D(cC)
-  
-         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)):
-             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))
-  
-         # Release accumulator buffer using raw mbarrier
-         cute.arch.mbarrier_arrive(acc_empty_mbar_ptr)
-  
-         # Deallocate TMEM
-         if warp_idx == EPILOG_WARP_IDS[0]:
-             cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False)
-         cute.arch.barrier()
-         if warp_idx == EPILOG_WARP_IDS[0]:
-             cute.arch.dealloc_tmem(acc_tmem_ptr, num_tmem_alloc_cols, is_two_cta=False)
-  
+     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,
  ⋯ 1 unchanged lines
      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: 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))
      )
  ⋯ 8 unchanged lines
      )
   
      # 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_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=(
  ⋯ 4 unchanged lines
          ),
      )
      initial_b = cute.make_tensor(
-         cute.make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16),
+         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=(
  ⋯ 4 unchanged lines
          ),
      )
   
-     # Setup sfa/sfb tensor
-     sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(initial_a.shape, sf_vec_size)
-     sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(initial_b.shape, sf_vec_size)
+     # 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)
+         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)
+         cute.make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=16,), sfb_layout)
   
      # Select MMA operation
      mma_op = tcgen05.MmaMXF4NVF4Op(
  ⋯ 11 unchanged lines
   
      # Compute A/B/SFA/SFB/C shared memory layout
      a_smem_layout_staged = sm100_utils.make_smem_layout_a(
-         tiled_mma, mma_tiler_mnk, ab_dtype, num_ab_stage,
+         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,
+         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,
+         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,
+         tiled_mma,
+         mma_tiler_mnk,
+         sf_vec_size,
+         num_ab_stage,
      )
      atom_thr_size = cute.size(tiled_mma.thr_id.shape)
   
  ⋯ 18 unchanged lines
          cluster_layout_vmnk.shape,
      )
      # Setup TMA for SFA
-     sfa_smem_layout = cute.slice_(sfa_smem_layout_staged, (None, None, None, 0))
+     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,
  ⋯ 4 unchanged lines
          internal_type=cutlass.Int16,
      )
      # Setup TMA for SFB
-     sfb_smem_layout = cute.slice_(sfb_smem_layout_staged, (None, None, None, 0))
+     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,
  ⋯ 13 unchanged lines
          a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size
      ) * atom_thr_size
   
-     # Store CTA shape information for each Group
+     # 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])
  ⋯ 4 unchanged lines
   
      # Launch the kernel
      kernel(
-         tiled_mma,
-         tma_atom_a,
-         tma_tensor_a,
-         tma_atom_b,
-         tma_tensor_b,
-         tma_atom_sfa,
-         tma_tensor_sfa,
-         tma_atom_sfb,
-         tma_tensor_sfb,
-         tensor_of_abc_ptrs,
-         tensor_of_sfasfb_ptrs,
-         tensor_of_tensormap,
-         tensor_of_problem_sizes,
-         a_smem_layout_staged,
-         b_smem_layout_staged,
-         sfa_smem_layout_staged,
-         sfb_smem_layout_staged,
-         cta_mn_list,
-         num_tma_load_bytes,
+         # 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],
+         block=[threads_per_cta, 1, 1],
          cluster=(1, 1, 1),
      )
      return
   
   
- # Global cache for compiled kernels
+ # 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):
-     """Compile the kernel once and cache it using problem_sizes as the key."""
+     """
+     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]
   
  ⋯ 6 unchanged lines
      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,
      )
  ⋯ 7 unchanged lines
          problem_sizes,
          num_groups
      )
+     # Store compiled kernel in cache with problem_sizes as key
      _compiled_kernel_cache[cache_key] = compiled_func
      return compiled_func
   
   
+ @dataclasses.dataclass
+ class _RuntimeCacheEntry:
+     compiled_func: object
+     problem_sizes: tuple[tuple[int, int, int, int], ...]
+     num_groups: int
+     total_num_clusters: int
+     tensor_of_problem_sizes: torch.Tensor
+     tensor_of_abc_ptrs: torch.Tensor
+     tensor_of_sfasfb_ptrs: torch.Tensor
+     tensor_of_tensormap: torch.Tensor
+     abc_ptrs_host: torch.Tensor
+     sfasfb_ptrs_host: torch.Tensor
+     last_abc_ptrs: tuple[tuple[int, int, int], ...] | None = None
+     last_sfasfb_ptrs: tuple[tuple[int, int], ...] | None = None
+  
+  
+ _runtime_cache: dict[tuple[tuple[int, int, int, int], ...], _RuntimeCacheEntry] = {}
+  
+  
+ def _get_or_create_entry(problem_sizes: list[tuple[int, int, int, int]]) -> _RuntimeCacheEntry:
+     ps_key = tuple(tuple(int(x) for x in ps) for ps in problem_sizes)
+     entry = _runtime_cache.get(ps_key)
+     if entry is not None:
+         return entry
+  
+     num_groups = len(problem_sizes)
+     total_num_clusters = 0
+     for m, n, _, _ in problem_sizes:
+         total_num_clusters += ceil_div(m, mma_tiler_mnk[0]) * ceil_div(n, mma_tiler_mnk[1])
+  
+     tensormap_shape = (total_num_clusters, num_tensormaps, bytes_per_tensormap // 8)
+  
+     tensor_of_problem_sizes = torch.tensor(ps_key, dtype=torch.int32, device="cuda")
+     tensor_of_abc_ptrs = torch.empty((num_groups, 3), dtype=torch.int64, device="cuda")
+     tensor_of_sfasfb_ptrs = torch.empty((num_groups, 2), dtype=torch.int64, device="cuda")
+     tensor_of_tensormap = torch.empty(tensormap_shape, dtype=torch.int64, device="cuda")
+  
+     abc_ptrs_host = torch.empty((num_groups, 3), dtype=torch.int64, device="cpu", pin_memory=True)
+     sfasfb_ptrs_host = torch.empty((num_groups, 2), dtype=torch.int64, device="cpu", pin_memory=True)
+  
+     compiled_func = compile_kernel(problem_sizes)
+  
+     entry = _RuntimeCacheEntry(
+         compiled_func=compiled_func,
+         problem_sizes=ps_key,
+         num_groups=num_groups,
+         total_num_clusters=total_num_clusters,
+         tensor_of_problem_sizes=tensor_of_problem_sizes,
+         tensor_of_abc_ptrs=tensor_of_abc_ptrs,
+         tensor_of_sfasfb_ptrs=tensor_of_sfasfb_ptrs,
+         tensor_of_tensormap=tensor_of_tensormap,
+         abc_ptrs_host=abc_ptrs_host,
+         sfasfb_ptrs_host=sfasfb_ptrs_host,
+     )
+     _runtime_cache[ps_key] = entry
+     return entry
+  
+  
+ def _maybe_update_ptrs(entry: _RuntimeCacheEntry, abc_tensors, sfasfb_reordered_tensors) -> None:
+     abc_ptrs = tuple(
+         (int(a.data_ptr()), int(b.data_ptr()), int(c.data_ptr())) for (a, b, c) in abc_tensors
  ⋯ diff truncated