submission 124468

# NVFP4 Block-Scaled GEMM - Persistent kernel with warp specialization
# Adapted from NVIDIA CUTLASS example


from typing import Union


import torch
from task import input_t, output_t


import cutlass
import cutlass.cute as cute
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.utils.blockscaled_layout as blockscaled_utils
from cutlass.cute.runtime import make_ptr




# Kernel configuration parameters
# Tile sizes for M, N dimensions (K computed dynamically)
mma_tiler_mn = (256, 128)
# 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
# Cluster shape
cluster_shape_mn = (2, 1)
# Use 2-CTA instructions (for mma_tiler_mn[0] == 256)
use_2cta_instrs = mma_tiler_mn[0] == 256
cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE


# Warp specialization IDs
epilog_warp_id = (0, 1, 2, 3)
mma_warp_id = 4
tma_warp_id = 5
threads_per_cta = 32 * len((mma_warp_id, tma_warp_id, *epilog_warp_id))


# Barriers
epilog_sync_barrier = pipeline.NamedBarrier(
   barrier_id=1,
   num_threads=32 * len(epilog_warp_id),
)
tmem_alloc_barrier = pipeline.NamedBarrier(
   barrier_id=2,
   num_threads=32 * len((mma_warp_id, *epilog_warp_id)),
)


# Memory configuration
smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
num_tmem_alloc_cols = 512
occupancy = 1
acc_dtype = cutlass.Float32




def ceil_div(a, b):
   return (a + b - 1) // b




