submission 404186

import cutlass
import cutlass.cute as cute
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.cute.nvgpu import cpasync, tcgen05
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
from cutlass.cutlass_dsl import CuTeDSL, dsl_user_op, T

import functools
from typing import Tuple, List, Union, Type

import torch
from task import input_t, output_t


class GroupGemm:
    """
    Block-scaled Group GEMM kernel for NVIDIA Blackwell (B200) GPU.
    
    This class encapsulates both host-side JIT compilation (@cute.jit __call__)
    and device-side kernel logic (@cute.kernel kernel) following the pattern
    used in optimized CUTLASS kernels.
    
    Architecture:
    - __init__: Store configuration (tile sizes, stages, warp roles)
    - _setup_attributes: Compute derived attributes from problem size
    - __call__ (@cute.jit): Host-side setup - create tensors, TMA atoms, launch kernel
    - kernel (@cute.kernel): Device-side warp-specialized mainloop and epilogue
    """
    
    def __init__(
        self,
        num_groups: int,
        mma_tiler_mnk: Tuple[int, int, int] = (128, 128, 256),
    ):
        """
        Initialize kernel configuration.
        
        Args:
            num_groups: Number of groups in this batch
            mma_tiler_mnk: Tile sizes for M, N, K dimensions
        """
        self.num_groups = num_groups
        self.mma_tiler_mnk = mma_tiler_mnk
        
        # Data types
        self.ab_dtype = cutlass.Float4E2M1FN
        self.sf_dtype = cutlass.Float8E4M3FN
        self.c_dtype = cutlass.Float16
        self.acc_dtype = cutlass.Float32
        self.sf_vec_size = 16
        
        # TMA descriptor configuration
        self.bytes_per_tensormap = 128
        self.num_tensormaps = 5  # a, b, sfa, sfb, c
        
        # MMA instruction shape
        self.mma_inst_shape_k = 64
        
        # Warp roles for warp-specialized mainloop/epilogue
        self.epilog_warp_id = (0, 1, 2, 3)
        self.mma_warp_id = 4
        self.tma_warp_id = 5
        self.threads_per_cta = 32 * 6  # 6 warps
        
        # Pipeline stages
        self.num_acc_stage = 1
        
        # SMEM capacity for B200 (SM100) - 228KB
        self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
        
        # Target occupancy (CTAs per SM)
        self.occupancy = 1
        
        # Total TMem columns
        self.num_tmem_alloc_cols = 512
        
        # Cluster shape (1x1 for group GEMM)
        self.cluster_shape_mn = (1, 1)
        
        # Barriers
        self.tmem_alloc_barrier = pipeline.NamedBarrier(
            barrier_id=1,
            num_threads=32 * 5,  # MMA + epilogue warps
        )
        self.epilog_sync_barrier = pipeline.NamedBarrier(
            barrier_id=2,
            num_threads=32 * 4,  # Epilogue warps
        )

    def _setup_attributes(self):
        """Compute derived attributes from configuration."""
        # Create trivial tiled MMA for attribute computation
        mma_op = tcgen05.MmaMXF4NVF4Op(
            self.sf_dtype,
            (self.mma_tiler_mnk[0], self.mma_tiler_mnk[1], self.mma_inst_shape_k),
            tcgen05.CtaGroup.ONE,
            tcgen05.OperandSource.SMEM,
        )
        tiled_mma = cute.make_tiled_mma(mma_op)
        
        # Cluster layouts
        self.cluster_layout_vmnk = cute.tiled_divide(
            cute.make_layout((1, 1, 1)),
            (tiled_mma.thr_id.shape,),
        )
        
        # SFB uses different tiled MMA for N dimension rounding
        mma_inst_shape_mn_sfb = (
            self.mma_tiler_mnk[0],
            cute.round_up(self.mma_tiler_mnk[1], 128),
        )
        tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
            self.ab_dtype,
            tcgen05.OperandMajorMode.K,
            tcgen05.OperandMajorMode.K,
            self.sf_dtype,
            self.sf_vec_size,
            tcgen05.CtaGroup.ONE,
            mma_inst_shape_mn_sfb,
        )
        self.cluster_layout_sfb_vmnk = cute.tiled_divide(
            cute.make_layout((1, 1, 1)),
            (tiled_mma_sfb.thr_id.shape,),
        )
        
        # CTA tile shape
        self.cta_tile_shape_mnk = (
            self.mma_tiler_mnk[0] // cute.size(tiled_mma.thr_id.shape),
            self.mma_tiler_mnk[1],
            self.mma_tiler_mnk[2],
        )
        
        # Epilogue tile
        self.c_layout = utils.LayoutEnum.ROW_MAJOR
        self.epi_tile = sm100_utils.compute_epilogue_tile_shape(
            self.cta_tile_shape_mnk,
            False,  # use_2cta_instrs
            self.c_layout,
            self.c_dtype,
        )
        
        # Compute optimal pipeline stages
        self.num_ab_stage = self._compute_num_ab_stages(tiled_mma)
        
        # SMEM layouts
        self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
            tiled_mma, self.mma_tiler_mnk, self.ab_dtype, self.num_ab_stage
        )
        self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
            tiled_mma, self.mma_tiler_mnk, self.ab_dtype, self.num_ab_stage
        )
        self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
            tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, self.num_ab_stage
        )
        self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(
            tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, self.num_ab_stage
        )
        self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi(
            self.c_dtype, self.c_layout, self.epi_tile, 1
        )
        
        return tiled_mma, tiled_mma_sfb

    def _compute_num_ab_stages(self, tiled_mma: cute.TiledMma) -> int:
        """Compute optimal number of A/B pipeline stages based on SMEM capacity."""
        # Compute SMEM size for single stage of each buffer
        a_smem_layout_one = sm100_utils.make_smem_layout_a(
            tiled_mma, self.mma_tiler_mnk, self.ab_dtype, 1
        )
        b_smem_layout_one = sm100_utils.make_smem_layout_b(
            tiled_mma, self.mma_tiler_mnk, self.ab_dtype, 1
        )
        sfa_smem_layout_one = blockscaled_utils.make_smem_layout_sfa(
            tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1
        )
        sfb_smem_layout_one = blockscaled_utils.make_smem_layout_sfb(
            tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1
        )
        c_smem_layout_one = sm100_utils.make_smem_layout_epi(
            self.c_dtype, self.c_layout, self.epi_tile, 1
        )
        
        # Bytes per stage for A, B, SFA, SFB
        ab_bytes_per_stage = (
            cute.size_in_bytes(self.ab_dtype, a_smem_layout_one)
            + cute.size_in_bytes(self.ab_dtype, b_smem_layout_one)
            + cute.size_in_bytes(self.sf_dtype, sfa_smem_layout_one)
            + cute.size_in_bytes(self.sf_dtype, sfb_smem_layout_one)
        )
        
        # Fixed overhead
        mbar_helpers_bytes = 2048
        c_bytes = cute.size_in_bytes(self.c_dtype, c_smem_layout_one)
        
        # Compute max AB stages that fit
        available_for_ab = self.smem_capacity // self.occupancy - mbar_helpers_bytes - c_bytes
        num_ab_stage = max(1, available_for_ab // ab_bytes_per_stage)
        
        # Cap at reasonable maximum
        return min(num_ab_stage, 8)

    @cute.jit
    def __call__(
        self,
        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]],
        num_groups: cutlass.Int32,
    ):
        """
        Host-side JIT entry point. Creates tensors, TMA atoms, and launches kernel.
        """
        # Setup derived attributes
        tiled_mma, tiled_mma_sfb = self._setup_attributes()
        
        # Create metadata tensors from pointers
        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, self.num_tensormaps, 16),
                stride=(self.num_tensormaps * 16, 16, 1)
            ),
        )
        
        # Create template tensors with max shapes for TMA descriptor setup
        max_m = cutlass.Int32(512)
        max_n = cutlass.Int32(7168)
        max_k = cutlass.Int32(7168)
        
        initial_a = cute.make_tensor(
            cute.make_ptr(self.ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16),
            cute.make_layout(
                (max_m, cute.assume(max_k, 32), cutlass.Int32(1)),
                stride=(cute.assume(max_k, 32), 1, cute.assume(max_m * max_k, 32)),
            ),
        )
        initial_b = cute.make_tensor(
            cute.make_ptr(self.ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16),
            cute.make_layout(
                (max_n, cute.assume(max_k, 32), cutlass.Int32(1)),
                stride=(cute.assume(max_k, 32), 1, cute.assume(max_n * max_k, 32)),
            ),
        )
        initial_c = cute.make_tensor(
            cute.make_ptr(self.c_dtype, 0, cute.AddressSpace.gmem, assumed_align=16),
            cute.make_layout(
                (max_m, max_n, cutlass.Int32(1)),
                stride=(cute.assume(max_n, 32), 1, cute.assume(max_m * max_n, 32)),
            ),
        )
        
        # Scale factor layouts
        sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(
            initial_a.shape, self.sf_vec_size
        )
        sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(
            initial_b.shape, self.sf_vec_size
        )
        initial_sfa = cute.make_tensor(
            cute.make_ptr(self.sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=16),
            sfa_layout
        )
        initial_sfb = cute.make_tensor(
            cute.make_ptr(self.sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=16),
            sfb_layout
        )
        
        atom_thr_size = cute.size(tiled_mma.thr_id.shape)
        
        # Setup TMA for A
        a_smem_layout = cute.slice_(self.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,
            self.mma_tiler_mnk,
            tiled_mma,
            self.cluster_layout_vmnk.shape,
        )
        
        # Setup TMA for B
        b_smem_layout = cute.slice_(self.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,
            self.mma_tiler_mnk,
            tiled_mma,
            self.cluster_layout_vmnk.shape,
        )
        
        # Setup TMA for SFA
        sfa_smem_layout = cute.slice_(self.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,
            self.mma_tiler_mnk,
            tiled_mma,
            self.cluster_layout_vmnk.shape,
            internal_type=cutlass.Int16,
        )
        
        # Setup TMA for SFB
        mma_inst_shape_mn_sfb = (
            self.mma_tiler_mnk[0],
            cute.round_up(self.mma_tiler_mnk[1], 128),
        )
        tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
            self.ab_dtype,
            tcgen05.OperandMajorMode.K,
            tcgen05.OperandMajorMode.K,
            self.sf_dtype,
            self.sf_vec_size,
            tcgen05.CtaGroup.ONE,
            mma_inst_shape_mn_sfb,
        )
        sfb_smem_layout = cute.slice_(self.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_inst_shape_mn_sfb[0], mma_inst_shape_mn_sfb[1], self.mma_tiler_mnk[2]),
            tiled_mma_sfb,
            self.cluster_layout_sfb_vmnk.shape,
            internal_type=cutlass.Int16,
        )
        
        # Compute TMA load bytes
        a_copy_size = cute.size_in_bytes(self.ab_dtype, a_smem_layout)
        b_copy_size = cute.size_in_bytes(self.ab_dtype, b_smem_layout)
        sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout)
        sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout)
        num_tma_load_bytes = (
            a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size
        ) * atom_thr_size
        
        # Setup TMA for C (store)
        epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0))
        tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom(
            cpasync.CopyBulkTensorTileS2GOp(),
            initial_c,
            epi_smem_layout,
            self.epi_tile,
        )
        
        # Store CTA shape information for each group
        cta_mn_list = []
        for group_idx, (m, n, k, l) in enumerate(problem_sizes):
            x, y = cute.ceil_div(problem_sizes[group_idx][:2], self.mma_tiler_mnk[0:2])
            cta_mn_list.append((x, y))
        
        # Compute grid size
        grid = (1, 1, total_num_clusters)
        
        # Launch the kernel
        self.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,
            tensor_of_abc_ptrs,
            tensor_of_sfasfb_ptrs,
            tensor_of_tensormap,
            tensor_of_problem_sizes,
            self.cluster_layout_vmnk,
            self.cluster_layout_sfb_vmnk,
            self.a_smem_layout_staged,
            self.b_smem_layout_staged,
            self.sfa_smem_layout_staged,
            self.sfb_smem_layout_staged,
            self.c_smem_layout_staged,
            self.epi_tile,
            cta_mn_list,
            num_tma_load_bytes,
            self.num_ab_stage,
        ).launch(
            grid=grid,
            block=[self.threads_per_cta, 1, 1],
            cluster=(1, 1, 1),
        )

    @cute.kernel
    def kernel(
        self,
        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,
        tensor_of_abc_ptrs: cute.Tensor,
        tensor_of_sfasfb_ptrs: cute.Tensor,
        tensormaps: cute.Tensor,
        tensor_of_problem_sizes: 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,
        cta_mn_list: List[Tuple[int, int]],
        num_tma_load_bytes: cutlass.Constexpr[int],
        num_ab_stage: cutlass.Constexpr[int],
    ):
        """
        Device-side kernel performing the Group GEMM computation.
        
        Warp specialization:
        - Warps 0-3: Epilogue (TMem → Registers → SMEM → GMEM via TMA)
        - Warp 4: MMA computation
        - Warp 5: TMA loads
        """
        warp_idx = cute.arch.warp_idx()
        warp_idx = cute.arch.make_warp_uniform(warp_idx)
        tidx, _, _ = cute.arch.thread_idx()

        # 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

        # Construct C Tensor for each CTA
        mC_mnl_iter = cute.make_ptr(
            self.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),)
        )
        real_mC_mnl = cute.make_tensor(mC_mnl_iter, mC_mnl_layout)
        
        # Use template tensor for partitioning
        gC_mnl = cute.local_tile(
            mC_mnl, cute.slice_(self.mma_tiler_mnk, (None, None, 0)), (None, None, None)
        )
        mma_tile_coord_mnl = (coord_x, coord_y, 0)

        # Define shared storage
        size_tensormap_in_i64 = (
            self.num_tensormaps * self.bytes_per_tensormap // 8
        )
        
        @cute.struct
        class SharedStorage:
            tensormap_buffer: cute.struct.MemRange[
                cutlass.Int64, size_tensormap_in_i64
            ]
            ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, num_ab_stage * 2]
            acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
            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 + self.bytes_per_tensormap // 8
        tensormap_sfa_smem_ptr = tensormap_b_smem_ptr + self.bytes_per_tensormap // 8
        tensormap_sfb_smem_ptr = tensormap_sfa_smem_ptr + self.bytes_per_tensormap // 8
        tensormap_c_smem_ptr = tensormap_sfb_smem_ptr + self.bytes_per_tensormap // 8

        sA = smem.allocate_tensor(
            element_type=self.ab_dtype,
            layout=a_smem_layout_staged.outer,
            byte_alignment=128,
            swizzle=a_smem_layout_staged.inner,
        )
        sB = smem.allocate_tensor(
            element_type=self.ab_dtype,
            layout=b_smem_layout_staged.outer,
            byte_alignment=128,
            swizzle=b_smem_layout_staged.inner,
        )
        sSFA = smem.allocate_tensor(
            element_type=self.sf_dtype,
            layout=sfa_smem_layout_staged,
            byte_alignment=128,
        )
        sSFB = smem.allocate_tensor(
            element_type=self.sf_dtype,
            layout=sfb_smem_layout_staged,
            byte_alignment=128,
        )
        sC = smem.allocate_tensor(
            element_type=self.c_dtype,
            layout=c_smem_layout_staged.outer,
            byte_alignment=128,
            swizzle=c_smem_layout_staged.inner,
        )

        # Update TMA descriptors with 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
        )
        tensormap_c_gmem_ptr = tensormap_manager.get_tensormap_ptr(
            tensormaps[(bidz, 4, None)].iterator
        )

        # Construct real tensors from runtime pointers
        mA_mkl_iter = cute.make_ptr(
            self.ab_dtype, tensor_of_abc_ptrs[group_idx, 0], cute.AddressSpace.gmem
        ).align(32)
        mB_nkl_iter = cute.make_ptr(
            self.ab_dtype, tensor_of_abc_ptrs[group_idx, 1], cute.AddressSpace.gmem
        ).align(32)
        sfa_mkl_iter = cute.make_ptr(
            self.sf_dtype, tensor_of_sfasfb_ptrs[group_idx, 0], cute.AddressSpace.gmem
        ).align(32)
        sfb_nkl_iter = cute.make_ptr(
            self.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),)
        )

        atom_shape = ((32, 4), (self.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)
        real_tensor_c = real_mC_mnl

        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.init_tensormap_from_atom(
                tma_atom_c, tensormap_c_smem_ptr, 0
            )
            tensormap_manager.update_tensormap(
                (real_tensor_a, real_tensor_b, real_tensor_sfa, real_tensor_sfb, real_tensor_c),
                (tma_atom_a, tma_atom_b, tma_atom_sfa, tma_atom_sfb, tma_atom_c),
                (tensormap_a_gmem_ptr, tensormap_b_gmem_ptr, tensormap_sfa_gmem_ptr, tensormap_sfb_gmem_ptr, tensormap_c_gmem_ptr),
                0,
                (tensormap_a_smem_ptr, tensormap_b_smem_ptr, tensormap_sfa_smem_ptr, tensormap_sfb_smem_ptr, tensormap_c_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)
            tensormap_manager.fence_tensormap_update(tensormap_c_gmem_ptr)

        cute.arch.barrier()

        if warp_idx == self.tma_warp_id:
            cpasync.prefetch_descriptor(tma_atom_a)
            cpasync.prefetch_descriptor(tma_atom_b)
            cpasync.prefetch_descriptor(tma_atom_sfa)
            cpasync.prefetch_descriptor(tma_atom_sfb)
            cpasync.prefetch_descriptor(tma_atom_c)

        mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape)
        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
        )

        ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
        num_tma_producer = 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)
        acc_pipeline_consumer_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread, len(self.epilog_warp_id)
        )
        acc_pipeline = pipeline.PipelineUmmaAsync.create(
            barrier_storage=storage.acc_full_mbar_ptr.data_ptr(),
            num_stages=self.num_acc_stage,
            producer_group=acc_pipeline_producer_group,
            consumer_group=acc_pipeline_consumer_group,
            cta_layout_vmnk=cluster_layout_vmnk,
        )

        tmem = utils.TmemAllocator(
            storage.tmem_holding_buf,
            barrier_for_retrieve=self.tmem_alloc_barrier,
            allocator_warp_id=self.epilog_warp_id[0],
        )

        pipeline_init_arrive(cluster_shape_mn=(1, 1), is_relaxed=True)

        # Tile real_tensor_a to compute k_block_cnt
        gA_mkl_real = cute.local_tile(
            real_tensor_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None)
        )

        # Tile template tensors for TMA partitioning
        gA_mkl = cute.local_tile(
            mA_mkl, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None)
        )
        gB_nkl = cute.local_tile(
            mB_nkl, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None)
        )
        gSFA_mkl = cute.local_tile(
            mSFA_mkl, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None)
        )
        gSFB_nkl = cute.local_tile(
            mSFB_nkl,
            cute.slice_((self.mma_tiler_mnk[0], cute.round_up(self.mma_tiler_mnk[1], 128), self.mma_tiler_mnk[2]), (0, None, None)),
            (None, None, None),
        )

        k_block_cnt = cute.size(gA_mkl_real, 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)

        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(self.mma_tiler_mnk[:2])
        tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape)

        pipeline_init_wait(cluster_shape_mn=(1, 1))

        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(self.mma_tiler_mnk[1] == 64):
            slice_n = mma_tile_coord_mnl[1] // 2
        tBgSFB_slice = tBgSFB[(None, slice_n, None, mma_tile_coord_mnl[2])]

        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
        )

        # TMA warp: Load tiles from global memory
        if warp_idx == self.tma_warp_id:
            ab_producer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer, num_ab_stage
            )
            tma_desc_a = tensormap_manager.get_tensormap_ptr(
                tensormap_a_gmem_ptr, cute.AddressSpace.generic
            )
            tma_desc_b = tensormap_manager.get_tensormap_ptr(
                tensormap_b_gmem_ptr, cute.AddressSpace.generic
            )
            tma_desc_sfa = tensormap_manager.get_tensormap_ptr(
                tensormap_sfa_gmem_ptr, cute.AddressSpace.generic
            )
            tma_desc_sfb = tensormap_manager.get_tensormap_ptr(
                tensormap_sfb_gmem_ptr, cute.AddressSpace.generic
            )
            for k_block_idx in cutlass.range(0, k_block_cnt, 1, unroll=1):
                ab_pipeline.producer_acquire(ab_producer_state)
                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),
                    tma_desc_ptr=tma_desc_a,
                    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),
                    tma_desc_ptr=tma_desc_b,
                    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),
                    tma_desc_ptr=tma_desc_sfa,
                    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),
                    tma_desc_ptr=tma_desc_sfb,
                    mcast_mask=sfb_full_mcast_mask,
                )
                ab_producer_state.advance()

        # MMA warp: Compute matrix multiplication
        elif warp_idx == self.mma_warp_id:
            tmem.wait_for_alloc()
            acc_tmem_ptr = tmem.retrieve_ptr(cutlass.Float32)
            tCtAcc = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
            sfa_tmem_ptr = cute.recast_ptr(
                acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc),
                dtype=self.sf_dtype,
            )
            tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
                tiled_mma,
                self.mma_tiler_mnk,
                self.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=self.sf_dtype,
            )
            tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
                tiled_mma,
                self.mma_tiler_mnk,
                self.sf_vec_size,
                cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)),
            )
            tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout)

            tiled_copy_s2t_sfa, tCsSFA_compact_s2t, tCtSFA_compact_s2t = (
                self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA)
            )
            tiled_copy_s2t_sfb, tCsSFB_compact_s2t, tCtSFB_compact_s2t = (
                self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB)
            )

            ab_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, num_ab_stage
            )
            acc_producer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer, self.num_acc_stage
            )

            tCtSFB_mma = tCtSFB
            if cutlass.const_expr(self.mma_tiler_mnk[1] == 64):
                offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2)
                shifted_ptr = cute.recast_ptr(
                    acc_tmem_ptr
                    + tcgen05.find_tmem_tensor_col_offset(tCtAcc)
                    + tcgen05.find_tmem_tensor_col_offset(tCtSFA)
                    + offset,
                    dtype=self.sf_dtype,
                )
                tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout)

            tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
            for _ in range(k_block_cnt):
                ab_pipeline.consumer_wait(ab_consumer_state)
                s2t_stage_coord = (None, None, None, None, ab_consumer_state.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)
                num_kphases = cute.size(tCrA, mode=[2])
                for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
                    kphase_coord = (None, None, kphase_idx, ab_consumer_state.index)
                    sf_kphase_coord = (None, None, kphase_idx)
                    tiled_mma.set(tcgen05.Field.SFA, tCtSFA[sf_kphase_coord].iterator)
                    tiled_mma.set(tcgen05.Field.SFB, tCtSFB_mma[sf_kphase_coord].iterator)
                    cute.gemm(tiled_mma, tCtAcc, tCrA[kphase_coord], tCrB[kphase_coord], tCtAcc)
                    tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
                ab_pipeline.consumer_release(ab_consumer_state)
                ab_consumer_state.advance()
            acc_pipeline.producer_commit(acc_producer_state)

        # Epilogue warps: Store results via TMA
        elif warp_idx in self.epilog_warp_id:
            tmem.allocate(self.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)
            
            # Setup epilogue copies
            tiled_copy_t2r, tTR_tAcc, tTR_rAcc = self.epilog_tmem_copy_and_partition(
                tidx, tCtAcc, tCgC, epi_tile
            )
            
            tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype)
            tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition(
                tiled_copy_t2r, tTR_rC, tidx, sC
            )
            
            _, bSG_sC, bSG_gC_partitioned = self.epilog_gmem_copy_and_partition(
                tma_atom_c, tCgC, epi_tile, sC
            )
            
            # TMA store pipeline
            c_producer_group = pipeline.CooperativeGroup(
                pipeline.Agent.Thread, 32 * len(self.epilog_warp_id),
            )
            c_pipeline = pipeline.PipelineTmaStore.create(
                num_stages=1, producer_group=c_producer_group,
            )
            
            if warp_idx == self.epilog_warp_id[0]:
                tensormap_manager.fence_tensormap_update(tensormap_c_gmem_ptr)
            
            acc_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.num_acc_stage
            )
            acc_pipeline.consumer_wait(acc_consumer_state)
            
            bSG_gC = bSG_gC_partitioned[(None, None, None, *mma_tile_coord_mnl)]
            tTR_tAcc_tile = tTR_tAcc[(None, None, None, None, None)]
            tTR_tAcc_grouped = cute.group_modes(tTR_tAcc_tile, 3, cute.rank(tTR_tAcc_tile))
            bSG_gC_grouped = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC))
            
            subtile_cnt = cute.size(tTR_tAcc_grouped.shape, mode=[3])
            for subtile_idx in range(subtile_cnt):
                tTR_tAcc_subtile = tTR_tAcc_grouped[(None, None, None, subtile_idx)]
                cute.copy(tiled_copy_t2r, tTR_tAcc_subtile, tTR_rAcc)
                
                acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load()
                tRS_rC.store(acc_vec.to(self.c_dtype))
                
                cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, 0)])
                
                cute.arch.fence_proxy(
                    cute.arch.ProxyKind.async_shared,
                    space=cute.arch.SharedSpace.shared_cta,
                )
                self.epilog_sync_barrier.arrive_and_wait()
                
                if warp_idx == self.epilog_warp_id[0]:
                    cute.copy(
                        tma_atom_c,
                        bSG_sC[(None, 0)],
                        bSG_gC_grouped[(None, subtile_idx)],
                        tma_desc_ptr=tensormap_manager.get_tensormap_ptr(
                            tensormap_c_gmem_ptr, cute.AddressSpace.generic,
                        ),
                    )
                    c_pipeline.producer_commit()
                    c_pipeline.producer_acquire()
                self.epilog_sync_barrier.arrive_and_wait()
            
            with cute.arch.elect_one():
                acc_pipeline.consumer_release(acc_consumer_state)
            
            tmem.relinquish_alloc_permit()
            tmem.free(acc_tmem_ptr)
            c_pipeline.producer_tail()

    def mainloop_s2t_copy_and_partition(
        self,
        sSF: cute.Tensor,
        tSF: cute.Tensor,
    ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
        """Setup SMEM → TMem copy for scale factors."""
        tCsSF_compact = cute.filter_zeros(sSF)
        tCtSF_compact = cute.filter_zeros(tSF)
        copy_atom_s2t = cute.make_copy_atom(
            tcgen05.Cp4x32x128bOp(tcgen05.CtaGroup.ONE),
            self.sf_dtype,
        )
        tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact)
        thr_copy_s2t = tiled_copy_s2t.get_slice(0)
        tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact)
        tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
            tiled_copy_s2t, tCsSF_compact_s2t_
        )
        tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact)
        return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t

    def epilog_tmem_copy_and_partition(
        self,
        tidx: cutlass.Int32,
        tAcc: cute.Tensor,
        tCgC: cute.Tensor,
        epi_tile: cute.Tile,
    ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
        """Setup TMem → Register copy for epilogue."""
        copy_atom_t2r = sm100_utils.get_tmem_load_op(
            self.cta_tile_shape_mnk,
            self.c_layout,
            self.c_dtype,
            self.acc_dtype,
            epi_tile,
            False,  # use_2cta_instrs
        )
        tAcc_epi = cute.flat_divide(tAcc[((None, None), 0, 0)], epi_tile)
        tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0)])
        thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
        tTR_tAcc = 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_fragment(tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype)
        return tiled_copy_t2r, tTR_tAcc, tTR_rAcc

    def epilog_smem_copy_and_partition(
        self,
        tiled_copy_t2r: cute.TiledCopy,
        tTR_rC: cute.Tensor,
        tidx: cutlass.Int32,
        sC: cute.Tensor,
    ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
        """Setup Register → SMEM copy for epilogue."""
        copy_atom_r2s = sm100_utils.get_smem_store_op(
            self.c_layout, self.c_dtype, self.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)
        return tiled_copy_r2s, tRS_rC, tRS_sC

    def epilog_gmem_copy_and_partition(
        self,
        tma_atom_c: cute.CopyAtom,
        tCgC: cute.Tensor,
        epi_tile: cute.Tile,
        sC: cute.Tensor,
    ) -> Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]:
        """Setup TMA store SMEM → GMEM for epilogue."""
        gC_epi = cute.flat_divide(tCgC[((None, None), 0, 0, None, None, None)], epi_tile)
        sC_for_tma_partition = cute.group_modes(sC, 0, 2)
        gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2)
        bSG_sC, bSG_gC = cpasync.tma_partition(
            tma_atom_c,
            0,
            cute.make_layout(1),
            sC_for_tma_partition,
            gC_for_tma_partition,
        )
        return tma_atom_c, bSG_sC, bSG_gC


