submission 383079

import cuda.bindings.driver as cuda

import os

import torch

from task import input_t, output_t

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

# TMA cache eviction policy constants
_TMA_CACHE_EVICT_NORMAL = 0x1000000000000000
_TMA_CACHE_EVICT_FIRST = 0x12F0000000000000
_TMA_CACHE_EVICT_LAST = 0x14F0000000000000

# -----------------------------------------------------------------------------
# Problem constants
# -----------------------------------------------------------------------------

# A/B are NVFP4 (e2m1) and scale factors are FP8 (e4m3fnuz). Output is FP16.
AB_DTYPE = cutlass.Float4E2M1FN
SF_DTYPE = cutlass.Float8E4M3FN
C_DTYPE = cutlass.Float16

SF_VEC_SIZE = 16


def _silu(x):
    # x * sigmoid(x) using tanh: sigmoid(x) = 0.5 * (1 + tanh(x/2))
    return x * 0.5 * (1.0 + cute.math.tanh(x * 0.5, fastmath=True))


# -----------------------------------------------------------------------------
# Persistent dual-GEMM kernel (warp-specialized) for SM100 (Blackwell)
# -----------------------------------------------------------------------------


class Sm100BlockScaledPersistentDualGemmKernel:
    """Persistent blockscaled dual GEMM:

    C = silu(A @ B1) * (A @ B2)

    This is adapted from CUTLASS CuTeDSL Blackwell persistent blockscaled GEMM
    example, extended to load/compute two B matrices and fuse SwiGLU-style
    epilogue.
    """

    def __init__(
        self,
        *,
        sf_vec_size: int,
        mma_tiler_mn: tuple[int, int],
        cluster_shape_mn: tuple[int, int],
        ab_stages: int = 0,
        prefetch_dist: int = -1,
    ):
        self.acc_dtype = cutlass.Float32
        self.sf_vec_size = sf_vec_size

        self.use_2cta_instrs = mma_tiler_mn[0] == 256
        self.cluster_shape_mn = cluster_shape_mn
        self.mma_tiler = (*mma_tiler_mn, 1)  # K is deferred

        self.cta_group = tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE

        # Warp specialization (same as CUTLASS example)
        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))

        # Barriers
        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)),
        )

        self.occupancy = 1
        self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")

        # TMEM capacity (SM100)
        self.num_tmem_alloc_cols = 512

        # Optional overrides
        # - ab_stages: 0 = auto, >0 = override (clamped to max that fits)
        # - prefetch_dist: -1 = auto (=num_ab_stage), 0 = disable,
        #                 >0 = explicit (clamped to <= num_ab_stage)
        self.ab_stages_param = ab_stages
        self.prefetch_dist_param = prefetch_dist

        self.prefetch_dist = 0
        self.prefetch_enabled = False

    def _setup_attributes(self):
        # MMA inst shapes (MN is configured, K is fixed by tcgen05 op)
        self.mma_inst_shape_mn = (self.mma_tiler[0], self.mma_tiler[1])
        self.mma_inst_shape_mn_sfb = (
            self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1),
            cute.round_up(self.mma_inst_shape_mn[1], 128),
        )

        tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
            self.a_dtype,
            self.a_major_mode,
            self.b_major_mode,
            self.sf_dtype,
            self.sf_vec_size,
            self.cta_group,
            self.mma_inst_shape_mn,
        )
        tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
            self.a_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,
        )

        # Compute mma/cluster/tile shapes
        mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
        mma_inst_tile_k = 4

        self.mma_tiler = (
            self.mma_inst_shape_mn[0],
            self.mma_inst_shape_mn[1],
            mma_inst_shape_k * mma_inst_tile_k,
        )
        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,
        )

        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.cta_tile_shape_mnk_sfb = (
            self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape),
            self.mma_tiler_sfb[1],
            self.mma_tiler_sfb[2],
        )

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

        # Multicast info
        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

        # Epilogue subtile
        self.epi_tile = sm100_utils.compute_epilogue_tile_shape(
            self.cta_tile_shape_mnk,
            self.use_2cta_instrs,
            self.c_layout,
            self.c_dtype,
        )
        self.epi_tile_n = cute.size(self.epi_tile[1])

        # Stages (dual-gemm: force 1 acc stage to fit TMEM)
        self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages_dual(
            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,
            self.ab_stages_param,
        )

        # Prefetch distance:
        #   -1 => auto (=num_ab_stage)
        #    0 => disabled
        #   >0 => explicit (clamped to <= num_ab_stage)
        if self.prefetch_dist_param < 0:
            self.prefetch_dist = self.num_ab_stage
        else:
            self.prefetch_dist = self.prefetch_dist_param
        if self.prefetch_dist > self.num_ab_stage:
            self.prefetch_dist = self.num_ab_stage
        if self.prefetch_dist < 0:
            self.prefetch_dist = 0
        self.prefetch_enabled = self.prefetch_dist > 0

        # Shared memory layouts
        self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
            tiled_mma,
            self.mma_tiler,
            self.a_dtype,
            self.num_ab_stage,
        )
        self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
            tiled_mma,
            self.mma_tiler,
            self.b_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,
        )

        # TMEM col accounting
        sf_atom_mn = 32
        self.num_sfa_tmem_cols = (self.cta_tile_shape_mnk[0] // sf_atom_mn) * 4
        self.num_sfb_tmem_cols = (self.cta_tile_shape_mnk_sfb[1] // sf_atom_mn) * 4

        # Accumulator cols per GEMM (include acc stage)
        self.num_accumulator_tmem_cols = self.cta_tile_shape_mnk[1] * self.num_acc_stage

        used_cols = (
            2 * self.num_accumulator_tmem_cols
            + self.num_sfa_tmem_cols
            + 2 * self.num_sfb_tmem_cols
        )
        if used_cols > self.num_tmem_alloc_cols:
            raise ValueError(
                f"TMEM overcommit: need {used_cols} cols but only have {self.num_tmem_alloc_cols}. "
                f"Try smaller mma_tiler_mn or fewer stages."
            )

    @cute.jit
    def __call__(
        self,
        a_tensor: cute.Tensor,
        b1_tensor: cute.Tensor,
        b2_tensor: cute.Tensor,
        sfa_tensor: cute.Tensor,
        sfb1_tensor: cute.Tensor,
        sfb2_tensor: cute.Tensor,
        c_tensor: cute.Tensor,
        max_active_clusters: int,
    ):
        # Infer dtypes/layouts
        self.a_dtype = a_tensor.element_type
        self.b_dtype = b1_tensor.element_type
        self.sf_dtype = sfa_tensor.element_type
        self.c_dtype = c_tensor.element_type

        # Layout major modes
        self.a_major_mode = utils.LayoutEnum.from_tensor(a_tensor).mma_major_mode()
        self.b_major_mode = utils.LayoutEnum.from_tensor(b1_tensor).mma_major_mode()
        self.c_layout = utils.LayoutEnum.from_tensor(c_tensor)

        # Require matching dtypes
        if cutlass.const_expr(self.a_dtype != self.b_dtype):
            raise TypeError(f"A/B type must match: {self.a_dtype} != {self.b_dtype}")

        # Attributes (tiler, stages, layouts, multicast)
        self._setup_attributes()

        # Re-wrap SFA/SFB tensors to the blockscaled SF atom layout
        sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(a_tensor.shape, self.sf_vec_size)
        sfa_tensor = cute.make_tensor(sfa_tensor.iterator, sfa_layout)

        sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(b1_tensor.shape, self.sf_vec_size)
        sfb1_tensor = cute.make_tensor(sfb1_tensor.iterator, sfb_layout)
        sfb2_tensor = cute.make_tensor(sfb2_tensor.iterator, sfb_layout)

        # Make tiled MMA
        tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
            self.a_dtype,
            self.a_major_mode,
            self.b_major_mode,
            self.sf_dtype,
            self.sf_vec_size,
            self.cta_group,
            self.mma_inst_shape_mn,
        )
        tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
            self.a_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,
        )
        atom_thr_size = cute.size(tiled_mma.thr_id.shape)

        # --- TMA atoms ---
        a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id)
        b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id)
        sfa_op = a_op
        sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, tiled_mma.thr_id)

        a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
        b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
        sfa_smem_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0))
        sfb_smem_layout = cute.slice_(self.sfb_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,
        )
        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,
        )

        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,
        )

        tma_atom_sfb1, tma_tensor_sfb1 = cute.nvgpu.make_tiled_tma_atom_B(
            sfb_op,
            sfb1_tensor,
            sfb_smem_layout,
            self.mma_tiler_sfb,
            tiled_mma_sfb,
            self.cluster_layout_sfb_vmnk.shape,
            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,
            tiled_mma_sfb,
            self.cluster_layout_sfb_vmnk.shape,
            internal_type=cutlass.Int16,
        )

        # Special handling for cta_tile_shape_n == 192 (SFB layout is stored in 128-col chunks)
        if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192):
            x = tma_tensor_sfb1.stride[0][1]
            y = cute.ceil_div(tma_tensor_sfb1.shape[0][1], 4)

            new_shape = (
                (tma_tensor_sfb1.shape[0][0], ((2, 2), y)),
                tma_tensor_sfb1.shape[1],
                tma_tensor_sfb1.shape[2],
            )
            # Use right multiplication for ScaledBasis (3 * x instead of x * 3)
            x_times_3 = 3 * x
            new_stride = (
                (tma_tensor_sfb1.stride[0][0], ((x, x), x_times_3)),
                tma_tensor_sfb1.stride[1],
                tma_tensor_sfb1.stride[2],
            )
            tma_tensor_sfb_new_layout = cute.make_layout(new_shape, stride=new_stride)
            tma_tensor_sfb1 = cute.make_tensor(
                tma_tensor_sfb1.iterator, tma_tensor_sfb_new_layout
            )
            tma_tensor_sfb2 = cute.make_tensor(
                tma_tensor_sfb2.iterator, tma_tensor_sfb_new_layout
            )

        # TMA load byte count (for barrier/pipeline)
        a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout)
        b_copy_size = cute.size_in_bytes(self.b_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

        # Setup TMA store for C
        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,
        )

        # Compute grid size (simple grid, non-persistent)
        # Calculate number of tiles (CTAs) needed, then account for cluster shape
        c_shape = cute.slice_(self.cta_tile_shape_mnk, (None, None, 0))
        gc = cute.zipped_divide(c_tensor, tiler=c_shape)
        num_ctas_mnl = gc[(0, (None, None, None))].shape
        # For non-persistent, grid = (num_clusters_m * cluster_m, num_clusters_n * cluster_n, l)
        num_clusters_m = cute.ceil_div(num_ctas_mnl[0], self.cluster_shape_mn[0])
        num_clusters_n = cute.ceil_div(num_ctas_mnl[1], self.cluster_shape_mn[1])
        grid = (
            num_clusters_m * self.cluster_shape_mn[0],
            num_clusters_n * self.cluster_shape_mn[1],
            num_ctas_mnl[2],
        )

        self.buffer_align_bytes = 1024

        # Define shared storage for kernel
        @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]
            acc_empty_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) -- B1
            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) -- B2
            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

        # Launch
        self.kernel(
            tiled_mma,
            tiled_mma_sfb,
            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,
            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),
            min_blocks_per_mp=1,
        )
        return

    # -------------------------------------------------------------------------
    # Device kernel
    # -------------------------------------------------------------------------

    @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_b1: cute.CopyAtom,
        mB_nkl1: cute.Tensor,
        tma_atom_b2: cute.CopyAtom,
        mB_nkl2: cute.Tensor,
        tma_atom_sfa: cute.CopyAtom,
        mSFA_mkl: cute.Tensor,
        tma_atom_sfb1: cute.CopyAtom,
        mSFB_nkl1: cute.Tensor,
        tma_atom_sfb2: cute.CopyAtom,
        mSFB_nkl2: cute.Tensor,
        tma_atom_c: cute.CopyAtom,
        mC_mnl: cute.Tensor,
        cluster_layout_vmnk: cute.Layout,
        cluster_layout_sfb_vmnk: cute.Layout,
        a_smem_layout_staged: cute.ComposedLayout,
        b_smem_layout_staged: cute.ComposedLayout,
        sfa_smem_layout_staged: cute.Layout,
        sfb_smem_layout_staged: cute.Layout,
        c_smem_layout_staged: cute.ComposedLayout | cute.Layout,
        epi_tile: cute.Tile,
    ):
        warp_idx = cute.arch.warp_idx()
        warp_idx = cute.arch.make_warp_uniform(warp_idx)

        # Prefetch TMA descriptors
        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)

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

        # Coords inside cluster
        bidx, bidy, bidz = cute.arch.block_idx()
        mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape)
        is_leader_cta = mma_tile_coord_v == 0
        cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
        block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
        block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord(cta_rank_in_cluster)

        tidx, tid_y, tid_z = cute.arch.thread_idx()

        # Shared memory alloc
        smem = utils.SmemAllocator()
        storage = smem.allocate(self.shared_storage)

        # Pipelines
        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.PipelineTmaUmma.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) * (2 if use_2cta_instrs else 1)
        acc_pipeline_consumer_group = pipeline.CooperativeGroup(
            pipeline.Agent.Thread,
            num_acc_consumer_threads,
        )
        acc_pipeline = pipeline.PipelineUmmaAsync.create(
            barrier_storage=storage.acc_full_mbar_ptr.data_ptr(),
            num_stages=self.num_acc_stage,
            producer_group=acc_pipeline_producer_group,
            consumer_group=acc_pipeline_consumer_group,
            cta_layout_vmnk=cluster_layout_vmnk,
            defer_sync=True,
        )

        # TMEM allocator
        tmem = utils.TmemAllocator(
            storage.tmem_holding_buf,
            barrier_for_retrieve=self.tmem_alloc_barrier,
            allocator_warp_id=self.epilog_warp_id[0],
            is_two_cta=use_2cta_instrs,
            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
        )

        # Cluster arrive
        pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)

        # Setup SMEM tensors
        sC = storage.sC.get_tensor(c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner)
        sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner)
        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)
        sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged)
        sSFB1 = storage.sSFB1.get_tensor(sfb_smem_layout_staged)
        sSFB2 = storage.sSFB2.get_tensor(sfb_smem_layout_staged)

        # Multicast masks
        a_full_mcast_mask = None
        b_full_mcast_mask = None
        sfa_full_mcast_mask = None
        sfb_full_mcast_mask = None
        if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs):
            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
        gA_mkl = cute.local_tile(mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None))
        gB_nkl1 = cute.local_tile(mB_nkl1, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None))
        gB_nkl2 = cute.local_tile(mB_nkl2, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None))
        gSFA_mkl = cute.local_tile(mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None))
        gSFB_nkl1 = cute.local_tile(
            mSFB_nkl1,
            cute.slice_(self.mma_tiler_sfb, (0, None, None)),
            (None, None, None),
        )
        gSFB_nkl2 = cute.local_tile(
            mSFB_nkl2,
            cute.slice_(self.mma_tiler_sfb, (0, None, None)),
            (None, None, None),
        )
        gC_mnl = cute.local_tile(mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None))

        # Partition global tensors for TiledMMA (needed for static TMA partitions)
        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)
        tCgB1 = thr_mma.partition_B(gB_nkl1)
        tCgB2 = thr_mma.partition_B(gB_nkl2)
        tCgSFA = thr_mma.partition_A(gSFA_mkl)
        tCgSFB1 = thr_mma_sfb.partition_B(gSFB_nkl1)
        tCgSFB2 = thr_mma_sfb.partition_B(gSFB_nkl2)
        tCgC = thr_mma.partition_C(gC_mnl)

        k_tile_cnt = cute.size(gA_mkl, mode=[3])

        # Partition global tensors for TMA
        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)
        tBsB1, tBgB1 = 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),
        )

        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)
        tBsSFB1, tBgSFB1 = 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 = 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 fragments
        tCrA = tiled_mma.make_fragment_A(sA)
        tCrB1 = tiled_mma.make_fragment_B(sB1)
        tCrB2 = tiled_mma.make_fragment_B(sB2)

        acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
        tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage))

        # Cluster wait before tmem alloc
        pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)

        # ---------------------------------------------------------------------
        # Specialized TMA warp (load A/B/SF)
        # ---------------------------------------------------------------------
        if warp_idx == self.tma_warp_id:
            # Simple grid scheduling (like shiyegao, no persistent scheduler)
            ab_producer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer, self.num_ab_stage
            )
            mma_tile_coord_mnl = (
                bidx // cute.size(tiled_mma.thr_id.shape),
                bidy,
                bidz,
            )
            tAgA_slice = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
            tBgB1_slice = tBgB1[(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])]
            tAgSFA_slice = tAgSFA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
            slice_n_sfb = mma_tile_coord_mnl[1]
            if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
                slice_n_sfb = mma_tile_coord_mnl[1] // 2
            tBgSFB1_slice = tBgSFB1[(None, slice_n_sfb, None, mma_tile_coord_mnl[2])]
            tBgSFB2_slice = tBgSFB2[(None, slice_n_sfb, None, mma_tile_coord_mnl[2])]
            # Prefetch: initial batch to prime the TMA pipeline
            if self.prefetch_enabled:
                for pf_k_tile in cutlass.range(
                    0, min(self.prefetch_dist, k_tile_cnt), unroll=1
                ):
                    cute.prefetch(tma_atom_a, tAgA_slice[(None, pf_k_tile)])
                    cute.prefetch(tma_atom_b1, tBgB1_slice[(None, pf_k_tile)])
                    cute.prefetch(tma_atom_b2, tBgB2_slice[(None, pf_k_tile)])
                    cute.prefetch(tma_atom_sfa, tAgSFA_slice[(None, pf_k_tile)])
                    cute.prefetch(tma_atom_sfb1, tBgSFB1_slice[(None, pf_k_tile)])
                    cute.prefetch(tma_atom_sfb2, tBgSFB2_slice[(None, pf_k_tile)])
            ab_producer_state.reset_count()
            peek_ab_empty_status = cutlass.Boolean(1)
            if ab_producer_state.count < k_tile_cnt:
                peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state)
            for k_tile_prefetch in cutlass.range(0, k_tile_cnt, 1, unroll=1):
                ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status)
                bar = ab_pipeline.producer_get_barrier(ab_producer_state)
                # TMA load A/B1/B2/SFA/SFB1/SFB2
                try:
                    cute.copy(
                        tma_atom_a,
                        tAgA_slice[(None, ab_producer_state.count)],
                        tAsA[(None, ab_producer_state.index)],
                        tma_bar_ptr=bar,
                        mcast_mask=a_full_mcast_mask,
                        cache_policy=cutlass.Int64(cutlass.Int64(_TMA_CACHE_EVICT_FIRST).ir_value()),
                    )
                except TypeError:
                    cute.copy(
                        tma_atom_a,
                        tAgA_slice[(None, ab_producer_state.count)],
                        tAsA[(None, ab_producer_state.index)],
                        tma_bar_ptr=bar,
                        mcast_mask=a_full_mcast_mask,
                    )
                try:
                    cute.copy(
                        tma_atom_b1,
                        tBgB1_slice[(None, ab_producer_state.count)],
                        tBsB1[(None, ab_producer_state.index)],
                        tma_bar_ptr=bar,
                        mcast_mask=b_full_mcast_mask,
                        cache_policy=cutlass.Int64(cutlass.Int64(_TMA_CACHE_EVICT_FIRST).ir_value()),
                    )
                except TypeError:
                    cute.copy(
                        tma_atom_b1,
                        tBgB1_slice[(None, ab_producer_state.count)],
                        tBsB1[(None, ab_producer_state.index)],
                        tma_bar_ptr=bar,
                        mcast_mask=b_full_mcast_mask,
                    )
                try:
                    cute.copy(
                        tma_atom_b2,
                        tBgB2_slice[(None, ab_producer_state.count)],
                        tBsB2[(None, ab_producer_state.index)],
                        tma_bar_ptr=bar,
                        mcast_mask=b_full_mcast_mask,
                        cache_policy=cutlass.Int64(cutlass.Int64(_TMA_CACHE_EVICT_FIRST).ir_value()),
                    )
                except TypeError:
                    cute.copy(
                        tma_atom_b2,
                        tBgB2_slice[(None, ab_producer_state.count)],
                        tBsB2[(None, ab_producer_state.index)],
                        tma_bar_ptr=bar,
                        mcast_mask=b_full_mcast_mask,
                    )
                try:
                    cute.copy(
                        tma_atom_sfa,
                        tAgSFA_slice[(None, ab_producer_state.count)],
                        tAsSFA[(None, ab_producer_state.index)],
                        tma_bar_ptr=bar,
                        mcast_mask=sfa_full_mcast_mask,
                        cache_policy=cutlass.Int64(cutlass.Int64(_TMA_CACHE_EVICT_FIRST).ir_value()),
                    )
                except TypeError:
                    cute.copy(
                        tma_atom_sfa,
                        tAgSFA_slice[(None, ab_producer_state.count)],
                        tAsSFA[(None, ab_producer_state.index)],
                        tma_bar_ptr=bar,
                        mcast_mask=sfa_full_mcast_mask,
                    )
                try:
                    cute.copy(
                        tma_atom_sfb1,
                        tBgSFB1_slice[(None, ab_producer_state.count)],
                        tBsSFB1[(None, ab_producer_state.index)],
                        tma_bar_ptr=bar,
                        mcast_mask=sfb_full_mcast_mask,
                        cache_policy=cutlass.Int64(cutlass.Int64(_TMA_CACHE_EVICT_FIRST).ir_value()),
                    )
                except TypeError:
                    cute.copy(
                        tma_atom_sfb1,
                        tBgSFB1_slice[(None, ab_producer_state.count)],
                        tBsSFB1[(None, ab_producer_state.index)],
                        tma_bar_ptr=bar,
                        mcast_mask=sfb_full_mcast_mask,
                    )
                try:
                    cute.copy(
                        tma_atom_sfb2,
                        tBgSFB2_slice[(None, ab_producer_state.count)],
                        tBsSFB2[(None, ab_producer_state.index)],
                        tma_bar_ptr=bar,
                        mcast_mask=sfb_full_mcast_mask,
                        cache_policy=cutlass.Int64(cutlass.Int64(_TMA_CACHE_EVICT_FIRST).ir_value()),
                    )
                except TypeError:
                    cute.copy(
                        tma_atom_sfb2,
                        tBgSFB2_slice[(None, ab_producer_state.count)],
                        tBsSFB2[(None, ab_producer_state.index)],
                        tma_bar_ptr=bar,
                        mcast_mask=sfb_full_mcast_mask,
                    )
                # Prefetch: rolling prefetch for next tiles
                if self.prefetch_enabled:
                    if k_tile_prefetch < k_tile_cnt - self.prefetch_dist:
                        future_k_tile = ab_producer_state.count + self.prefetch_dist
                        cute.prefetch(tma_atom_a, tAgA_slice[(None, future_k_tile)])
                        cute.prefetch(tma_atom_b1, tBgB1_slice[(None, future_k_tile)])
                        cute.prefetch(tma_atom_b2, tBgB2_slice[(None, future_k_tile)])
                        cute.prefetch(tma_atom_sfa, tAgSFA_slice[(None, future_k_tile)])
                        cute.prefetch(tma_atom_sfb1, tBgSFB1_slice[(None, future_k_tile)])
                        cute.prefetch(tma_atom_sfb2, tBgSFB2_slice[(None, future_k_tile)])
                ab_producer_state.advance()
                peek_ab_empty_status = cutlass.Boolean(1)
                if ab_producer_state.count < k_tile_cnt:
                    peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state)
            ab_pipeline.producer_tail(ab_producer_state)

        # ---------------------------------------------------------------------
        # Specialized MMA warp
        # ---------------------------------------------------------------------
        if warp_idx == self.mma_warp_id:
            tmem.wait_for_alloc()

            acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)

            # Acc1/Acc2 base (MMA, MMA_M, MMA_N, STAGE)
            tCtAcc1_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
            acc2_ptr = acc_tmem_ptr + self.num_accumulator_tmem_cols
            tCtAcc2_base = cute.make_tensor(acc2_ptr, tCtAcc_fake.layout)

            # SFA/SFB TMEM tensors
            sfa_tmem_ptr = cute.recast_ptr(
                acc_tmem_ptr + 2 * self.num_accumulator_tmem_cols,
                dtype=self.sf_dtype,
            )
            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)

            sfb1_tmem_ptr = cute.recast_ptr(
                acc_tmem_ptr + 2 * self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols,
                dtype=self.sf_dtype,
            )
            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(sfb1_tmem_ptr, tCtSFB_layout)

            sfb2_tmem_ptr = cute.recast_ptr(
                acc_tmem_ptr
                + 2 * self.num_accumulator_tmem_cols
                + self.num_sfa_tmem_cols
                + self.num_sfb_tmem_cols,
                dtype=self.sf_dtype,
            )
            tCtSFB2 = cute.make_tensor(sfb2_tmem_ptr, tCtSFB_layout)

            # S2T copy/partition
            (
                tiled_copy_s2t_sfa,
                tCsSFA_compact_s2t,
                tCtSFA_compact_s2t,
            ) = self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA)
            (
                tiled_copy_s2t_sfb,
                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)

            # Simple grid scheduling (like shiyegao, no persistent scheduler)
            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
            )
            mma_tile_coord_mnl = (
                bidx // cute.size(tiled_mma.thr_id.shape),
                bidy,
                bidz,
            )
            acc_stage_index = acc_producer_state.index
            tCtAcc1 = tCtAcc1_base[(None, None, None, acc_stage_index)]
            tCtAcc2 = tCtAcc2_base[(None, None, None, acc_stage_index)]
            ab_consumer_state.reset_count()
            peek_ab_full_status = cutlass.Boolean(1)
            if ab_consumer_state.count < k_tile_cnt and is_leader_cta:
                peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state)
            if is_leader_cta:
                acc_pipeline.producer_acquire(acc_producer_state)
            # Adjust SFB pointers for tiles where SFB is stored in 128-col chunks
            tCtSFB1_mma = tCtSFB1
            tCtSFB2_mma = tCtSFB2
            if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192):
                offset = cutlass.Int32(2) if mma_tile_coord_mnl[1] % 2 == 1 else cutlass.Int32(0)
                shifted_ptr1 = cute.recast_ptr(
                    acc_tmem_ptr + 2 * self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols + offset,
                    dtype=self.sf_dtype,
                )
                tCtSFB1_mma = cute.make_tensor(shifted_ptr1, tCtSFB_layout)
                shifted_ptr2 = cute.recast_ptr(
                    acc_tmem_ptr
                    + 2 * self.num_accumulator_tmem_cols
                    + self.num_sfa_tmem_cols
                    + self.num_sfb_tmem_cols
                    + offset,
                    dtype=self.sf_dtype,
                )
                tCtSFB2_mma = cute.make_tensor(shifted_ptr2, tCtSFB_layout)
            elif cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
                offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2)
                shifted_ptr1 = cute.recast_ptr(
                    acc_tmem_ptr + 2 * self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols + offset,
                    dtype=self.sf_dtype,
                )
                tCtSFB1_mma = cute.make_tensor(shifted_ptr1, tCtSFB_layout)
                shifted_ptr2 = cute.recast_ptr(
                    acc_tmem_ptr
                    + 2 * self.num_accumulator_tmem_cols
                    + self.num_sfa_tmem_cols
                    + self.num_sfb_tmem_cols
                    + offset,
                    dtype=self.sf_dtype,
                )
                tCtSFB2_mma = cute.make_tensor(shifted_ptr2, tCtSFB_layout)
            tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
            for k_tile in range(k_tile_cnt):
                if is_leader_cta:
                    ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status)
                    s2t_stage_coord = (None, None, None, None, ab_consumer_state.index)
                    cute.copy(
                        tiled_copy_s2t_sfa,
                        tCsSFA_compact_s2t[s2t_stage_coord],
                        tCtSFA_compact_s2t,
                    )
                    cute.copy(
                        tiled_copy_s2t_sfb,
                        tCsSFB1_compact_s2t[s2t_stage_coord],
                        tCtSFB1_compact_s2t,
                    )
                    cute.copy(
                        tiled_copy_s2t_sfb2,
                        tCsSFB2_compact_s2t[s2t_stage_coord],
                        tCtSFB2_compact_s2t,
                    )
                    num_kblocks = cute.size(tCrA, mode=[2])
                    for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):
                        kblock_coord = (None, None, kblock_idx, ab_consumer_state.index)
                        sf_kblock_coord = (None, None, kblock_idx)
                        tiled_mma.set(tcgen05.Field.SFA, tCtSFA[sf_kblock_coord].iterator)
                        tiled_mma.set(tcgen05.Field.SFB, tCtSFB1_mma[sf_kblock_coord].iterator)
                        cute.gemm(
                            tiled_mma,
                            tCtAcc1,
                            tCrA[kblock_coord],
                            tCrB1[kblock_coord],
                            tCtAcc1,
                        )
                        tiled_mma.set(tcgen05.Field.SFB, tCtSFB2_mma[sf_kblock_coord].iterator)
                        cute.gemm(
                            tiled_mma,
                            tCtAcc2,
                            tCrA[kblock_coord],
                            tCrB2[kblock_coord],
                            tCtAcc2,
                        )
                        tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
                    ab_pipeline.consumer_release(ab_consumer_state)
                ab_consumer_state.advance()
                peek_ab_full_status = cutlass.Boolean(1)
                if ab_consumer_state.count < k_tile_cnt:
                    if is_leader_cta:
                        peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state)
            if is_leader_cta:
                acc_pipeline.producer_commit(acc_producer_state)
            acc_producer_state.advance()
            acc_pipeline.producer_tail(acc_producer_state)

        # ---------------------------------------------------------------------
        # Specialized epilogue warps (also allocate TMEM)
        # ---------------------------------------------------------------------
        if warp_idx < self.mma_warp_id:
            tmem.allocate(self.num_tmem_alloc_cols)
            tmem.wait_for_alloc()

            acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
            tCtAcc1_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
            tCtAcc2_base = cute.make_tensor(
                acc_tmem_ptr + self.num_accumulator_tmem_cols,
                tCtAcc_fake.layout,
            )

            epi_tidx = tidx

            # Partition for tmem -> reg for both accumulators
            tiled_copy_t2r_1, tTR_tAcc1_base, tTR_rAcc1 = self.epilog_tmem_copy_and_partition(
                epi_tidx, tCtAcc1_base, tCgC, epi_tile, use_2cta_instrs
            )
            tiled_copy_t2r_2, tTR_tAcc2_base, tTR_rAcc2 = self.epilog_tmem_copy_and_partition(
                epi_tidx, tCtAcc2_base, tCgC, epi_tile, use_2cta_instrs
            )

            # Use one copy atom (they should be identical)
            tiled_copy_t2r = tiled_copy_t2r_1

            tTR_rC = cute.make_rmem_tensor(tTR_rAcc1.shape, self.c_dtype)
            tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition(
                tiled_copy_t2r,
                tTR_rC,
                epi_tidx,
                sC,
            )

            tma_atom_c, bSG_sC, bSG_gC_partitioned = self.epilog_gmem_copy_and_partition(
                epi_tidx,
                tma_atom_c,
                tCgC,
                epi_tile,
                sC,
            )

            # Simple grid scheduling (like shiyegao, no persistent scheduler)
            acc_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.num_acc_stage
            )
            c_producer_group = pipeline.CooperativeGroup(
                pipeline.Agent.Thread,
                32 * len(self.epilog_warp_id),
            )
            c_pipeline = pipeline.PipelineTmaStore.create(
                num_stages=self.num_c_stage,
                producer_group=c_producer_group,
            )
            mma_tile_coord_mnl = (
                bidx // cute.size(tiled_mma.thr_id.shape),
                bidy,
                bidz,
            )
            bSG_gC = bSG_gC_partitioned[(None, None, None, *mma_tile_coord_mnl)]
            acc_stage_index = acc_consumer_state.index
            tTR_tAcc1 = tTR_tAcc1_base[(None, None, None, None, None, acc_stage_index)]
            tTR_tAcc2 = tTR_tAcc2_base[(None, None, None, None, None, acc_stage_index)]
            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])
            num_prev_subtiles = cutlass.Int32(0)
            for subtile_idx in cutlass.range(subtile_cnt):
                real_subtile_idx = subtile_idx
                # Load both accumulators
                tTR_tAcc1_mn = tTR_tAcc1[(None, None, None, real_subtile_idx)]
                tTR_tAcc2_mn = tTR_tAcc2[(None, None, None, real_subtile_idx)]
                cute.copy(tiled_copy_t2r, tTR_tAcc1_mn, tTR_rAcc1)
                cute.copy(tiled_copy_t2r, tTR_tAcc2_mn, tTR_rAcc2)
                # Fused epilogue: silu(acc1) * acc2
                x = tiled_copy_r2s.retile(tTR_rAcc1).load()
                y = tiled_copy_r2s.retile(tTR_rAcc2).load()
                out = _silu(x) * y
                tRS_rC.store(out.to(self.c_dtype))
                # Store to SMEM
                c_buffer = (num_prev_subtiles + real_subtile_idx) % self.num_c_stage
                cute.copy(
                    tiled_copy_r2s,
                    tRS_rC,
                    tRS_sC[(None, None, None, c_buffer)],
                )
                cute.arch.fence_proxy(
                    cute.arch.ProxyKind.async_shared,
                    space=cute.arch.SharedSpace.shared_cta,
                )
                self.epilog_sync_barrier.arrive_and_wait()
                # TMA store to GMEM (one warp)
                if warp_idx == self.epilog_warp_id[0]:
                    try:
                        cute.copy(
                            tma_atom_c,
                            bSG_sC[(None, c_buffer)],
                            bSG_gC[(None, real_subtile_idx)],
                            cache_policy=cutlass.Int64(cutlass.Int64(_TMA_CACHE_EVICT_FIRST).ir_value()),
                        )
                    except TypeError:
                        cute.copy(
                            tma_atom_c,
                            bSG_sC[(None, c_buffer)],
                            bSG_gC[(None, real_subtile_idx)],
                        )
                    c_pipeline.producer_commit()
                    c_pipeline.producer_acquire()
                self.epilog_sync_barrier.arrive_and_wait()
            # Release accumulator buffer
            with cute.arch.elect_one():
                acc_pipeline.consumer_release(acc_consumer_state)
            acc_consumer_state.advance()

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

        return

    # -------------------------------------------------------------------------
    # Helper functions (mostly reused from CUTLASS example)
    # -------------------------------------------------------------------------

    def mainloop_s2t_copy_and_partition(
        self,
        sSF: cute.Tensor,
        tSF: cute.Tensor,
    ):
        tCsSF_compact = cute.filter_zeros(sSF)
        tCtSF_compact = cute.filter_zeros(tSF)

        copy_atom_s2t = cute.make_copy_atom(tcgen05.Cp4x32x128bOp(self.cta_group), 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,
        gC_mnl: cute.Tensor,
        epi_tile: cute.Tile,
        use_2cta_instrs: cutlass.Boolean | bool,
    ):
        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,
            use_2cta_instrs,
        )
        tAcc_epi = cute.flat_divide(tAcc[((None, None), 0, 0, None)], epi_tile)
        tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)])

        thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
        tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi)

        gC_mnl_epi = cute.flat_divide(
            gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile
        )
        tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi)

        tTR_rAcc = cute.make_rmem_tensor(
            tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape,
            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,
    ):
        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,
        tidx: cutlass.Int32,
        atom: cute.CopyAtom | cute.TiledCopy,
        gC_mnl: cute.Tensor,
        epi_tile: cute.Tile,
        sC: cute.Tensor,
    ):
        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)

        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_dual(
        tiled_mma: cute.TiledMma,
        mma_tiler_mnk: tuple[int, int, int],
        a_dtype,
        b_dtype,
        epi_tile: cute.Tile,
        c_dtype,
        c_layout: utils.LayoutEnum,
        sf_dtype,
        sf_vec_size: int,
        smem_capacity: int,
        occupancy: int,
        ab_stages: int = 0,
    ):
        # Dual-GEMM: accumulator stage count depends on N tile to fit TMEM.
        # For N=64 we can afford 2 stages; for N>=128 we keep 1 stage.
        num_acc_stage = 2 if mma_tiler_mnk[1] == 64 else 1
        num_c_stage = 2

        a_smem_layout_stage_one = sm100_utils.make_smem_layout_a(
            tiled_mma,
            mma_tiler_mnk,
            a_dtype,
            1,
        )
        b_smem_layout_stage_one = sm100_utils.make_smem_layout_b(
            tiled_mma,
            mma_tiler_mnk,
            b_dtype,
            1,
        )
        sfa_smem_layout_stage_one = blockscaled_utils.make_smem_layout_sfa(
            tiled_mma,
            mma_tiler_mnk,
            sf_vec_size,
            1,
        )
        sfb_smem_layout_stage_one = blockscaled_utils.make_smem_layout_sfb(
            tiled_mma,
            mma_tiler_mnk,
            sf_vec_size,
            1,
        )
        c_smem_layout_stage_one = sm100_utils.make_smem_layout_epi(
            c_dtype,
            c_layout,
            epi_tile,
            1,
        )

        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_stage_one)
            + cute.size_in_bytes(sf_dtype, sfa_smem_layout_stage_one)
            + 2 * cute.size_in_bytes(sf_dtype, sfb_smem_layout_stage_one)
        )
        mbar_helpers_bytes = 1024
        c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_stage_one)
        c_bytes = c_bytes_per_stage * num_c_stage

        max_ab_stage = (
            smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes)
        ) // ab_bytes_per_stage
        if max_ab_stage < 1:
            max_ab_stage = 1

        num_ab_stage = max_ab_stage
        if ab_stages > 0:
            num_ab_stage = min(ab_stages, max_ab_stage)

        num_c_stage += (
            smem_capacity
            - occupancy * ab_bytes_per_stage * num_ab_stage
            - occupancy * (mbar_helpers_bytes + c_bytes)
        ) // (occupancy * c_bytes_per_stage)
        if num_c_stage < 1:
            num_c_stage = 1

        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],
        max_active_clusters: int,
    ):
        c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0))
        gc = cute.zipped_divide(c, tiler=c_shape)
        num_ctas_mnl = gc[(0, (None, None, None))].shape
        cluster_shape_mnl = (*cluster_shape_mn, 1)

        tile_sched_params = utils.PersistentTileSchedulerParams(num_ctas_mnl, cluster_shape_mnl)
        grid = utils.StaticPersistentTileScheduler.get_grid_shape(tile_sched_params, max_active_clusters)
        return tile_sched_params, grid