@cute.kernel
def kernel(
   tiled_mma: cute.TiledMma,
   tiled_mma_sfb: cute.TiledMma,
   tma_atom_a: cute.CopyAtom,
   mA_mkl: cute.Tensor,
   tma_atom_b: cute.CopyAtom,
   mB_nkl: cute.Tensor,
   tma_atom_sfa: cute.CopyAtom,
   mSFA_mkl: cute.Tensor,
   tma_atom_sfb: cute.CopyAtom,
   mSFB_nkl: cute.Tensor,
   tma_atom_c: cute.CopyAtom,
   mC_mnl: cute.Tensor,
   cluster_layout_vmnk: cute.Layout,
   cluster_layout_sfb_vmnk: cute.Layout,
   a_smem_layout_staged: cute.ComposedLayout,
   b_smem_layout_staged: cute.ComposedLayout,
   sfa_smem_layout_staged: cute.Layout,
   sfb_smem_layout_staged: cute.Layout,
   c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout],
   epi_tile: cute.Tile,
   tile_sched_params: utils.PersistentTileSchedulerParams,
   mma_tiler: cutlass.Constexpr,
   mma_tiler_sfb: cutlass.Constexpr,
   cta_tile_shape_mnk: cutlass.Constexpr,
   num_ab_stage: cutlass.Constexpr[int],
   num_acc_stage: cutlass.Constexpr[int],
   num_c_stage: cutlass.Constexpr[int],
   num_tma_load_bytes: cutlass.Constexpr[int],
   num_mcast_ctas_a: cutlass.Constexpr[int],
   num_mcast_ctas_b: cutlass.Constexpr[int],
   is_a_mcast: cutlass.Constexpr[bool],
   is_b_mcast: cutlass.Constexpr[bool],
   overlapping_accum: cutlass.Constexpr[bool],
   num_accumulator_tmem_cols: cutlass.Constexpr[int],
   num_sfa_tmem_cols: cutlass.Constexpr[int],
   num_sf_tmem_cols: cutlass.Constexpr[int],
   epi_tile_n: cutlass.Constexpr[int],
   iter_acc_early_release_in_epilogue: cutlass.Constexpr[int],
   c_layout: cutlass.Constexpr,
   shared_storage: cutlass.Constexpr,
):
   """GPU device kernel with warp specialization and persistent scheduling."""
   warp_idx = cute.arch.warp_idx()
   warp_idx = cute.arch.make_warp_uniform(warp_idx)


   # Prefetch TMA descriptors
   if warp_idx == tma_warp_id:
       cpasync.prefetch_descriptor(tma_atom_a)
       cpasync.prefetch_descriptor(tma_atom_b)
       cpasync.prefetch_descriptor(tma_atom_sfa)
       cpasync.prefetch_descriptor(tma_atom_sfb)
       cpasync.prefetch_descriptor(tma_atom_c)


   use_2cta = cute.size(tiled_mma.thr_id.shape) == 2


   bidx, bidy, bidz = cute.arch.block_idx()
   mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape)
   is_leader_cta = mma_tile_coord_v == 0
   cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
   block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
   block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord(cta_rank_in_cluster)
   tidx, _, _ = cute.arch.thread_idx()


   # Allocate shared storage
   smem = utils.SmemAllocator()
   storage = smem.allocate(shared_storage)


   # Initialize pipelines
   ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
   num_tma_producer = num_mcast_ctas_a + num_mcast_ctas_b - 1
   ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_tma_producer)
   ab_pipeline = pipeline.PipelineTmaUmma.create(
       barrier_storage=storage.ab_full_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,
       cta_layout_vmnk=cluster_layout_vmnk,
       defer_sync=True,
   )


   acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
   num_acc_consumer_threads = len(epilog_warp_id) * (2 if use_2cta else 1)
   acc_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_acc_consumer_threads)
   acc_pipeline = pipeline.PipelineUmmaAsync.create(
       barrier_storage=storage.acc_full_mbar_ptr.data_ptr(),
       num_stages=num_acc_stage,
       producer_group=acc_pipeline_producer_group,
       consumer_group=acc_pipeline_consumer_group,
       cta_layout_vmnk=cluster_layout_vmnk,
       defer_sync=True,
   )


   tmem = utils.TmemAllocator(
       storage.tmem_holding_buf,
       barrier_for_retrieve=tmem_alloc_barrier,
       allocator_warp_id=epilog_warp_id[0],
       is_two_cta=use_2cta,
       two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
   )


   pipeline_init_arrive(cluster_shape_mn=cluster_shape_mn, is_relaxed=True)


   # Setup SMEM tensors
   sC = storage.sC.get_tensor(c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner)
   sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner)
   sB = storage.sB.get_tensor(b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner)
   sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged)
   sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged)


   # Multicast masks
   a_full_mcast_mask = None
   b_full_mcast_mask = None
   sfa_full_mcast_mask = None
   sfb_full_mcast_mask = None
   if cutlass.const_expr(is_a_mcast or is_b_mcast or use_2cta):
       a_full_mcast_mask = cpasync.create_tma_multicast_mask(
           cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
       )
       b_full_mcast_mask = cpasync.create_tma_multicast_mask(
           cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1
       )
       sfa_full_mcast_mask = cpasync.create_tma_multicast_mask(
           cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
       )
       sfb_full_mcast_mask = cpasync.create_tma_multicast_mask(
           cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_vmnk, mcast_mode=1
       )


   # Partition global tensors
   gA_mkl = cute.local_tile(mA_mkl, cute.slice_(mma_tiler, (None, 0, None)), (None, None, None))
   gB_nkl = cute.local_tile(mB_nkl, cute.slice_(mma_tiler, (0, None, None)), (None, None, None))
   gSFA_mkl = cute.local_tile(mSFA_mkl, cute.slice_(mma_tiler, (None, 0, None)), (None, None, None))
   gSFB_nkl = cute.local_tile(mSFB_nkl, cute.slice_(mma_tiler_sfb, (0, None, None)), (None, None, None))
   gC_mnl = cute.local_tile(mC_mnl, cute.slice_(mma_tiler, (None, None, 0)), (None, None, None))
   k_tile_cnt = cute.size(gA_mkl, mode=[3])


   thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
   thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v)
   tCgA = thr_mma.partition_A(gA_mkl)
   tCgB = thr_mma.partition_B(gB_nkl)
   tCgSFA = thr_mma.partition_A(gSFA_mkl)
   tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl)
   tCgC = thr_mma.partition_C(gC_mnl)


   # TMA partitions
   a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape)
   tAsA, tAgA = cpasync.tma_partition(
       tma_atom_a, block_in_cluster_coord_vmnk[2], a_cta_layout,
       cute.group_modes(sA, 0, 3), cute.group_modes(tCgA, 0, 3),
   )
   b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape)
   tBsB, tBgB = cpasync.tma_partition(
       tma_atom_b, block_in_cluster_coord_vmnk[1], b_cta_layout,
       cute.group_modes(sB, 0, 3), cute.group_modes(tCgB, 0, 3),
   )
   sfa_cta_layout = a_cta_layout
   tAsSFA, tAgSFA = cpasync.tma_partition(
       tma_atom_sfa, block_in_cluster_coord_vmnk[2], sfa_cta_layout,
       cute.group_modes(sSFA, 0, 3), cute.group_modes(tCgSFA, 0, 3),
   )
   tAsSFA = cute.filter_zeros(tAsSFA)
   tAgSFA = cute.filter_zeros(tAgSFA)
   sfb_cta_layout = cute.make_layout(cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape)
   tBsSFB, tBgSFB = cpasync.tma_partition(
       tma_atom_sfb, block_in_cluster_coord_sfb_vmnk[1], sfb_cta_layout,
       cute.group_modes(sSFB, 0, 3), cute.group_modes(tCgSFB, 0, 3),
   )
   tBsSFB = cute.filter_zeros(tBsSFB)
   tBgSFB = cute.filter_zeros(tBgSFB)


   tCrA = tiled_mma.make_fragment_A(sA)
   tCrB = tiled_mma.make_fragment_B(sB)
   acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2])
   if cutlass.const_expr(overlapping_accum):
       num_acc_stage_overlapped = 2
       tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, num_acc_stage_overlapped))
       tCtAcc_fake = cute.make_tensor(
           tCtAcc_fake.iterator,
           cute.make_layout(
               tCtAcc_fake.shape,
               stride=(
                   tCtAcc_fake.stride[0],
                   tCtAcc_fake.stride[1],
                   tCtAcc_fake.stride[2],
                   (256 - num_sf_tmem_cols) * tCtAcc_fake.stride[0][1]
               )
           )
       )
   else:
       tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, num_acc_stage))


   pipeline_init_wait(cluster_shape_mn=cluster_shape_mn)


   # TMA warp
   if warp_idx == tma_warp_id:
       tile_sched = utils.StaticPersistentTileScheduler.create(
           tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
       )
       work_tile = tile_sched.initial_work_tile_info()
       ab_producer_state = pipeline.make_pipeline_state(
           pipeline.PipelineUserType.Producer, num_ab_stage
       )


       while work_tile.is_valid_tile:
           cur_tile_coord = work_tile.tile_idx
           mma_tile_coord_mnl = (
               cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),
               cur_tile_coord[1],
               cur_tile_coord[2],
           )
           tAgA_slice = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
           tBgB_slice = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
           tAgSFA_slice = tAgSFA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
           slice_n = mma_tile_coord_mnl[1]
           if cutlass.const_expr(cta_tile_shape_mnk[1] == 64):
               slice_n = mma_tile_coord_mnl[1] // 2
           tBgSFB_slice = tBgSFB[(None, slice_n, None, mma_tile_coord_mnl[2])]


           ab_producer_state.reset_count()
           peek_ab_empty_status = cutlass.Boolean(1)
           if ab_producer_state.count < k_tile_cnt:
               peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state)


           for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):
               ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status)
               cute.copy(tma_atom_a, tAgA_slice[(None, ab_producer_state.count)],
                        tAsA[(None, ab_producer_state.index)],
                        tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
                        mcast_mask=a_full_mcast_mask)
               cute.copy(tma_atom_b, tBgB_slice[(None, ab_producer_state.count)],
                        tBsB[(None, ab_producer_state.index)],
                        tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
                        mcast_mask=b_full_mcast_mask)
               cute.copy(tma_atom_sfa, tAgSFA_slice[(None, ab_producer_state.count)],
                        tAsSFA[(None, ab_producer_state.index)],
                        tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
                        mcast_mask=sfa_full_mcast_mask)
               cute.copy(tma_atom_sfb, tBgSFB_slice[(None, ab_producer_state.count)],
                        tBsSFB[(None, ab_producer_state.index)],
                        tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
                        mcast_mask=sfb_full_mcast_mask)
               ab_producer_state.advance()
               peek_ab_empty_status = cutlass.Boolean(1)
               if ab_producer_state.count < k_tile_cnt:
                   peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state)


           tile_sched.advance_to_next_work()
           work_tile = tile_sched.get_current_work()


       ab_pipeline.producer_tail(ab_producer_state)


   # MMA warp
   if warp_idx == mma_warp_id:
       tmem.wait_for_alloc()
       acc_tmem_ptr = tmem.retrieve_ptr(acc_dtype)
       tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)


       sfa_tmem_ptr = cute.recast_ptr(acc_tmem_ptr + num_accumulator_tmem_cols, dtype=sf_dtype)
       tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
           tiled_mma, mma_tiler, 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 + num_accumulator_tmem_cols + num_sfa_tmem_cols, dtype=sf_dtype
       )
       tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
           tiled_mma, mma_tiler, sf_vec_size,
           cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)),
       )
       tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout)


       # S2T copy for SFA
       tCsSFA_compact = cute.filter_zeros(sSFA)
       tCtSFA_compact = cute.filter_zeros(tCtSFA)
       copy_atom_s2t_sfa = cute.make_copy_atom(tcgen05.Cp4x32x128bOp(cta_group), sf_dtype)
       tiled_copy_s2t_sfa = tcgen05.make_s2t_copy(copy_atom_s2t_sfa, 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)


       # S2T copy for SFB
       tCsSFB_compact = cute.filter_zeros(sSFB)
       tCtSFB_compact = cute.filter_zeros(tCtSFB)
       copy_atom_s2t_sfb = cute.make_copy_atom(tcgen05.Cp4x32x128bOp(cta_group), sf_dtype)
       tiled_copy_s2t_sfb = tcgen05.make_s2t_copy(copy_atom_s2t_sfb, 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)


       tile_sched = utils.StaticPersistentTileScheduler.create(
           tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
       )
       work_tile = tile_sched.initial_work_tile_info()
       ab_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, num_ab_stage)
       acc_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, num_acc_stage)


       while work_tile.is_valid_tile:
           cur_tile_coord = work_tile.tile_idx
           mma_tile_coord_mnl = (
               cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),
               cur_tile_coord[1],
               cur_tile_coord[2],
           )
           if cutlass.const_expr(overlapping_accum):
               acc_stage_index = acc_producer_state.phase ^ 1
           else:
               acc_stage_index = acc_producer_state.index
           tCtAcc = tCtAcc_base[(None, None, None, acc_stage_index)]


           ab_consumer_state.reset_count()
           peek_ab_full_status = cutlass.Boolean(1)
           if ab_consumer_state.count < k_tile_cnt and is_leader_cta:
               peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state)


           if is_leader_cta:
               acc_pipeline.producer_acquire(acc_producer_state)


           tCtSFB_mma = tCtSFB
           if cutlass.const_expr(cta_tile_shape_mnk[1] == 192):
               offset = cutlass.Int32(2) if mma_tile_coord_mnl[1] % 2 == 1 else cutlass.Int32(0)
               shifted_ptr = cute.recast_ptr(
                   acc_tmem_ptr + num_accumulator_tmem_cols + num_sfa_tmem_cols + offset,
                   dtype=sf_dtype,
               )
               tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout)
           elif cutlass.const_expr(cta_tile_shape_mnk[1] == 64):
               offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2)
               shifted_ptr = cute.recast_ptr(
                   acc_tmem_ptr + num_accumulator_tmem_cols + num_sfa_tmem_cols + offset,
                   dtype=sf_dtype,
               )
               tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout)


           tiled_mma.set(tcgen05.Field.ACCUMULATE, False)


           for k_tile in range(k_tile_cnt):
               if is_leader_cta:
                   ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status)
                   s2t_stage_coord = (None, None, None, None, ab_consumer_state.index)
                   cute.copy(tiled_copy_s2t_sfa, tCsSFA_compact_s2t[s2t_stage_coord], tCtSFA_compact_s2t)
                   cute.copy(tiled_copy_s2t_sfb, tCsSFB_compact_s2t[s2t_stage_coord], tCtSFB_compact_s2t)


                   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_consumer_state.index)
                       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_mma[sf_kblock_coord].iterator)
                       cute.gemm(tiled_mma, tCtAcc, tCrA[kblock_coord], tCrB[kblock_coord], tCtAcc)
                       tiled_mma.set(tcgen05.Field.ACCUMULATE, True)


                   ab_pipeline.consumer_release(ab_consumer_state)


               ab_consumer_state.advance()
               peek_ab_full_status = cutlass.Boolean(1)
               if ab_consumer_state.count < k_tile_cnt and is_leader_cta:
                   peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state)


           if is_leader_cta:
               acc_pipeline.producer_commit(acc_producer_state)
           acc_producer_state.advance()


           tile_sched.advance_to_next_work()
           work_tile = tile_sched.get_current_work()


       acc_pipeline.producer_tail(acc_producer_state)


   # Epilogue warps
   if warp_idx < mma_warp_id:
       tmem.allocate(num_tmem_alloc_cols)
       tmem.wait_for_alloc()
       acc_tmem_ptr = tmem.retrieve_ptr(acc_dtype)
       tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)


       # Epilogue TMEM copy setup
       copy_atom_t2r = sm100_utils.get_tmem_load_op(
           cta_tile_shape_mnk, c_layout, c_dtype, acc_dtype, epi_tile, use_2cta,
       )
       tAcc_epi = cute.flat_divide(tCtAcc_base[((None, None), 0, 0, None)], epi_tile)
       tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)])
       thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
       tTR_tAcc_base = thr_copy_t2r.partition_S(tAcc_epi)
       gC_mnl_epi = cute.flat_divide(tCgC[((None, None), 0, 0, None, None, None)], epi_tile)
       tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi)
       tTR_rAcc = cute.make_rmem_tensor(tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, acc_dtype)
       tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, c_dtype)


       # R2S copy setup
       copy_atom_r2s = sm100_utils.get_smem_store_op(c_layout, c_dtype, acc_dtype, tiled_copy_t2r)
       tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r)
       thr_copy_r2s = tiled_copy_r2s.get_slice(tidx)
       tRS_sC = thr_copy_r2s.partition_D(sC)
       tRS_rC = tiled_copy_r2s.retile(tTR_rC)


       # GMEM copy setup
       gC_epi = cute.flat_divide(tCgC[((None, None), 0, 0, None, None, None)], epi_tile)
       sC_for_tma = cute.group_modes(sC, 0, 2)
       gC_for_tma = cute.group_modes(gC_epi, 0, 2)
       bSG_sC, bSG_gC_partitioned = cpasync.tma_partition(tma_atom_c, 0, cute.make_layout(1), sC_for_tma, gC_for_tma)


       tile_sched = utils.StaticPersistentTileScheduler.create(
           tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
       )
       work_tile = tile_sched.initial_work_tile_info()
       acc_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, num_acc_stage)


       c_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(epilog_warp_id))
       c_pipeline = pipeline.PipelineTmaStore.create(num_stages=num_c_stage, producer_group=c_producer_group)


       while work_tile.is_valid_tile:
           cur_tile_coord = work_tile.tile_idx
           mma_tile_coord_mnl = (
               cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),
               cur_tile_coord[1],
               cur_tile_coord[2],
           )
           bSG_gC = bSG_gC_partitioned[(None, None, None, *mma_tile_coord_mnl)]


           if cutlass.const_expr(overlapping_accum):
               acc_stage_index = acc_consumer_state.phase
               reverse_subtile = cutlass.Boolean(True) if acc_stage_index == 0 else cutlass.Boolean(False)
           else:
               acc_stage_index = acc_consumer_state.index


           tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_stage_index)]
           acc_pipeline.consumer_wait(acc_consumer_state)


           tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
           bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC))


           subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
           num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt
           for subtile_idx in cutlass.range(subtile_cnt):
               real_subtile_idx = subtile_idx
               if cutlass.const_expr(overlapping_accum):
                   if reverse_subtile:
                       real_subtile_idx = cta_tile_shape_mnk[1] // epi_tile_n - 1 - subtile_idx


               tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)]
               cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc)


               if cutlass.const_expr(overlapping_accum):
                   if subtile_idx == iter_acc_early_release_in_epilogue:
                       cute.arch.fence_view_async_tmem_load()
                       with cute.arch.elect_one():
                           acc_pipeline.consumer_release(acc_consumer_state)
                       acc_consumer_state.advance()


               acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load()
               acc_vec = acc_vec.to(c_dtype)
               tRS_rC.store(acc_vec)


               c_buffer = (num_prev_subtiles + real_subtile_idx) % num_c_stage
               cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)])
               cute.arch.fence_proxy(cute.arch.ProxyKind.async_shared, space=cute.arch.SharedSpace.shared_cta)
               epilog_sync_barrier.arrive_and_wait()


               if warp_idx == epilog_warp_id[0]:
                   cute.copy(tma_atom_c, bSG_sC[(None, c_buffer)], bSG_gC[(None, real_subtile_idx)])
                   c_pipeline.producer_commit()
                   c_pipeline.producer_acquire()
               epilog_sync_barrier.arrive_and_wait()


           if cutlass.const_expr(not overlapping_accum):
               with cute.arch.elect_one():
                   acc_pipeline.consumer_release(acc_consumer_state)
               acc_consumer_state.advance()


           tile_sched.advance_to_next_work()
           work_tile = tile_sched.get_current_work()


       tmem.relinquish_alloc_permit()
       epilog_sync_barrier.arrive_and_wait()
       tmem.free(acc_tmem_ptr)
       c_pipeline.producer_tail()


   return




