submission 380589

import os
import torch
from task import input_t, output_t

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

import cuda.bindings.driver as cuda

# Ensure CuTeDSL targets Blackwell B200 (sm_100a) unless explicitly overridden.
# If this is not set, the DSL may compile for an "unknown" target and the kernel
# can trap with CUDA_ERROR_ILLEGAL_INSTRUCTION on-device.
os.environ.setdefault("CUTE_DSL_ARCH", "sm_100a")

import cutlass
import cutlass.cute as cute
import cutlass.cute.math as cute_math
from cutlass.cute.runtime import make_ptr
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.torch as cutlass_torch
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


mma_inst_shape_k = 64
ab_dtype = cutlass.Float4E2M1FN
sf_dtype = cutlass.Float8E4M3FN
c_dtype = cutlass.Float16
sf_vec_size = 16


class Sm100BlockScaledDenseDualGemmKernel:
    def __init__(
        self,
        mma_tiler_mn: Tuple[int, int],
        cluster_shape_mn: Tuple[int, int],
    ):
        self.ab_dtype = cutlass.Float4E2M1FN
        self.sf_dtype = cutlass.Float8E4M3FN
        self.acc_dtype = cutlass.Float32
        self.c_dtype = cutlass.Float16
        self.sf_vec_size = 16

        self.epilog_warp_id = (0, 1, 2, 3)
        self.mma_warp_id = 4
        self.tma_warp_id = 5
        self.threads_per_cta = 32 * len((self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id))

        self.mma_tiler_mn = mma_tiler_mn
        self.cluster_shape_mn = cluster_shape_mn

        self.occupancy = 1

        # Keep bring-up conservative; can be tuned later.
        self.force_num_ab_stage: Optional[int] = None
        self.force_num_c_stage: Optional[int] = None

        self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
        self.num_tmem_alloc_cols = 512

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

    def _setup_attributes(self):

        tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
            self.ab_dtype,
            self.a_major_mode,
            self.b_major_mode,
            self.sf_dtype,
            self.sf_vec_size,
            tcgen05.CtaGroup.ONE,
            self.mma_tiler_mn,
        )

        mma_inst_tile_k = 4
        mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
        
        self.mma_tiler = (
            self.mma_tiler_mn[0],
            self.mma_tiler_mn[1],
            mma_inst_shape_k * mma_inst_tile_k,
        )
        self.cta_tile_shape_mnk = (
            self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
            self.mma_tiler[1],
            self.mma_tiler[2],
        )
        self.cluster_layout_vmnk = cute.tiled_divide(
            cute.make_layout((*self.cluster_shape_mn, 1)),
            (tiled_mma.thr_id.shape,),
        )

        self.mma_inst_shape_mn_sfb = (
            self.mma_tiler_mn[0],
            cute.round_up(self.mma_tiler_mn[1], 128),
        )
        
        # Create a specific TiledMMA for SFB using the rounded shape
        tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
            self.ab_dtype,
            self.a_major_mode,
            self.b_major_mode,
            self.sf_dtype,
            self.sf_vec_size,
            tcgen05.CtaGroup.ONE,
            self.mma_inst_shape_mn_sfb,
        )

        # Create specific tiler for SFB
        self.mma_tiler_sfb = (
            self.mma_inst_shape_mn_sfb[0],
            self.mma_inst_shape_mn_sfb[1],
            mma_inst_shape_k * mma_inst_tile_k,
        )

        # Create specific cluster layout for SFB
        self.cluster_layout_sfb_vmnk = cute.tiled_divide(
            cute.make_layout((*self.cluster_shape_mn, 1)),
            (tiled_mma_sfb.thr_id.shape,),
        )

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

        self.epi_tile = sm100_utils.compute_epilogue_tile_shape(
            self.cta_tile_shape_mnk,
            False,
            self.c_layout,
            self.c_dtype,
        )
        self.epi_tile_n = cute.size(self.epi_tile[1])
        self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages(
            tiled_mma,
            self.mma_tiler,
            self.a_dtype,
            self.b_dtype,
            self.epi_tile,
            self.c_dtype,
            self.c_layout,
            self.sf_dtype,
            self.sf_vec_size,
            self.smem_capacity,
            self.occupancy,
        )

        if self.force_num_ab_stage is not None:
            self.num_ab_stage = self.force_num_ab_stage
        if self.force_num_c_stage is not None:
            self.num_c_stage = self.force_num_c_stage

        self.prefetch_stage = self.num_ab_stage

        self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
            tiled_mma,
            self.mma_tiler,
            self.ab_dtype,
            self.num_ab_stage,
        )

        self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
            tiled_mma,
            self.mma_tiler,
            self.ab_dtype,
            self.num_ab_stage,
        )

        self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
            tiled_mma,
            self.mma_tiler,
            self.sf_vec_size,
            self.num_ab_stage,
        )

        
        self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(
            tiled_mma,
            self.mma_tiler,
            self.sf_vec_size,
            self.num_ab_stage,
        )
        self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi(
            self.c_dtype,
            self.c_layout,
            self.epi_tile,
            self.num_c_stage,
        )

    @cute.jit
    def __call__(
        self,
        a_ptr: cute.Pointer,
        b1_ptr: cute.Pointer,
        b2_ptr: cute.Pointer,
        sfa_ptr: cute.Pointer,
        sfb1_ptr: cute.Pointer,
        sfb2_ptr: cute.Pointer,
        c_ptr: cute.Pointer,
        m: cutlass.Int32,
        n: cutlass.Int32,
        k: cutlass.Int32,
        l: cutlass.Int32,
    ):
        self.a_dtype: Type[cutlass.Numeric] = a_ptr.value_type
        self.b_dtype: Type[cutlass.Numeric] = b1_ptr.value_type
        self.sf_dtype: Type[cutlass.Numeric] = sfa_ptr.value_type
        self.c_dtype: Type[cutlass.Numeric] = c_ptr.value_type

        self.a_major_mode, self.b_major_mode, self.c_layout = (
            tcgen05.OperandMajorMode.K,
            tcgen05.OperandMajorMode.K,
            utils.LayoutEnum.ROW_MAJOR,
        )
        self._setup_attributes()

        a_tensor = cute.make_tensor(
            a_ptr,
            cute.make_ordered_layout(
                (cute.assume(m, 32), k, l), order=(1, 0, 2)
            ),
        )
        b1_tensor = cute.make_tensor(
            b1_ptr,
            cute.make_ordered_layout(
                (cute.assume(n, 32), k, l), order=(1, 0, 2)
            ),
        )
        b2_tensor = cute.make_tensor(
            b2_ptr,
            cute.make_ordered_layout(
                (cute.assume(n, 32), k, l), order=(1, 0, 2)
            ),
        )

        c_tensor = cute.make_tensor(
            c_ptr,
            cute.make_ordered_layout(
                (m, cute.assume(n, 32), l), order=(1, 0, 2)
            )
        )

        sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(
            a_tensor.shape, self.sf_vec_size
        )
        sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)

        sfb1_layout = blockscaled_utils.tile_atom_to_shape_SF(
            b1_tensor.shape, self.sf_vec_size
        )
        sfb2_layout = blockscaled_utils.tile_atom_to_shape_SF(
            b2_tensor.shape, self.sf_vec_size
        )
        sfb1_tensor = cute.make_tensor(sfb1_ptr, sfb1_layout)
        sfb2_tensor = cute.make_tensor(sfb2_ptr, sfb2_layout)

        # Standard TiledMMA
        tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
            self.ab_dtype,
            self.a_major_mode,
            self.b_major_mode,
            self.sf_dtype,
            self.sf_vec_size,
            tcgen05.CtaGroup.ONE,
            self.mma_tiler_mn,
        )

        # SFB Specific TiledMMA (Re-created here or stored in self)
        tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
            self.ab_dtype,
            self.a_major_mode,
            self.b_major_mode,
            self.sf_dtype,
            self.sf_vec_size,
            tcgen05.CtaGroup.ONE,
            self.mma_inst_shape_mn_sfb,
        )

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

        b_op = sm100_utils.cluster_shape_to_tma_atom_B(
            self.cluster_shape_mn, tiled_mma.thr_id
        )
        b_smem_layout = cute.slice_(
            self.b_smem_layout_staged, (None, None, None, 0)
        )
        tma_atom_b1, tma_tensor_b1 = cute.nvgpu.make_tiled_tma_atom_B(
            b_op,
            b1_tensor,
            b_smem_layout,
            self.mma_tiler,
            tiled_mma,
            self.cluster_layout_vmnk.shape,
        )
        tma_atom_b2, tma_tensor_b2 = cute.nvgpu.make_tiled_tma_atom_B(
            b_op,
            b2_tensor,
            b_smem_layout,
            self.mma_tiler,
            tiled_mma,
            self.cluster_layout_vmnk.shape,
        )

        sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(
            self.cluster_shape_mn, tiled_mma.thr_id
        )
        # TODO: (Q) is this only used for the tma copy calculation? Otherwise,
        # we use the _staged layout for everything else?
        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(
            sfa_op,
            sfa_tensor,
            sfa_smem_layout,
            self.mma_tiler,
            tiled_mma,
            self.cluster_layout_vmnk.shape,
            internal_type=cutlass.Int16,
        )

        sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(
            self.cluster_shape_mn, tiled_mma.thr_id
        )
        sfb_smem_layout = cute.slice_(
            self.sfb_smem_layout_staged, (None, None, None, 0)
        )
        tma_atom_sfb1, tma_tensor_sfb1 = cute.nvgpu.make_tiled_tma_atom_B(
            sfb_op,
            sfb1_tensor,
            sfb_smem_layout,
            self.mma_tiler_sfb,           # Use SFB tiler
            tiled_mma_sfb,                # Use SFB tiled_mma
            self.cluster_layout_sfb_vmnk.shape, # Use SFB cluster layout
            internal_type=cutlass.Int16,
        )
        tma_atom_sfb2, tma_tensor_sfb2 = cute.nvgpu.make_tiled_tma_atom_B(
            sfb_op,
            sfb2_tensor,
            sfb_smem_layout,
            self.mma_tiler_sfb,           # Use SFB tiler
            tiled_mma_sfb,                # Use SFB tiled_mma
            self.cluster_layout_sfb_vmnk.shape, # Use SFB cluster layout
            internal_type=cutlass.Int16,
        )

        atom_thr_size = cute.size(tiled_mma.thr_id.shape)
        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)
        self.num_tma_load_bytes = (
            a_copy_size
            + 2 * b_copy_size
            + sfa_copy_size
            + 2 * sfb_copy_size
        ) * atom_thr_size

        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(),
            c_tensor,
            epi_smem_layout,
            self.epi_tile,
        )

        grid = self._compute_grid(
            c_tensor, self.cta_tile_shape_mnk, self.cluster_shape_mn
        )

        self.buffer_align_bytes = 1024

        @cute.struct
        class SharedStorage:
            ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage]
            ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage]
            acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
            tmem_dealloc_mbar_ptr: cutlass.Int64
            tmem_holding_buf: cutlass.Int32
            # (EPI_TILE_M, EPI_TILE_N, STAGE)
            sC: cute.struct.Align[
                cute.struct.MemRange[
                    self.c_dtype, cute.cosize(self.c_smem_layout_staged.outer)
                ],
                self.buffer_align_bytes,
            ]
            # (MMA, MMA_M, MMA_K, STAGE)
            sA: cute.struct.Align[
                cute.struct.MemRange[
                    self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)
                ],
                self.buffer_align_bytes,
            ]
            # (MMA, MMA_N, MMA_K, STAGE)
            sB1: cute.struct.Align[
                cute.struct.MemRange[
                    self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)
                ],
                self.buffer_align_bytes,
            ]
            # (MMA, MMA_N, MMA_K, STAGE)
            sB2: cute.struct.Align[
                cute.struct.MemRange[
                    self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)
                ],
                self.buffer_align_bytes,
            ]
            # (MMA, MMA_M, MMA_K, STAGE)
            sSFA: cute.struct.Align[
                cute.struct.MemRange[
                    self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged)
                ],
                self.buffer_align_bytes,
            ]
            # (MMA, MMA_N, MMA_K, STAGE)
            sSFB1: cute.struct.Align[
                cute.struct.MemRange[
                    self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)
                ],
                self.buffer_align_bytes,
            ]
            # (MMA, MMA_N, MMA_K, STAGE)
            sSFB2: cute.struct.Align[
                cute.struct.MemRange[
                    self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)
                ],
                self.buffer_align_bytes,
            ]

        self.shared_storage = SharedStorage

        self.kernel(
            tiled_mma,
            tiled_mma_sfb,                  # Pass specialized SFB MMA
            tma_atom_a,
            tma_tensor_a,
            tma_atom_b1,
            tma_tensor_b1,
            tma_atom_b2,
            tma_tensor_b2,
            tma_atom_sfa,
            tma_tensor_sfa,
            tma_atom_sfb1,
            tma_tensor_sfb1,
            tma_atom_sfb2,
            tma_tensor_sfb2,
            tma_atom_c,
            tma_tensor_c,
            self.cluster_layout_vmnk,
            self.cluster_layout_sfb_vmnk,   # Pass specialized SFB Cluster Layout
            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,
        ).launch(
            grid=grid,
            block=[self.threads_per_cta, 1, 1],
            cluster=(*self.cluster_shape_mn, 1),
            smem=self.shared_storage.size_in_bytes(),
        )
        return

    # GPU device kernel
    @cute.kernel
    def kernel(
        self,
        tiled_mma: cute.TiledMma,
        tiled_mma_sfb: cute.TiledMma,       # Receive SFB Tiled MMA
        tma_atom_a: cute.CopyAtom,
        mA_mkl: cute.Tensor,
        tma_atom_b1: cute.CopyAtom,
        mB1_nkl: cute.Tensor,
        tma_atom_b2: cute.CopyAtom,
        mB2_nkl: cute.Tensor,
        tma_atom_sfa: cute.CopyAtom,
        mSFA_mkl: cute.Tensor,
        tma_atom_sfb1: cute.CopyAtom,
        mSFB1_nkl: cute.Tensor,
        tma_atom_sfb2: cute.CopyAtom,
        mSFB2_nkl: cute.Tensor,
        tma_atom_c: Optional[cute.CopyAtom],
        mC_mnl: cute.Tensor,
        cluster_layout_vmnk: cute.Layout,
        cluster_layout_sfb_vmnk: cute.Layout, # Receive SFB Cluster 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, None],
        epi_tile: cute.Tile,
    ):
        warp_idx = cute.arch.warp_idx()
        warp_idx = cute.arch.make_warp_uniform(warp_idx)

        # Prefetch descriptors with dedicated TMA warp
        if warp_idx == self.tma_warp_id:
            cpasync.prefetch_descriptor(tma_atom_a)
            cpasync.prefetch_descriptor(tma_atom_b1)
            cpasync.prefetch_descriptor(tma_atom_b2)
            cpasync.prefetch_descriptor(tma_atom_sfa)
            cpasync.prefetch_descriptor(tma_atom_sfb1)
            cpasync.prefetch_descriptor(tma_atom_sfb2)
            cpasync.prefetch_descriptor(tma_atom_c)

        bidx, bidy, bidz = cute.arch.block_idx()
        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
        )

        cta_coord = (bidx, bidy, bidz)
        mma_tile_coord_mnl = (
            cta_coord[0] // cute.size(tiled_mma.thr_id.shape),
            cta_coord[1],
            cta_coord[2],
        )

        tidx, _, _ = cute.arch.thread_idx()

        smem = utils.SmemAllocator()
        storage = smem.allocate(self.shared_storage)

        ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
        num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
        ab_pipeline_consumer_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread, num_tma_producer
        )
        ab_pipeline = pipeline.PipelineTmaAsync.create(
            barrier_storage=storage.ab_full_mbar_ptr.data_ptr(),
            num_stages=self.num_ab_stage,
            producer_group=ab_pipeline_producer_group,
            consumer_group=ab_pipeline_consumer_group,
            tx_count=self.num_tma_load_bytes,
            cta_layout_vmnk=cluster_layout_vmnk,
            defer_sync=True,
        )

        acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
        num_acc_consumer_threads = len(self.epilog_warp_id)
        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=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],
        )

        cute.arch.cluster_arrive_relaxed()

        # (EPI_TILE_M, EPI_TILE_N, STAGE)=((8,16),(32,1),(1,3))
        sC = storage.sC.get_tensor(
            c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner
        )
        # (MMA, MMA_M, MMA_K, STAGE)= ((128,64),1,4,7)
        sA = storage.sA.get_tensor(
            a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner
        )
        # (MMA, MMA_N, MMA_K, STAGE)=((64,64),1,4,7)
        sB1 = storage.sB1.get_tensor(
            b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner
        )
        sB2 = storage.sB2.get_tensor(
            b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner
        )
        # (MMA, MMA_M, MMA_K, STAGE)= ((((32,4),1),(16,4)),1,4,7)
        sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged)
        # (MMA, MMA_N, MMA_K, STAGE)= ((((32,4),1),(16,4)),1,4,7)
        sSFB1 = storage.sSFB1.get_tensor(sfb_smem_layout_staged)
        sSFB2 = storage.sSFB2.get_tensor(sfb_smem_layout_staged)

        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
        )

        #
        # Local_tile partition global tensors
        #
        # (bM, bK, RestM, RestK, RestL)=(128,256,?,?,?)
        gA_mkl = cute.local_tile(
            mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
        )
        # (bN, bK, RestN, RestK, RestL)=(64,256,?,?,?)
        gB1_nkl = cute.local_tile(
            mB1_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
        )
        gB2_nkl = cute.local_tile(
            mB2_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
        )
        # (bM, bK, RestM, RestK, RestL)=((32,4),(16,4,4),?,?,(1,?))
        gSFA_mkl = cute.local_tile(
            mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
        )
        # (bN, bK, RestN, RestK, RestL)=((32,4),(16,4,4),?,?,(1,?))
        gSFB1_nkl = cute.local_tile(
            mSFB1_nkl,
            cute.slice_(self.mma_tiler_sfb, (0, None, None)),
            (None, None, None),
        )
        gSFB2_nkl = cute.local_tile(
            mSFB2_nkl,
            cute.slice_(self.mma_tiler_sfb, (0, None, None)),
            (None, None, None),
        )
        # (bM, bN, RestM, RestN, RestL)=(128,64,?,?,?)
        gC_mnl = cute.local_tile(
            mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
        )
        k_block_cnt = cute.size(gA_mkl, mode=[3])

        thr_mma = tiled_mma.get_slice(0)
        thr_mma_sfb = tiled_mma_sfb.get_slice(0) # Get slice for SFB

        #
        # Partition global tensor for TiledMMA_A/B/C
        #
        thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
        thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v)
        # (MMA, MMA_M, MMA_K, RestM, RestK, RestL)= ((128,64),1,4,?,?,?)
        tCgA = thr_mma.partition_A(gA_mkl)
        # (MMA, MMA_N, MMA_K, RestN, RestK, RestL)=((64,64),1,4,?,?,?)
        tCgB1 = thr_mma.partition_B(gB1_nkl)
        tCgB2 = thr_mma.partition_B(gB2_nkl)
        # (MMA, MMA_M, MMA_K, RestM, RestK, RestL)=(((32,4),(16,4)),1,4,?,?,(1,?))
        tCgSFA = thr_mma.partition_A(gSFA_mkl)
        # (MMA, MMA_N, MMA_K, RestN, RestK, RestL)=(((32,4),(16,4)),1,4,?,?,(1,?))
        tCgSFB1 = thr_mma_sfb.partition_B(gSFB1_nkl)
        tCgSFB2 = thr_mma_sfb.partition_B(gSFB2_nkl)
        # (MMA, MMA_M, MMA_N, RestM, RestN, RestL)=((128,64),1,1,?,?,?)
        tCgC = thr_mma.partition_C(gC_mnl)

        #
        # Partition global/shared tensor for TMA load A/B
        #
        # TMA load A partition_S/D
        a_cta_layout = cute.make_layout(
            cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape
        )
        # ((atom_v, rest_v), STAGE)
        # ((atom_v=(256,128), rest_v), RestM, RestK, RestL)
        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),
        )
        # TMA load B partition_S/D
        b_cta_layout = cute.make_layout(
            cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape
        )
        # ((atom_v, rest_v), STAGE)
        # ((atom_v=(256,64), rest_v), RestN, RestK, RestL)
        tBsB, tBgB = cpasync.tma_partition(
            tma_atom_b1,
            block_in_cluster_coord_vmnk[1],
            b_cta_layout,
            cute.group_modes(sB1, 0, 3),
            cute.group_modes(tCgB1, 0, 3),
        )
        tBsB2, tBgB2 = cpasync.tma_partition(
            tma_atom_b2,
            block_in_cluster_coord_vmnk[1],
            b_cta_layout,
            cute.group_modes(sB2, 0, 3),
            cute.group_modes(tCgB2, 0, 3),
        )
        #  TMALDG_SFA partition_S/D
        sfa_cta_layout = a_cta_layout
        # ((atom_v, rest_v), STAGE)
        # ((atom_v=(512,4), rest_v=16->1(after filter)), RestM, RestK, RestL)
        tAsSFA, tAgSFA = cute.nvgpu.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)

        # TMALDG_SFB partition_S/D
        sfb_cta_layout = cute.make_layout(
            cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape
        )
        # ((atom_v, rest_v), STAGE)
        # ((atom_v=(512,4), rest_v=16->1(after filter)), RestM, RestK, RestL)
        tBsSFB1, tBgSFB1 = cute.nvgpu.cpasync.tma_partition(
            tma_atom_sfb1,
            block_in_cluster_coord_sfb_vmnk[1],
            sfb_cta_layout,
            cute.group_modes(sSFB1, 0, 3),
            cute.group_modes(tCgSFB1, 0, 3),
        )
        tBsSFB1 = cute.filter_zeros(tBsSFB1)
        tBgSFB1 = cute.filter_zeros(tBgSFB1)

        tBsSFB2, tBgSFB2 = cute.nvgpu.cpasync.tma_partition(
            tma_atom_sfb2,
            block_in_cluster_coord_sfb_vmnk[1],
            sfb_cta_layout,
            cute.group_modes(sSFB2, 0, 3),
            cute.group_modes(tCgSFB2, 0, 3),
        )
        tBsSFB2 = cute.filter_zeros(tBsSFB2)
        tBgSFB2 = cute.filter_zeros(tBgSFB2)

        # (MMA, MMA_M, MMA_K, STAGE) = (1,1,4,7)
        tCrA = tiled_mma.make_fragment_A(sA)
        # (MMA, MMA_N, MMA_K, STAGE) = (1,1,4,7)
        tCrB1 = tiled_mma.make_fragment_B(sB1)
        tCrB2 = tiled_mma.make_fragment_B(sB2)
        # (MMA, MMA_M, MMA_N)=((128, 64), 1, 1)
        acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
        # (MMA, MMA_M, MMA_N)=((128,64),1,1)
        tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape)

        cute.arch.cluster_wait()

        # ---------- TMA warp: AB producer ----------
        if warp_idx == self.tma_warp_id:
            ab_producer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer, self.num_ab_stage
            )
            # ((atom_v, rest_v), RestK)= (((256,128),1),?)
            tAgA_slice = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
            # ((atom_v, rest_v), RestK)=(((256,64),1),?)
            tBgB_slice = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
            tBgB2_slice = tBgB2[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
            # ((atom_v, rest_v), RestK)=(((512,4),1),?)
            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.cta_tile_shape_mnk[1] == 64):
                slice_n = mma_tile_coord_mnl[1] // 2
            # ((atom_v, rest_v), RestK)=(((512,4),1),?)
            tBgSFB1_slice = tBgSFB1[(None, slice_n, None, mma_tile_coord_mnl[2])]
            tBgSFB2_slice = tBgSFB2[(None, slice_n, None, mma_tile_coord_mnl[2])]

            for prefetch_tile in cutlass.range(0, self.prefetch_stage, unroll=1):
                cute.prefetch(
                    tma_atom_a,
                    tAgA_slice[(None, prefetch_tile)]
                )
                cute.prefetch(
                    tma_atom_b1,
                    tBgB_slice[(None, prefetch_tile)]
                )
                cute.prefetch(
                    tma_atom_b2,
                    tBgB2_slice[(None, prefetch_tile)]
                )
                cute.prefetch(
                    tma_atom_sfa,
                    tAgSFA_slice[(None, prefetch_tile)]
                )
                cute.prefetch(
                    tma_atom_sfb1,
                    tBgSFB1_slice[(None, prefetch_tile)]
                )
                cute.prefetch(
                    tma_atom_sfb2,
                    tBgSFB2_slice[(None, prefetch_tile)]
                )

            peek_ab_empty_status = ab_pipeline.producer_try_acquire(
                ab_producer_state
            )

            for k_block_idx in cutlass.range(0, k_block_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_b1,
                    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_b2,
                    tBgB2_slice[(None, ab_producer_state.count)],
                    tBsB2[(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_sfb1,
                    tBgSFB1_slice[(None, ab_producer_state.count)],
                    tBsSFB1[(None, ab_producer_state.index)],
                    tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
                    mcast_mask=sfb_full_mcast_mask,
                )
                cute.copy(
                    tma_atom_sfb2,
                    tBgSFB2_slice[(None, ab_producer_state.count)],
                    tBsSFB2[(None, ab_producer_state.index)],
                    tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
                    mcast_mask=sfb_full_mcast_mask,
                )

                if k_block_idx < k_block_cnt - self.prefetch_stage:
                    next_k_idx = ab_producer_state.count + self.prefetch_stage
                    cute.prefetch(
                        tma_atom_a,
                        tAgA_slice[(None, next_k_idx)]
                    )
                    cute.prefetch(
                        tma_atom_b1,
                        tBgB_slice[(None, next_k_idx)]
                    )
                    cute.prefetch(
                        tma_atom_b2,
                        tBgB2_slice[(None, next_k_idx)]
                    )
                    cute.prefetch(
                        tma_atom_sfa,
                        tAgSFA_slice[(None, next_k_idx)]
                    )
                    cute.prefetch(
                        tma_atom_sfb1,
                        tBgSFB1_slice[(None, next_k_idx)]
                    )
                    cute.prefetch(
                        tma_atom_sfb2,
                        tBgSFB2_slice[(None, next_k_idx)]
                    )

                ab_producer_state.advance()
                if ab_producer_state.count < k_block_cnt:
                    peek_ab_empty_status = ab_pipeline.producer_try_acquire(
                        ab_producer_state
                    )

            ab_pipeline.producer_tail(ab_producer_state)

        # ---------- MMA warp: AB consumer + GEMM + ACC producer ----------
        elif warp_idx == self.mma_warp_id:

            tmem.wait_for_alloc()
            #
            # Retrieving tensor memory ptr and make accumulator/SFA/SFB tensor
            #
            acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) # tcgen05.find_tmem_tensor_col_offset(tCtAcc) = 64
            # Make accumulator tmem tensors (dual)
            # (MMA, MMA_M, MMA_N, STAGE)= ((128,64),1,1)
            tCtAcc1 = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
            acc2_tmem_ptr = acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc1)
            tCtAcc2 = cute.make_tensor(acc2_tmem_ptr, tCtAcc_fake.layout)
            # Make SFA tmem tensor    
            sfa_tmem_ptr = cute.recast_ptr(
                acc2_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc2),
                dtype=self.sf_dtype,
            )
            # (MMA, MMA_M, MMA_K)=((((32,4),4),(16,4)),1,4)
            tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
                tiled_mma,
                self.mma_tiler,
                self.sf_vec_size,
                cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)),
            )
            tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) #tcgen05.find_tmem_tensor_col_offset(tCtSFA) = 16
            # Make SFB tmem tensor
            sfb_tmem_ptr = cute.recast_ptr(
                acc2_tmem_ptr
                + tcgen05.find_tmem_tensor_col_offset(tCtAcc2)
                + tcgen05.find_tmem_tensor_col_offset(tCtSFA),
                dtype=self.sf_dtype,
            )
            # (MMA, MMA_N, MMA_K)=((((32,4),4),(16,4)),1,4)
            tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
                tiled_mma,
                self.mma_tiler,
                self.sf_vec_size,
                cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)),
            )
            tCtSFB1 = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) #tcgen05.find_tmem_tensor_col_offset(tCtSFB)=16

            # IMPORTANT: tcgen05.find_tmem_tensor_col_offset(...) is expressed in TMEM
            # column units; do the pointer arithmetic on the original Float32 TMEM
            # pointer and only then recast to sf dtype. Doing `+ offset` on a recast
            # sf pointer applies element-based pointer arithmetic and corrupts layout.
            sfb2_tmem_ptr = cute.recast_ptr(
                acc2_tmem_ptr
                + tcgen05.find_tmem_tensor_col_offset(tCtAcc2)
                + tcgen05.find_tmem_tensor_col_offset(tCtSFA)
                + tcgen05.find_tmem_tensor_col_offset(tCtSFB1),
                dtype=self.sf_dtype,
            )
            tCtSFB2 = cute.make_tensor(sfb2_tmem_ptr, tCtSFB_layout)
            #
            # Partition for S2T copy of SFA/SFB
            #
            # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)=((((32, 1, 1), 4), 1), 1, 1, 4, 7)
            # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)=(((32, 16, 4), 1), 1, 1, 4)
            tiled_copy_s2t_sfa, tCsSFA_compact_s2t, tCtSFA_compact_s2t = (
                self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA)
            )
            # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)=((((32, 1, 1), 4), 1), 1, 1, 4, 7)
            # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)=(((32, 16, 4), 1), 1, 1, 4)
            tiled_copy_s2t_sfb1, tCsSFB1_compact_s2t, tCtSFB1_compact_s2t = (
                self.mainloop_s2t_copy_and_partition(sSFB1, tCtSFB1)
            )
            tiled_copy_s2t_sfb2, tCsSFB2_compact_s2t, tCtSFB2_compact_s2t = (
                self.mainloop_s2t_copy_and_partition(sSFB2, tCtSFB2)
            )
            ab_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.num_ab_stage
            )
            acc_producer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer, self.num_acc_stage
            )

            # Peek initial full
            peek_ab_full_status = ab_pipeline.consumer_try_wait(
                ab_consumer_state
            )
            
            tCtSFB1_mma = tCtSFB1
            tCtSFB2_mma = tCtSFB2
            if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
                # Move in increments of 64 columns of SFB
                offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2)
                # IMPORTANT: match top2's addressing pattern by applying the
                # column-offset arithmetic before the recast.
                shifted_ptr1 = cute.recast_ptr(
                    acc2_tmem_ptr
                    + tcgen05.find_tmem_tensor_col_offset(tCtAcc2)
                    + tcgen05.find_tmem_tensor_col_offset(tCtSFA)
                    + offset,
                    dtype=self.sf_dtype,
                )
                shifted_ptr2 = cute.recast_ptr(
                    acc2_tmem_ptr
                    + tcgen05.find_tmem_tensor_col_offset(tCtAcc2)
                    + tcgen05.find_tmem_tensor_col_offset(tCtSFA)
                    + tcgen05.find_tmem_tensor_col_offset(tCtSFB1)
                    + offset,
                    dtype=self.sf_dtype,
                )
                tCtSFB1_mma = cute.make_tensor(shifted_ptr1, tCtSFB_layout)
                tCtSFB2_mma = cute.make_tensor(shifted_ptr2, tCtSFB_layout)


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

            for _ in range(k_block_cnt):
                ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status)

                s2t_stage_coord = (
                    None,
                    None,
                    None,
                    None,
                    ab_consumer_state.index,
                )
                tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord]
                tCsSFB1_compact_s2t_staged = tCsSFB1_compact_s2t[s2t_stage_coord]
                tCsSFB2_compact_s2t_staged = tCsSFB2_compact_s2t[s2t_stage_coord]
                cute.copy(
                    tiled_copy_s2t_sfa,
                    tCsSFA_compact_s2t_staged,
                    tCtSFA_compact_s2t,
                )
                cute.copy(
                    tiled_copy_s2t_sfb1,
                    tCsSFB1_compact_s2t_staged,
                    tCtSFB1_compact_s2t,
                )
                cute.copy(
                    tiled_copy_s2t_sfb2,
                    tCsSFB2_compact_s2t_staged,
                    tCtSFB2_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,
                        tCtSFB1_mma[sf_kphase_coord].iterator,
                    )
                    cute.gemm(
                        tiled_mma,
                        tCtAcc1,
                        tCrA[kphase_coord],
                        tCrB1[kphase_coord],
                        tCtAcc1,
                    )
                    tiled_mma.set(
                        tcgen05.Field.SFB,
                        tCtSFB2_mma[sf_kphase_coord].iterator,
                    )
                    cute.gemm(
                        tiled_mma,
                        tCtAcc2,
                        tCrA[kphase_coord],
                        tCrB2[kphase_coord],
                        tCtAcc2,
                    )
                    tiled_mma.set(tcgen05.Field.ACCUMULATE, True)

                ab_pipeline.consumer_release(ab_consumer_state)

                ab_consumer_state.advance()
                if ab_consumer_state.count < k_block_cnt:
                    peek_ab_full_status = ab_pipeline.consumer_try_wait(
                        ab_consumer_state
                    )

            acc_pipeline.producer_commit(acc_producer_state)
        
        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(self.acc_dtype)
            tCtAcc1 = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
            acc2_tmem_ptr = acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc1)
            tCtAcc2 = cute.make_tensor(acc2_tmem_ptr, tCtAcc_fake.layout)

            tiled_copy_t2r, tTR_tAcc1, tTR_rAcc1 = self.epilog_tmem_copy_and_partition(
                tidx, tCtAcc1, tCgC, epi_tile
            )
            _, tTR_tAcc2, tTR_rAcc2 = self.epilog_tmem_copy_and_partition(
                tidx, tCtAcc2, tCgC, epi_tile
            )

            tTR_rC = cute.make_fragment(tTR_rAcc1.shape, self.c_dtype)
            tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition(
                tiled_copy_t2r, tTR_rC, tidx, sC
            )
            tma_atom_c, bSG_sC, bSG_gC = self.epilog_gmem_copy_and_partition(
                tidx, tma_atom_c, tCgC, epi_tile, sC
            )
            bSG_gC = bSG_gC[(None, None, None, *mma_tile_coord_mnl)]

            acc_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.num_acc_stage
            )

            acc_pipeline.consumer_wait(acc_consumer_state)

            tTR_tAcc1 = cute.group_modes(tTR_tAcc1, 3, cute.rank(tTR_tAcc1))
            tTR_tAcc2 = cute.group_modes(tTR_tAcc2, 3, cute.rank(tTR_tAcc2))
            bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC))

            subtile_cnt = cute.size(tTR_tAcc1.shape, mode=[3])

            for subtile_idx in range(subtile_cnt):
                tTR_tAcc1_mn = tTR_tAcc1[(None, None, None, subtile_idx)]
                tTR_tAcc2_mn = tTR_tAcc2[(None, None, None, subtile_idx)]
                cute.copy(tiled_copy_t2r, tTR_tAcc1_mn, tTR_rAcc1)
                cute.copy(tiled_copy_t2r, tTR_tAcc2_mn, tTR_rAcc2)

                x = tTR_rAcc1.load()
                y = tTR_rAcc2.load()
                # silu(x) = x / (1 + exp(-x))
                # Prefer non-fastmath here to match PyTorch reference more closely.
                silu_x = x / (cutlass.Float32(1.0) + cute_math.exp(-x, fastmath=True))
                # silu_x = x / (cutlass.Float32(1.0) + cute_math.exp(-x, fastmath=False))
                out = (silu_x * y).to(self.c_dtype)
                tRS_rC.store(out)

                cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, subtile_idx)])
                cute.arch.fence_view_async_shared()

                # Ensure all epilogue warps have finished writing sC for this subtile
                # before warp0 triggers the TMA store.
                self.epilog_sync_barrier.arrive_and_wait()

                if warp_idx == self.epilog_warp_id[0]:
                    cute.copy(
                        tma_atom_c,
                        bSG_sC[(None, subtile_idx)],
                        bSG_gC[(None, subtile_idx)],
                    )

                # Keep epilogue warps in sync across subtiles.
                # self.epilog_sync_barrier.arrive_and_wait()
                    
            tmem.relinquish_alloc_permit()
            tmem.free(acc_tmem_ptr)
    
    def mainloop_s2t_copy_and_partition(
        self,
        sSF: cute.Tensor,
        tSF: cute.Tensor,
    ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
        """
        Make tiledCopy for smem to tmem load for scale factor tensor, then use it to partition smem memory (source) and tensor memory (destination).

        :param sSF: The scale factor tensor in smem
        :type sSF: cute.Tensor
        :param tSF: The scale factor tensor in tmem
        :type tSF: cute.Tensor

        :return: A tuple containing (tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t) where:
            - tiled_copy_s2t: The tiled copy operation for smem to tmem load for scale factor tensor(s2t)
            - tCsSF_compact_s2t: The partitioned scale factor tensor in smem
            - tSF_compact_s2t: The partitioned scale factor tensor in tmem
        :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]
        """
        # (MMA, MMA_MN, MMA_K, STAGE)
        tCsSF_compact = cute.filter_zeros(sSF)
        # (MMA, MMA_MN, MMA_K)
        tCtSF_compact = cute.filter_zeros(tSF)

        # Make S2T CopyAtom and tiledCopy
        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)

        # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
        tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact)
        # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
        tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
            tiled_copy_s2t, tCsSF_compact_s2t_
        )
        # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
        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,
        gC_mnl: cute.Tensor,
        epi_tile: cute.Tile,
    ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
        """
        Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination).

        :param tidx: The thread index in epilogue warp groups
        :type tidx: cutlass.Int32
        :param tAcc: The accumulator tensor to be copied and partitioned
        :type tAcc: cute.Tensor
        :param gC_mnl: The global tensor C
        :type gC_mnl: cute.Tensor
        :param epi_tile: The epilogue tiler
        :type epi_tile: cute.Tile

        :return: A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where:
            - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r)
            - tTR_tAcc: The partitioned accumulator tensor
            - tTR_rAcc: The accumulated tensor in register used to hold t2r results
        :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]
        """
        # Make tiledCopy for tensor memory load
        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
        )
        # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N)
        tAcc_epi = cute.flat_divide(
            tAcc[((None, None), 0, 0)],
            epi_tile,
        )
        # (EPI_TILE_M, EPI_TILE_N)
        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)
        # (T2R, T2R_M, T2R_N, EPI_M, EPI_M)
        tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi)

        # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL)
        gC_mnl_epi = cute.flat_divide(
            gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile
        )
        # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL)
        tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi)
        # (T2R, T2R_M, T2R_N)
        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]:
        """
        Make tiledCopy for shared memory store, then use it to partition register array (source) and shared memory (destination).

        :param tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r)
        :type tiled_copy_t2r: cute.TiledCopy
        :param tTR_rC: The partitioned accumulator tensor
        :type tTR_rC: cute.Tensor
        :param tidx: The thread index in epilogue warp groups
        :type tidx: cutlass.Int32
        :param sC: The shared memory tensor to be copied and partitioned
        :type sC: cute.Tensor

        :return: A tuple containing (tiled_copy_r2s, tRS_rC, tRS_sC) where:
            - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s)
            - tRS_rC: The partitioned tensor C (register source)
            - tRS_sC: The partitioned tensor C (smem destination)
        :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]
        """
        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)
        # (R2S, R2S_M, R2S_N, PIPE_D)
        thr_copy_r2s = tiled_copy_r2s.get_slice(tidx)
        tRS_sC = thr_copy_r2s.partition_D(sC)
        # (R2S, R2S_M, R2S_N)
        tRS_rC = tiled_copy_r2s.retile(tTR_rC)
        return tiled_copy_r2s, tRS_rC, tRS_sC

    def epilog_gmem_copy_and_partition(
        self,
        tidx: cutlass.Int32,
        atom: Union[cute.CopyAtom, cute.TiledCopy],
        gC_mnl: cute.Tensor,
        epi_tile: cute.Tile,
        sC: cute.Tensor,
    ) -> Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]:
        """Make tiledCopy for global memory store, then use it to:
        partition shared memory (source) and global memory (destination) for TMA store version.

        :param tidx: The thread index in epilogue warp groups
        :type tidx: cutlass.Int32
        :param atom: The copy_atom_c to be used for TMA store version, or tiled_copy_t2r for none TMA store version
        :type atom: cute.CopyAtom or cute.TiledCopy
        :param gC_mnl: The global tensor C
        :type gC_mnl: cute.Tensor
        :param epi_tile: The epilogue tiler
        :type epi_tile: cute.Tile
        :param sC: The shared memory tensor to be copied and partitioned
        :type sC: cute.Tensor

        :return: A tuple containing (tma_atom_c, bSG_sC, bSG_gC) where:
            - tma_atom_c: The TMA copy atom
            - bSG_sC: The partitioned shared memory tensor C
            - bSG_gC: The partitioned global tensor C
        :rtype: Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]
        """
        # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL)
        gC_epi = cute.flat_divide(
            gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile
        )

        tma_atom_c = atom
        sC_for_tma_partition = cute.group_modes(sC, 0, 2)
        gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2)
        # ((ATOM_V, REST_V), EPI_M, EPI_N)
        # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL)
        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
    
    @staticmethod
    def _compute_stages(
        tiled_mma: cute.TiledMma,
        mma_tiler_mnk: Tuple[int, int, int],
        a_dtype: Type[cutlass.Numeric],
        b_dtype: Type[cutlass.Numeric],
        epi_tile: cute.Tile,
        c_dtype: Type[cutlass.Numeric],
        c_layout: utils.LayoutEnum,
        sf_dtype: Type[cutlass.Numeric],
        sf_vec_size: int,
        smem_capacity: int,
        occupancy: int,
    ) -> Tuple[int, int, int]:
        """Computes the number of stages for A/B/C operands based on heuristics.

        :param tiled_mma: The tiled MMA object defining the core computation.
        :type tiled_mma: cute.TiledMma
        :param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler.
        :type mma_tiler_mnk: tuple[int, int, int]
        :param a_dtype: Data type of operand A.
        :type a_dtype: type[cutlass.Numeric]
        :param b_dtype: Data type of operand B.
        :type b_dtype: type[cutlass.Numeric]
        :param epi_tile: The epilogue tile shape.
        :type epi_tile: cute.Tile
        :param c_dtype: Data type of operand C (output).
        :type c_dtype: type[cutlass.Numeric]
        :param c_layout: Layout enum of operand C.
        :type c_layout: utils.LayoutEnum
        :param sf_dtype: Data type of Scale factor.
        :type sf_dtype: type[cutlass.Numeric]
        :param sf_vec_size: Scale factor vector size.
        :type sf_vec_size: int
        :param smem_capacity: Total available shared memory capacity in bytes.
        :type smem_capacity: int
        :param occupancy: Target number of CTAs per SM (occupancy).
        :type occupancy: int

        :return: A tuple containing the computed number of stages for:
                 (ACC stages, A/B operand stages, C stages)
        :rtype: tuple[int, int, int]
        """
        # ACC stages
        num_acc_stage = 1

        # Default C stages
        num_c_stage = 2

        # Calculate smem layout and size for one stage of A, B, SFA, SFB and C
        a_smem_layout_stage_one = sm100_utils.make_smem_layout_a(
            tiled_mma,
            mma_tiler_mnk,
            a_dtype,
            1,  # a tmp 1 stage is provided
        )
        b_smem_layout_staged_one = sm100_utils.make_smem_layout_b(
            tiled_mma,
            mma_tiler_mnk,
            b_dtype,
            1,  # a tmp 1 stage is provided
        )
        sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa(
            tiled_mma,
            mma_tiler_mnk,
            sf_vec_size,
            1,  # a tmp 1 stage is provided
        )
        sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb(
            tiled_mma,
            mma_tiler_mnk,
            sf_vec_size,
            1,  # a tmp 1 stage is provided
        )

        c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi(
            c_dtype,
            c_layout,
            epi_tile,
            1,
        )

        # Dual GEMM loads two B operands and two SFB scale-factor tensors per stage.
        ab_bytes_per_stage = (
            cute.size_in_bytes(a_dtype, a_smem_layout_stage_one)
            + 2 * cute.size_in_bytes(b_dtype, b_smem_layout_staged_one)
            + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one)
            + 2 * 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

        # Calculate A/B/SFA/SFB stages:
        # Start with total smem per CTA (capacity / occupancy)
        # Subtract reserved bytes and initial C stages bytes
        # Divide remaining by bytes needed per A/B/SFA/SFB stage
        num_ab_stage = (
            smem_capacity - (mbar_helpers_bytes + c_bytes)
        ) // ab_bytes_per_stage

        # Refine epilogue stages:
        # Calculate remaining smem after allocating for A/B/SFA/SFB stages and reserved bytes
        # Add remaining unused smem to epilogue
        num_c_stage += (
            smem_capacity
            -  ab_bytes_per_stage * num_ab_stage
            -  (mbar_helpers_bytes + c_bytes)
        ) // (c_bytes_per_stage)

        return num_acc_stage, num_ab_stage, num_c_stage

    @staticmethod
    def _compute_grid(
        c: cute.Tensor,
        cta_tile_shape_mnk: Tuple[int, int, int],
        cluster_shape_mn: Tuple[int, int],
    ) -> Tuple[int, int, int]:
        """Compute grid shape for the output tensor C.

        :param c: The output tensor C
        :type c: cute.Tensor
        :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile.
        :type cta_tile_shape_mnk: tuple[int, int, int]
        :param cluster_shape_mn: Shape of each cluster in M, N dimensions.
        :type cluster_shape_mn: tuple[int, int]

        :return: Grid shape for kernel launch.
        :rtype: tuple[int, int, int]
        """

        grid = (
                cute.ceil_div(c.layout.shape[0], cta_tile_shape_mnk[0]),
                cute.ceil_div(c.layout.shape[1], cta_tile_shape_mnk[1]),
                c.layout.shape[2],
        )
        
        return grid
    