# -----------------------------------------------------------------------------
# Compilation / dispatch
# -----------------------------------------------------------------------------


_COMPILED: dict[tuple[int, int, int, int, int, int], callable] = {}


def _get_max_active_clusters(cluster_shape_mn: tuple[int, int]) -> int:
    cluster_size = cluster_shape_mn[0] * cluster_shape_mn[1]
    hardware_info = cutlass.utils.HardwareInfo()
    return hardware_info.get_max_active_clusters(cluster_size)


def _compile_for_config(
    mma_m: int,
    mma_n: int,
    cluster_m: int,
    cluster_n: int,
    ab_stages: int,
    prefetch_dist: int,
):
    key = (mma_m, mma_n, cluster_m, cluster_n, ab_stages, prefetch_dist)
    if key in _COMPILED:
        return _COMPILED[key]

    gemm = Sm100BlockScaledPersistentDualGemmKernel(
        sf_vec_size=SF_VEC_SIZE,
        mma_tiler_mn=(mma_m, mma_n),
        cluster_shape_mn=(cluster_m, cluster_n),
        ab_stages=ab_stages,
        prefetch_dist=prefetch_dist,
    )

    max_active_clusters = _get_max_active_clusters((cluster_m, cluster_n))

    @cute.jit
    def my_kernel(
        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,
        problem_size: tuple,
    ):
        m, n, k, l = problem_size

        a_tensor = cute.make_tensor(
            a_ptr,
            cute.make_layout(
                (m, cute.assume(k, 32), l),
                stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32)),
            ),
        )
        b1_tensor = cute.make_tensor(
            b1_ptr,
            cute.make_layout(
                (n, cute.assume(k, 32), l),
                stride=(cute.assume(k, 32), 1, cute.assume(n * k, 32)),
            ),
        )
        b2_tensor = cute.make_tensor(
            b2_ptr,
            cute.make_layout(
                (n, cute.assume(k, 32), l),
                stride=(cute.assume(k, 32), 1, cute.assume(n * k, 32)),
            ),
        )
        c_tensor = cute.make_tensor(
            c_ptr,
            cute.make_layout(
                (cute.assume(m, 32), n, l),
                stride=(n, 1, m * n),
            ),
        )

        # Scale factor tensors are reinterpreted internally into the blockscaled SF layout.
        sfa_tensor = cute.make_tensor(sfa_ptr, cute.make_layout(1))
        sfb1_tensor = cute.make_tensor(sfb1_ptr, cute.make_layout(1))
        sfb2_tensor = cute.make_tensor(sfb2_ptr, cute.make_layout(1))

        gemm(
            a_tensor,
            b1_tensor,
            b2_tensor,
            sfa_tensor,
            sfb1_tensor,
            sfb2_tensor,
            c_tensor,
            max_active_clusters,
        )
        return

    # Compile once (use dummy pointers)
    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)

    compiled = cute.compile(
        my_kernel,
        a_ptr,
        b1_ptr,
        b2_ptr,
        sfa_ptr,
        sfb1_ptr,
        sfb2_ptr,
        c_ptr,
        (0, 0, 0, 0),
        options="--opt-level 3",
    )

    _COMPILED[key] = compiled
    return compiled