@cute.jit
def my_kernel(
   a_ptr: cute.Pointer,
   b_ptr: cute.Pointer,
   sfa_ptr: cute.Pointer,
   sfb_ptr: cute.Pointer,
   c_ptr: cute.Pointer,
   problem_size: tuple,
   max_active_clusters: cutlass.Constexpr,
):
   """Host-side JIT function to prepare tensors and launch GPU kernel."""
   m, n, k, l = problem_size


   # Create tensors from pointers
   a_tensor = cute.make_tensor(
       a_ptr,
       cute.make_layout(
           (m, cute.assume(k, 32), l),
           stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32)),
       ),
   )
   b_tensor = cute.make_tensor(
       b_ptr,
       cute.make_layout(
           (n, cute.assume(k, 32), l),
           stride=(cute.assume(k, 32), 1, cute.assume(n * k, 32)),
       ),
   )
   c_tensor = cute.make_tensor(
       c_ptr, cute.make_layout((cute.assume(m, 32), n, l), stride=(n, 1, m * n))
   )


   # Setup sfa/sfb tensor
   sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(a_tensor.shape, sf_vec_size)
   sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)
   sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(b_tensor.shape, sf_vec_size)
   sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout)


   # Compute MMA configuration
   mma_inst_shape_mn = mma_tiler_mn
   mma_inst_shape_mn_sfb = (
       mma_inst_shape_mn[0] // (2 if use_2cta_instrs else 1),
       cute.round_up(mma_inst_shape_mn[1], 128),
   )


   tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
       ab_dtype,
       utils.LayoutEnum.from_tensor(a_tensor).mma_major_mode(),
       utils.LayoutEnum.from_tensor(b_tensor).mma_major_mode(),
       sf_dtype,
       sf_vec_size,
       cta_group,
       mma_inst_shape_mn,
   )
   tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
       ab_dtype,
       utils.LayoutEnum.from_tensor(a_tensor).mma_major_mode(),
       utils.LayoutEnum.from_tensor(b_tensor).mma_major_mode(),
       sf_dtype,
       sf_vec_size,
       tcgen05.CtaGroup.ONE,
       mma_inst_shape_mn_sfb,
   )


   mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
   mma_inst_tile_k = 4
   mma_tiler = (mma_inst_shape_mn[0], mma_inst_shape_mn[1], mma_inst_shape_k * mma_inst_tile_k)
   mma_tiler_sfb = (mma_inst_shape_mn_sfb[0], mma_inst_shape_mn_sfb[1], mma_inst_shape_k * mma_inst_tile_k)


   cta_tile_shape_mnk = (
       mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
       mma_tiler[1],
       mma_tiler[2],
   )
   cta_tile_shape_mnk_sfb = (
       mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape),
       mma_tiler_sfb[1],
       mma_tiler_sfb[2],
   )


   cluster_layout_vmnk = cute.tiled_divide(
       cute.make_layout((*cluster_shape_mn, 1)),
       (tiled_mma.thr_id.shape,),
   )
   cluster_layout_sfb_vmnk = cute.tiled_divide(
       cute.make_layout((*cluster_shape_mn, 1)),
       (tiled_mma_sfb.thr_id.shape,),
   )


   num_mcast_ctas_a = cute.size(cluster_layout_vmnk.shape[2])
   num_mcast_ctas_b = cute.size(cluster_layout_vmnk.shape[1])
   is_a_mcast = num_mcast_ctas_a > 1
   is_b_mcast = num_mcast_ctas_b > 1


   c_layout = utils.LayoutEnum.from_tensor(c_tensor)
   epi_tile = sm100_utils.compute_epilogue_tile_shape(
       cta_tile_shape_mnk, use_2cta_instrs, c_layout, c_dtype,
   )
   epi_tile_n = cute.size(epi_tile[1])


   # Compute stages
   num_acc_stage = 1 if mma_tiler[1] == 256 else 2
   num_c_stage = 2
   a_smem_layout_stage_one = sm100_utils.make_smem_layout_a(tiled_mma, mma_tiler, ab_dtype, 1)
   b_smem_layout_staged_one = sm100_utils.make_smem_layout_b(tiled_mma, mma_tiler, ab_dtype, 1)
   sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa(tiled_mma, mma_tiler, sf_vec_size, 1)
   sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb(tiled_mma, mma_tiler, sf_vec_size, 1)
   c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi(c_dtype, c_layout, epi_tile, 1)
   ab_bytes_per_stage = (
       cute.size_in_bytes(ab_dtype, a_smem_layout_stage_one) +
       cute.size_in_bytes(ab_dtype, b_smem_layout_staged_one) +
       cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one) +
       cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one)
   )
   mbar_helpers_bytes = 1024
   c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one)
   c_bytes = c_bytes_per_stage * num_c_stage
   num_ab_stage = (smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes)) // ab_bytes_per_stage
   num_c_stage += (
       smem_capacity - occupancy * ab_bytes_per_stage * num_ab_stage - occupancy * (mbar_helpers_bytes + c_bytes)
   ) // (occupancy * c_bytes_per_stage)


   # Compute SMEM layouts
   a_smem_layout_staged = sm100_utils.make_smem_layout_a(tiled_mma, mma_tiler, ab_dtype, num_ab_stage)
   b_smem_layout_staged = sm100_utils.make_smem_layout_b(tiled_mma, mma_tiler, ab_dtype, num_ab_stage)
   sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(tiled_mma, mma_tiler, sf_vec_size, num_ab_stage)
   sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(tiled_mma, mma_tiler, sf_vec_size, num_ab_stage)
   c_smem_layout_staged = sm100_utils.make_smem_layout_epi(c_dtype, c_layout, epi_tile, num_c_stage)


   overlapping_accum = num_acc_stage == 1
   sf_atom_mn = 32
   num_sfa_tmem_cols = (cta_tile_shape_mnk[0] // sf_atom_mn) * 4
   num_sfb_tmem_cols = (cta_tile_shape_mnk_sfb[1] // sf_atom_mn) * 4
   num_sf_tmem_cols = num_sfa_tmem_cols + num_sfb_tmem_cols
   num_accumulator_tmem_cols = (
       cta_tile_shape_mnk[1] * num_acc_stage
       if not overlapping_accum
       else cta_tile_shape_mnk[1] * 2 - num_sf_tmem_cols
   )
   iter_acc_early_release_in_epilogue = num_sf_tmem_cols // epi_tile_n


   atom_thr_size = cute.size(tiled_mma.thr_id.shape)


   # Setup TMA for A
   a_op = sm100_utils.cluster_shape_to_tma_atom_A(cluster_shape_mn, tiled_mma.thr_id)
   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(
       a_op, a_tensor, a_smem_layout, mma_tiler, tiled_mma, cluster_layout_vmnk.shape,
   )


   # Setup TMA for B
   b_op = sm100_utils.cluster_shape_to_tma_atom_B(cluster_shape_mn, tiled_mma.thr_id)
   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(
       b_op, b_tensor, b_smem_layout, mma_tiler, tiled_mma, cluster_layout_vmnk.shape,
   )


   # Setup TMA for SFA
   sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(cluster_shape_mn, tiled_mma.thr_id)
   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(
       sfa_op, sfa_tensor, sfa_smem_layout, mma_tiler, tiled_mma,
       cluster_layout_vmnk.shape, internal_type=cutlass.Int16,
   )


   # Setup TMA for SFB
   sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(cluster_shape_mn, tiled_mma.thr_id)
   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(
       sfb_op, sfb_tensor, sfb_smem_layout, mma_tiler_sfb, tiled_mma_sfb,
       cluster_layout_sfb_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


   # Setup TMA store for C
   epi_smem_layout = cute.slice_(c_smem_layout_staged, (None, None, 0))
   tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom(
       cpasync.CopyBulkTensorTileS2GOp(), c_tensor, epi_smem_layout, epi_tile,
   )


   # Compute grid
   c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0))
   gc = cute.zipped_divide(c_tensor, tiler=c_shape)
   num_ctas_mnl = gc[(0, (None, None, None))].shape
   cluster_shape_mnl = (*cluster_shape_mn, 1)
   tile_sched_params = utils.PersistentTileSchedulerParams(num_ctas_mnl, cluster_shape_mnl)
   grid = utils.StaticPersistentTileScheduler.get_grid_shape(tile_sched_params, max_active_clusters)


   buffer_align_bytes = 1024


   # Define shared storage
   @cute.struct
   class SharedStorage:
       ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, num_ab_stage]
       ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, num_ab_stage]
       acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, num_acc_stage]
       acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, num_acc_stage]
       tmem_dealloc_mbar_ptr: cutlass.Int64
       tmem_holding_buf: cutlass.Int32
       sC: cute.struct.Align[
           cute.struct.MemRange[c_dtype, cute.cosize(c_smem_layout_staged.outer)],
           buffer_align_bytes,
       ]
       sA: cute.struct.Align[
           cute.struct.MemRange[ab_dtype, cute.cosize(a_smem_layout_staged.outer)],
           buffer_align_bytes,
       ]
       sB: cute.struct.Align[
           cute.struct.MemRange[ab_dtype, cute.cosize(b_smem_layout_staged.outer)],
           buffer_align_bytes,
       ]
       sSFA: cute.struct.Align[
           cute.struct.MemRange[sf_dtype, cute.cosize(sfa_smem_layout_staged)],
           buffer_align_bytes,
       ]
       sSFB: cute.struct.Align[
           cute.struct.MemRange[sf_dtype, cute.cosize(sfb_smem_layout_staged)],
           buffer_align_bytes,
       ]


   # Launch kernel
   kernel(
       tiled_mma, tiled_mma_sfb,
       tma_atom_a, tma_tensor_a,
       tma_atom_b, tma_tensor_b,
       tma_atom_sfa, tma_tensor_sfa,
       tma_atom_sfb, tma_tensor_sfb,
       tma_atom_c, tma_tensor_c,
       cluster_layout_vmnk,
       cluster_layout_sfb_vmnk,
       a_smem_layout_staged,
       b_smem_layout_staged,
       sfa_smem_layout_staged,
       sfb_smem_layout_staged,
       c_smem_layout_staged,
       epi_tile,
       tile_sched_params,
       mma_tiler,
       mma_tiler_sfb,
       cta_tile_shape_mnk,
       num_ab_stage,
       num_acc_stage,
       num_c_stage,
       num_tma_load_bytes,
       num_mcast_ctas_a,
       num_mcast_ctas_b,
       is_a_mcast,
       is_b_mcast,
       overlapping_accum,
       num_accumulator_tmem_cols,
       num_sfa_tmem_cols,
       num_sf_tmem_cols,
       epi_tile_n,
       iter_acc_early_release_in_epilogue,
       c_layout,
       SharedStorage,
   ).launch(
       grid=grid,
       block=[threads_per_cta, 1, 1],
       cluster=(*cluster_shape_mn, 1),
       min_blocks_per_mp=1,
   )
   return