_compiled_kernel_cache = None


def compile_kernel():
    global _compiled_kernel_cache
    if _compiled_kernel_cache is not None:
        return _compiled_kernel_cache
    
    a_ptr = make_ptr(
        ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16
    )
    b1_ptr = make_ptr(
        ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16
    )
    b2_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
    )
    sfb1_ptr = make_ptr(
        sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32
    )
    sfb2_ptr = make_ptr(
        sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32
    )

    # Use a legality-safe cluster shape for all official shapes.
    # my_kernel = Sm100BlockScaledDenseDualGemmKernel((128, 64), (1, 1))
    # my_kernel = Sm100BlockScaledDenseDualGemmKernel((128, 64), (1, 2))
    # my_kernel = Sm100BlockScaledDenseDualGemmKernel((128, 128), (1, 1))
    # my_kernel = Sm100BlockScaledDenseDualGemmKernel((128, 64), (4, 1))
    my_kernel = Sm100BlockScaledDenseDualGemmKernel((128, 64), (2, 2))
    _compiled_kernel_cache = cute.compile(
        my_kernel, a_ptr, b1_ptr, b2_ptr, sfa_ptr, sfb1_ptr, sfb2_ptr, c_ptr, 0, 0, 0, 0
    )
    return _compiled_kernel_cache


def custom_kernel(data: input_t) -> output_t:
    a, b1, b2, _, _, _, sfa_permuted, sfb1_permuted, sfb2_permuted, c = data

    compiled = compile_kernel()

    m, k_half, l = a.shape
    n, _, _ = b1.shape
    k = k_half * 2

    a_ptr = make_ptr(ab_dtype, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
    b1_ptr = make_ptr(ab_dtype, b1.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
    b2_ptr = make_ptr(ab_dtype, b2.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)
    sfb1_ptr = make_ptr(sf_dtype, sfb1_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32)
    sfb2_ptr = make_ptr(sf_dtype, sfb2_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32)

    compiled(a_ptr, b1_ptr, b2_ptr, sfa_ptr, sfb1_ptr, sfb2_ptr, c_ptr, m, n, k, l)
    return c