# -----------------------------------------------------------------------------
# Config selection / tuning
# -----------------------------------------------------------------------------

# Config tuple:
#   (mma_m, mma_n, cluster_m, cluster_n, ab_stages, prefetch_dist)
# where:
#   - ab_stages: 0 = auto; >0 = override (clamped to max that fits)
#   - prefetch_dist: -1 = auto (=num_ab_stage); 0 = disable;
#                    >0 = explicit (clamped to <= num_ab_stage)
_NVFP4_DUAL_GEMM_FORCE_CONFIG_ENV = "NVFP4_DUAL_GEMM_FORCE_CONFIG"


def _parse_force_config() -> tuple[int, int, int, int, int, int] | None:
    cfg = os.environ.get(_NVFP4_DUAL_GEMM_FORCE_CONFIG_ENV)
    if not cfg:
        return None

    parts = [p.strip() for p in cfg.split(",") if p.strip()]
    try:
        values = [int(p) for p in parts]
    except ValueError as e:
        raise ValueError(
            f"{_NVFP4_DUAL_GEMM_FORCE_CONFIG_ENV} must be comma-separated ints, got: {cfg!r}"
        ) from e

    if len(values) == 4:
        mma_m, mma_n, cluster_m, cluster_n = values
        return (mma_m, mma_n, cluster_m, cluster_n, 0, -1)

    if len(values) == 6:
        mma_m, mma_n, cluster_m, cluster_n, ab_stages, prefetch_dist = values
        return (mma_m, mma_n, cluster_m, cluster_n, ab_stages, prefetch_dist)

    raise ValueError(
        f"{_NVFP4_DUAL_GEMM_FORCE_CONFIG_ENV} must have 4 or 6 ints, got {len(values)}: {cfg!r}"
    )