# Global cache for compiled kernel
_compiled_kernel_cache = None




def compile_kernel():
   """Compile the kernel once and cache it."""
   global _compiled_kernel_cache
   if _compiled_kernel_cache is not None:
       return _compiled_kernel_cache


   # Compute max active clusters for optimal persistent scheduling
   hardware_info = cutlass.utils.HardwareInfo()
   max_active_clusters = hardware_info.get_max_active_clusters(
       cluster_shape_mn[0] * cluster_shape_mn[1]
   )


   # Create CuTe pointers for A/B/C/SFA/SFB via make_ptr with address 0
   a_ptr = make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
   b_ptr = make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
   c_ptr = make_ptr(c_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
   sfa_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32)
   sfb_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32)


   # Compile the kernel with optimization
   _compiled_kernel_cache = cute.compile(
       my_kernel, a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (0, 0, 0, 0),
       max_active_clusters,
       options="--opt-level 2"
   )
   return _compiled_kernel_cache




def custom_kernel(data: input_t) -> output_t:
   """Execute the block-scaled GEMM kernel."""
   a, b, _, _, sfa_permuted, sfb_permuted, c = data


   # Ensure kernel is compiled
   compiled_func = compile_kernel()


   # Get dimensions from MxKxL layout
   m, k, l = a.shape
   n, _, _ = b.shape
   # Torch use e2m1_x2 data type, thus k is halved
   k = k * 2


   # Create CuTe pointers
   a_ptr = make_ptr(ab_dtype, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
   b_ptr = make_ptr(ab_dtype, b.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
   c_ptr = make_ptr(c_dtype, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
   sfa_ptr = make_ptr(sf_dtype, sfa_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32)
   sfb_ptr = make_ptr(sf_dtype, sfb_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32)


   # Execute the compiled kernel
   compiled_func(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l))


   return c