# Kernel cache
_compiled_kernel_cache = {}


def compile_kernel(problem_sizes: List[Tuple[int, int, int, int]]):
    """Compile kernel for given problem sizes."""
    global _compiled_kernel_cache
    
    num_groups = len(problem_sizes)
    cache_key = f"{num_groups}"
    
    if cache_key in _compiled_kernel_cache:
        return _compiled_kernel_cache[cache_key]
    
    # Create pointers for compilation
    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,
    )
    cute_ptr_of_tensor_of_tensormap = make_ptr(
        cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16,
    )
    total_num_clusters = cutlass.Int32(1)
    
    # Create kernel instance and compile
    gemm = GroupGemm(num_groups=num_groups)
    compiled_func = cute.compile(
        gemm,
        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,
    )
    
    _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.
    """
    abc_tensors, _, sfasfb_reordered_tensors, problem_sizes = data

    # Compile kernel for this batch
    compiled_func = compile_kernel(problem_sizes)

    # Extract raw data pointers
    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)
    ):
        abc_ptrs.append((a.data_ptr(), b.data_ptr(), c.data_ptr()))
        sfasfb_ptrs.append((sfa_reordered.data_ptr(), sfb_reordered.data_ptr()))

    # Create metadata tensors
    tensor_of_problem_sizes = torch.tensor(
        problem_sizes, dtype=torch.int32, device="cuda"
    )
    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 grid size
    mma_tiler_mnk = (128, 128, 256)
    cta_tile_shape_mn = [128, mma_tiler_mnk[1]]
    cluster_tile_shape_mn = tuple(
        x * y for x, y in zip(cta_tile_shape_mn, (1, 1))
    )
    
    total_num_clusters = 0
    num_groups = len(problem_sizes)
    for m, n, _, _ in problem_sizes:
        num_clusters_mn = tuple(
            (x + y - 1) // y for x, y in zip((m, n), cluster_tile_shape_mn)
        )
        total_num_clusters += functools.reduce(lambda x, y: x * y, num_clusters_mn)

    # Allocate tensormap buffer
    bytes_per_tensormap = 128
    num_tensormaps = 5
    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
    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 kernel
    compiled_func(
        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,
    )

    res = []
    for i in range(num_groups):
        res.append(abc_tensors[i][2])
    return res