# Tuned configs for the public benchmark shapes (m, n, k).
# Update via local search/profiling on Blackwell.
_TUNED_CONFIGS: dict[tuple[int, int, int], tuple[int, int, int, int, int, int]] = {
    (256, 4096, 7168): (256, 64, 2, 1, 0, -1),
    (512, 4096, 7168): (256, 128, 4, 1, 0, -1),
    (256, 3072, 4096): (256, 64, 2, 1, 0, -1),
    (512, 3072, 7168): (256, 128, 4, 1, 0, -1),
}


def _select_config(m: int, n: int, k: int) -> tuple[int, int, int, int, int, int]:
    forced = _parse_force_config()
    if forced is not None:
        mma_m, mma_n, cluster_m, cluster_n, ab_stages, prefetch_dist = forced
        if (m % mma_m != 0) or (n % mma_n != 0):
            raise ValueError(
                f"Forced config {forced} is incompatible with problem (m={m}, n={n}, k={k})"
            )
        return forced

    tuned = _TUNED_CONFIGS.get((m, n, k))
    if tuned is not None:
        return tuned

    # Fallback heuristic (non-benchmark shapes).
    mma_m = 256 if (m % 256 == 0) else 128
    mma_n = 128 if (n % 128 == 0 and n >= 3072) else 64
    cluster_m = 2 if mma_m == 256 else 1
    cluster_n = 1
    ab_stages = 0
    prefetch_dist = -1
    return (mma_m, mma_n, cluster_m, cluster_n, ab_stages, prefetch_dist)


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

    # Torch uses float4_e2m1fn_x2, so logical K is doubled.
    m = c.shape[0]
    n = c.shape[1]
    l = c.shape[2]
    k = a.shape[1] * 2

    mma_m, mma_n, cluster_m, cluster_n, ab_stages, prefetch_dist = _select_config(m, n, k)
    compiled = _compile_for_config(mma_m, mma_n, cluster_m, cluster_n, ab_stages, prefetch_dist)

    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