submission 445235

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

import functools
from typing import Tuple, List, Dict, Any

import torch
from task import input_t, output_t

# -----------------------------------------------------------------------------
# Tunables
# -----------------------------------------------------------------------------
bytes_per_tensormap = 128
num_tensormaps = 4

# NVFP4 UMMA op requires M-mode=128 (cannot be 64)
mma_tiler_mnk = (128, 128, 256)
mma_inst_shape_k = 64

ab_dtype = cutlass.Float4E2M1FN
sf_dtype = cutlass.Float8E4M3FN
c_dtype = cutlass.Float16

sf_vec_size = 16

threads_per_cta = 128

num_acc_stage = 1
num_ab_stage = 2  # try 2/3/4 later

# safe (can tune later)
num_tmem_alloc_cols = 512


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


# -----------------------------------------------------------------------------
# 1) Init tensormaps kernel (one CTA per group)
#    Writes tensormap descriptors into tensormaps[group, 0..3, :]
# -----------------------------------------------------------------------------
@cute.kernel
def init_tensormaps_kernel(
    tma_atom_a: cute.CopyAtom,
    tma_atom_b: cute.CopyAtom,
    tma_atom_sfa: cute.CopyAtom,
    tma_atom_sfb: cute.CopyAtom,
    tensor_of_abc_ptrs: cute.Tensor,        # [G,3] int64
    tensor_of_sfasfb_ptrs: cute.Tensor,     # [G,2] int64
    tensor_of_problem_sizes: cute.Tensor,   # [G,4] int32
    tensormaps: cute.Tensor,                # [G,4,16] int64
):
    warp_idx = cute.arch.warp_idx()
    warp_idx = cute.arch.make_warp_uniform(warp_idx)

    _, _, bidz = cute.arch.block_idx()
    group_idx = bidz

    m = tensor_of_problem_sizes[group_idx, 0]
    n = tensor_of_problem_sizes[group_idx, 1]
    k = tensor_of_problem_sizes[group_idx, 2]
    l = cutlass.Int32(1)

    # Shared buffer for building descriptors
    size_tensormap_in_i64 = num_tensormaps * bytes_per_tensormap // 8

    @cute.struct
    class SharedStorage:
        tensormap_buffer: cute.struct.MemRange[cutlass.Int64, size_tensormap_in_i64]

    smem = utils.SmemAllocator()
    storage = smem.allocate(SharedStorage)

    tensormap_smem_ptr = storage.tensormap_buffer.data_ptr()
    tensormap_a_smem_ptr = tensormap_smem_ptr
    tensormap_b_smem_ptr = tensormap_a_smem_ptr + bytes_per_tensormap // 8
    tensormap_sfa_smem_ptr = tensormap_b_smem_ptr + bytes_per_tensormap // 8
    tensormap_sfb_smem_ptr = tensormap_sfa_smem_ptr + bytes_per_tensormap // 8

    # Target gmem descriptor locations
    tensormap_manager = utils.TensorMapManager(utils.TensorMapUpdateMode.SMEM, 128)

    tensormap_a_gmem_ptr = tensormap_manager.get_tensormap_ptr(
        tensormaps[(group_idx, 0, None)].iterator
    )
    tensormap_b_gmem_ptr = tensormap_manager.get_tensormap_ptr(
        tensormaps[(group_idx, 1, None)].iterator
    )
    tensormap_sfa_gmem_ptr = tensormap_manager.get_tensormap_ptr(
        tensormaps[(group_idx, 2, None)].iterator
    )
    tensormap_sfb_gmem_ptr = tensormap_manager.get_tensormap_ptr(
        tensormaps[(group_idx, 3, None)].iterator
    )

    # Real pointers
    mA_mkl_iter = cute.make_ptr(
        ab_dtype, tensor_of_abc_ptrs[group_idx, 0], cute.AddressSpace.gmem
    ).align(32)
    mB_nkl_iter = cute.make_ptr(
        ab_dtype, tensor_of_abc_ptrs[group_idx, 1], cute.AddressSpace.gmem
    ).align(32)
    sfa_mkl_iter = cute.make_ptr(
        sf_dtype, tensor_of_sfasfb_ptrs[group_idx, 0], cute.AddressSpace.gmem
    ).align(32)
    sfb_nkl_iter = cute.make_ptr(
        sf_dtype, tensor_of_sfasfb_ptrs[group_idx, 1], cute.AddressSpace.gmem
    ).align(32)

    # Layouts
    mA_mkl_layout = cute.make_layout(
        (m, k, l), stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32))
    )
    mB_nkl_layout = cute.make_layout(
        (n, k, l), stride=(cute.assume(k, 32), 1, cute.assume(n * k, 32))
    )

    # SFA/SFB special layout (cublas doc)
    atom_shape = ((32, 4), (sf_vec_size, 4))
    atom_stride = ((16, 4), (0, 1))
    sfa_layout = cute.tile_to_shape(
        cute.make_layout(atom_shape, stride=atom_stride),
        mA_mkl_layout.shape,
        (2, 1, 3),
    )
    sfb_layout = cute.tile_to_shape(
        cute.make_layout(atom_shape, stride=atom_stride),
        mB_nkl_layout.shape,
        (2, 1, 3),
    )

    real_tensor_a = cute.make_tensor(mA_mkl_iter, mA_mkl_layout)
    real_tensor_b = cute.make_tensor(mB_nkl_iter, mB_nkl_layout)
    real_tensor_sfa = cute.make_tensor(sfa_mkl_iter, sfa_layout)
    real_tensor_sfb = cute.make_tensor(sfb_nkl_iter, sfb_layout)

    # Only warp0 builds & writes descriptors
    if warp_idx == 0:
        tensormap_manager.init_tensormap_from_atom(tma_atom_a, tensormap_a_smem_ptr, 0)
        tensormap_manager.init_tensormap_from_atom(tma_atom_b, tensormap_b_smem_ptr, 0)
        tensormap_manager.init_tensormap_from_atom(
            tma_atom_sfa, tensormap_sfa_smem_ptr, 0
        )
        tensormap_manager.init_tensormap_from_atom(
            tma_atom_sfb, tensormap_sfb_smem_ptr, 0
        )

        tensormap_manager.update_tensormap(
            (real_tensor_a, real_tensor_b, real_tensor_sfa, real_tensor_sfb),
            (tma_atom_a, tma_atom_b, tma_atom_sfa, tma_atom_sfb),
            (
                tensormap_a_gmem_ptr,
                tensormap_b_gmem_ptr,
                tensormap_sfa_gmem_ptr,
                tensormap_sfb_gmem_ptr,
            ),
            0,
            (
                tensormap_a_smem_ptr,
                tensormap_b_smem_ptr,
                tensormap_sfa_smem_ptr,
                tensormap_sfb_smem_ptr,
            ),
        )

        tensormap_manager.fence_tensormap_update(tensormap_a_gmem_ptr)
        tensormap_manager.fence_tensormap_update(tensormap_b_gmem_ptr)
        tensormap_manager.fence_tensormap_update(tensormap_sfa_gmem_ptr)
        tensormap_manager.fence_tensormap_update(tensormap_sfb_gmem_ptr)

    pass


# -----------------------------------------------------------------------------
# 2) Compute kernel (NO tensormap update inside CTA)
# -----------------------------------------------------------------------------
@cute.kernel
def gemm_kernel(
    tiled_mma: cute.TiledMma,
    tma_atom_a: cute.CopyAtom,
    mA_mkl: cute.Tensor,  # proxy
    tma_atom_b: cute.CopyAtom,
    mB_nkl: cute.Tensor,  # proxy
    tma_atom_sfa: cute.CopyAtom,
    mSFA_mkl: cute.Tensor,  # proxy
    tma_atom_sfb: cute.CopyAtom,
    mSFB_nkl: cute.Tensor,  # proxy
    tensor_of_abc_ptrs: cute.Tensor,
    tensor_of_sfasfb_ptrs: cute.Tensor,
    tensormaps: cute.Tensor,              # [G,4,16] pre-initialized
    tensor_of_problem_sizes: cute.Tensor,
    tensor_of_cluster_meta: cute.Tensor,  # [total_clusters,3] -> (g, tx, ty)
    a_smem_layout_staged: cute.ComposedLayout,
    b_smem_layout_staged: cute.ComposedLayout,
    sfa_smem_layout_staged: cute.Layout,
    sfb_smem_layout_staged: cute.Layout,
    num_tma_load_bytes: cutlass.Constexpr[int],
):
    warp_idx = cute.arch.warp_idx()
    warp_idx = cute.arch.make_warp_uniform(warp_idx)
    tidx, _, _ = cute.arch.thread_idx()

    _, _, bidz = cute.arch.block_idx()

    group_idx = tensor_of_cluster_meta[bidz, 0]
    coord_x = tensor_of_cluster_meta[bidz, 1]
    coord_y = tensor_of_cluster_meta[bidz, 2]

    m = tensor_of_problem_sizes[group_idx, 0]
    n = tensor_of_problem_sizes[group_idx, 1]
    k = tensor_of_problem_sizes[group_idx, 2]
    l = cutlass.Int32(1)

    # C tensor
    mC_mnl_iter = cute.make_ptr(
        c_dtype, tensor_of_abc_ptrs[group_idx, 2], cute.AddressSpace.gmem
    ).align(32)
    mC_mnl_layout = cute.make_layout(
        (m, n, l),
        stride=(cute.assume(n, 32), 1, cute.assume(m * n, 32)),
    )
    mC_mnl = cute.make_tensor(mC_mnl_iter, mC_mnl_layout)

    gC_mnl = cute.local_tile(
        mC_mnl, cute.slice_(mma_tiler_mnk, (None, None, 0)), (coord_x, coord_y, 0)
    )

    # Shared storage (no tensormap buffer now)
    @cute.struct
    class SharedStorage:
        ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, num_ab_stage * 2]
        acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, num_acc_stage * 2]
        tmem_holding_buf: cutlass.Int32

    smem = utils.SmemAllocator()
    storage = smem.allocate(SharedStorage)

    # SMEM tensors
    sA = smem.allocate_tensor(
        element_type=ab_dtype,
        layout=a_smem_layout_staged.outer,
        byte_alignment=128,
        swizzle=a_smem_layout_staged.inner,
    )
    sB = smem.allocate_tensor(
        element_type=ab_dtype,
        layout=b_smem_layout_staged.outer,
        byte_alignment=128,
        swizzle=b_smem_layout_staged.inner,
    )
    sSFA = smem.allocate_tensor(
        element_type=sf_dtype,
        layout=sfa_smem_layout_staged,
        byte_alignment=128,
    )
    sSFB = smem.allocate_tensor(
        element_type=sf_dtype,
        layout=sfb_smem_layout_staged,
        byte_alignment=128,
    )

    # Pipelines
    ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
    ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1)
    ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create(
        barrier_storage=storage.ab_mbar_ptr.data_ptr(),
        num_stages=num_ab_stage,
        producer_group=ab_pipeline_producer_group,
        consumer_group=ab_pipeline_consumer_group,
        tx_count=num_tma_load_bytes,
    ).make_participants()

    acc_producer, acc_consumer = pipeline.PipelineUmmaAsync.create(
        barrier_storage=storage.acc_mbar_ptr.data_ptr(),
        num_stages=num_acc_stage,
        producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
        consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, threads_per_cta),
    ).make_participants()

    # Proxy partitioning
    gA_mkl = cute.local_tile(
        mA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None)
    )
    gB_nkl = cute.local_tile(
        mB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None)
    )
    gSFA_mkl = cute.local_tile(
        mSFA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None)
    )
    gSFB_nkl = cute.local_tile(
        mSFB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None)
    )

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

    # Read prebuilt tensormaps for this group
    tensormap_manager = utils.TensorMapManager(utils.TensorMapUpdateMode.SMEM, 128)

    tensormap_a_gmem_ptr = tensormap_manager.get_tensormap_ptr(
        tensormaps[(group_idx, 0, None)].iterator
    )
    tensormap_b_gmem_ptr = tensormap_manager.get_tensormap_ptr(
        tensormaps[(group_idx, 1, None)].iterator
    )
    tensormap_sfa_gmem_ptr = tensormap_manager.get_tensormap_ptr(
        tensormaps[(group_idx, 2, None)].iterator
    )
    tensormap_sfb_gmem_ptr = tensormap_manager.get_tensormap_ptr(
        tensormaps[(group_idx, 3, None)].iterator
    )

    tma_desc_a = tensormap_manager.get_tensormap_ptr(
        tensormap_a_gmem_ptr, cute.AddressSpace.generic
    )
    tma_desc_b = tensormap_manager.get_tensormap_ptr(
        tensormap_b_gmem_ptr, cute.AddressSpace.generic
    )
    tma_desc_sfa = tensormap_manager.get_tensormap_ptr(
        tensormap_sfa_gmem_ptr, cute.AddressSpace.generic
    )
    tma_desc_sfb = tensormap_manager.get_tensormap_ptr(
        tensormap_sfb_gmem_ptr, cute.AddressSpace.generic
    )

    # TMA partition
    tAsA, tAgA = cpasync.tma_partition(
        tma_atom_a,
        0,
        cute.make_layout(1),
        cute.group_modes(sA, 0, 3),
        cute.group_modes(tCgA, 0, 3),
    )
    tBsB, tBgB = cpasync.tma_partition(
        tma_atom_b,
        0,
        cute.make_layout(1),
        cute.group_modes(sB, 0, 3),
        cute.group_modes(tCgB, 0, 3),
    )
    tAsSFA, tAgSFA = cpasync.tma_partition(
        tma_atom_sfa,
        0,
        cute.make_layout(1),
        cute.group_modes(sSFA, 0, 3),
        cute.group_modes(tCgSFA, 0, 3),
    )
    tAsSFA = cute.filter_zeros(tAsSFA)
    tAgSFA = cute.filter_zeros(tAgSFA)

    tBsSFB, tBgSFB = cpasync.tma_partition(
        tma_atom_sfb,
        0,
        cute.make_layout(1),
        cute.group_modes(sSFB, 0, 3),
        cute.group_modes(tCgSFB, 0, 3),
    )
    tBsSFB = cute.filter_zeros(tBsSFB)
    tBgSFB = cute.filter_zeros(tBgSFB)

    # MMA fragments / TMEM
    tCrA = tiled_mma.make_fragment_A(sA)
    tCrB = tiled_mma.make_fragment_B(sB)

    acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2])
    tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape)

    tmem_alloc_barrier = pipeline.NamedBarrier(barrier_id=1, num_threads=threads_per_cta)
    tmem = utils.TmemAllocator(storage.tmem_holding_buf, barrier_for_retrieve=tmem_alloc_barrier)
    tmem.allocate(num_tmem_alloc_cols)
    tmem.wait_for_alloc()
    acc_tmem_ptr = tmem.retrieve_ptr(cutlass.Float32)
    tCtAcc = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)

    # SFA/SFB in TMEM
    sfa_tmem_ptr = cute.recast_ptr(
        acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc),
        dtype=sf_dtype,
    )
    tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
        tiled_mma,
        mma_tiler_mnk,
        sf_vec_size,
        cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)),
    )
    tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout)

    sfb_tmem_ptr = cute.recast_ptr(
        acc_tmem_ptr
        + tcgen05.find_tmem_tensor_col_offset(tCtAcc)
        + tcgen05.find_tmem_tensor_col_offset(tCtSFA),
        dtype=sf_dtype,
    )
    tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
        tiled_mma,
        mma_tiler_mnk,
        sf_vec_size,
        cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)),
    )
    tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout)

    # S2T copy for SFA/SFB (keep original pattern)
    copy_atom_s2t = cute.make_copy_atom(
        tcgen05.Cp4x32x128bOp(tcgen05.CtaGroup.ONE),
        sf_dtype,
    )

    tCsSFA_compact = cute.filter_zeros(sSFA)
    tCtSFA_compact = cute.filter_zeros(tCtSFA)
    tiled_copy_s2t_sfa = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFA_compact)
    thr_copy_s2t_sfa = tiled_copy_s2t_sfa.get_slice(0)
    tCsSFA_compact_s2t_ = thr_copy_s2t_sfa.partition_S(tCsSFA_compact)
    tCsSFA_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy_s2t_sfa, tCsSFA_compact_s2t_)
    tCtSFA_compact_s2t = thr_copy_s2t_sfa.partition_D(tCtSFA_compact)

    tCsSFB_compact = cute.filter_zeros(sSFB)
    tCtSFB_compact = cute.filter_zeros(tCtSFB)
    tiled_copy_s2t_sfb = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFB_compact)
    thr_copy_s2t_sfb = tiled_copy_s2t_sfb.get_slice(0)
    tCsSFB_compact_s2t_ = thr_copy_s2t_sfb.partition_S(tCsSFB_compact)
    tCsSFB_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy_s2t_sfb, tCsSFB_compact_s2t_)
    tCtSFB_compact_s2t = thr_copy_s2t_sfb.partition_D(tCtSFB_compact)

    # k_tile_cnt
    k_tile_cnt = k // cutlass.Int32(mma_tiler_mnk[2])

    # Slice tile coords
    tAgA = tAgA[(None, coord_x, None, 0)]
    tBgB = tBgB[(None, coord_y, None, 0)]
    tAgSFA = tAgSFA[(None, coord_x, None, 0)]
    tBgSFB = tBgSFB[(None, coord_y, None, 0)]

    # Main loop
    if warp_idx == 0:
        acc_empty = acc_producer.acquire_and_advance()
        tiled_mma.set(tcgen05.Field.ACCUMULATE, False)

        # prime tile0
        ab_empty0 = ab_producer.acquire_and_advance()
        cute.copy(
            tma_atom_a,
            tAgA[(None, 0)],
            tAsA[(None, ab_empty0.index)],
            tma_bar_ptr=ab_empty0.barrier,
            tma_desc_ptr=tma_desc_a,
        )
        cute.copy(
            tma_atom_b,
            tBgB[(None, 0)],
            tBsB[(None, ab_empty0.index)],
            tma_bar_ptr=ab_empty0.barrier,
            tma_desc_ptr=tma_desc_b,
        )
        cute.copy(
            tma_atom_sfa,
            tAgSFA[(None, 0)],
            tAsSFA[(None, ab_empty0.index)],
            tma_bar_ptr=ab_empty0.barrier,
            tma_desc_ptr=tma_desc_sfa,
        )
        cute.copy(
            tma_atom_sfb,
            tBgSFB[(None, 0)],
            tBsSFB[(None, ab_empty0.index)],
            tma_bar_ptr=ab_empty0.barrier,
            tma_desc_ptr=tma_desc_sfb,
        )

        for k_tile in range(k_tile_cnt):
            # prefetch next
            if k_tile + 1 < k_tile_cnt:
                ab_empty = ab_producer.acquire_and_advance()
                kt = k_tile + 1
                cute.copy(
                    tma_atom_a,
                    tAgA[(None, kt)],
                    tAsA[(None, ab_empty.index)],
                    tma_bar_ptr=ab_empty.barrier,
                    tma_desc_ptr=tma_desc_a,
                )
                cute.copy(
                    tma_atom_b,
                    tBgB[(None, kt)],
                    tBsB[(None, ab_empty.index)],
                    tma_bar_ptr=ab_empty.barrier,
                    tma_desc_ptr=tma_desc_b,
                )
                cute.copy(
                    tma_atom_sfa,
                    tAgSFA[(None, kt)],
                    tAsSFA[(None, ab_empty.index)],
                    tma_bar_ptr=ab_empty.barrier,
                    tma_desc_ptr=tma_desc_sfa,
                )
                cute.copy(
                    tma_atom_sfb,
                    tBgSFB[(None, kt)],
                    tBsSFB[(None, ab_empty.index)],
                    tma_bar_ptr=ab_empty.barrier,
                    tma_desc_ptr=tma_desc_sfb,
                )

            ab_full = ab_consumer.wait_and_advance()

            # S2T
            s2t_stage_coord = (None, None, None, None, ab_full.index)
            cute.copy(
                tiled_copy_s2t_sfa,
                tCsSFA_compact_s2t[s2t_stage_coord],
                tCtSFA_compact_s2t,
            )
            cute.copy(
                tiled_copy_s2t_sfb,
                tCsSFB_compact_s2t[s2t_stage_coord],
                tCtSFB_compact_s2t,
            )

            # GEMM
            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_full.index)

                sf_kblock_coord = (None, None, kblock_idx)
                tiled_mma.set(tcgen05.Field.SFA, tCtSFA[sf_kblock_coord].iterator)
                tiled_mma.set(tcgen05.Field.SFB, tCtSFB[sf_kblock_coord].iterator)

                cute.gemm(
                    tiled_mma,
                    tCtAcc,
                    tCrA[kblock_coord],
                    tCrB[kblock_coord],
                    tCtAcc,
                )
                tiled_mma.set(tcgen05.Field.ACCUMULATE, True)

            ab_full.release()

        acc_empty.commit()

    # Epilogue
    op = tcgen05.Ld32x32bOp(tcgen05.Repetition.x128, tcgen05.Pack.NONE)
    copy_atom_t2r = cute.make_copy_atom(op, cutlass.Float32)
    tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tCtAcc[None, 0, 0])
    thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
    tDtAcc = thr_copy_t2r.partition_S(tCtAcc[None, 0, 0])
    tDgC = thr_copy_t2r.partition_D(tCgC[None, 0, 0])

    tDrAcc = cute.make_rmem_tensor(tDgC.shape, cutlass.Float32)
    tDrC = cute.make_rmem_tensor(tDgC.shape, c_dtype)

    tmem.relinquish_alloc_permit()
    acc_full = acc_consumer.wait_and_advance()

    cute.copy(tiled_copy_t2r, tDtAcc, tDrAcc)
    acc_vec = tDrAcc.load()
    tDrC.store(acc_vec.to(c_dtype))

    # Store: N is multiple of 128, so only M-tail
    simt_atom = cute.make_copy_atom(
        cute.nvgpu.CopyUniversalOp(), c_dtype, num_bits_per_copy=16
    )

    thread_layout = cute.make_layout((1, threads_per_cta), stride=(threads_per_cta, 1))
    value_layout = cute.make_layout((1, 1))
    tiled_copy_r2g = cute.make_tiled_copy_tv(simt_atom, thread_layout, value_layout)
    thr_copy_r2g = tiled_copy_r2g.get_slice(tidx)

    cC = cute.make_identity_tensor(gC_mnl.shape)
    tDcC = thr_copy_r2g.partition_D(cC)

    residue_m = mC_mnl.shape[0] - cutlass.Int32(coord_x) * cutlass.Int32(mma_tiler_mnk[0])

    if residue_m >= cutlass.Int32(mma_tiler_mnk[0]):
        cute.copy(simt_atom, cute.flatten(tDrC), cute.flatten(tDgC))
    else:
        tDpC = cute.make_rmem_tensor(tDrC.shape, cutlass.Boolean)
        residue_n = cutlass.Int32(mma_tiler_mnk[1])
        for i in range(cute.size(tDrC.shape)):
            tDpC[i] = cute.elem_less(tDcC[i], (residue_n, residue_m))
        cute.copy(
            simt_atom,
            cute.flatten(tDrC),
            cute.flatten(tDgC),
            pred=cute.flatten(tDpC),
        )

    acc_full.release()

    cute.arch.barrier()
    tmem.free(acc_tmem_ptr)
    pass


# -----------------------------------------------------------------------------
# Host-side JIT wrappers
# -----------------------------------------------------------------------------
@cute.jit
def init_tensormaps_jit(
    ptr_of_tensor_of_problem_sizes: cute.Pointer,
    ptr_of_tensor_of_abc_ptrs: cute.Pointer,
    ptr_of_tensor_of_sfasfb_ptrs: cute.Pointer,
    ptr_of_tensor_of_tensormap: cute.Pointer,
    num_groups: cutlass.Int32,
):
    tensor_of_abc_ptrs = cute.make_tensor(
        ptr_of_tensor_of_abc_ptrs, cute.make_layout((num_groups, 3), stride=(3, 1))
    )
    tensor_of_sfasfb_ptrs = cute.make_tensor(
        ptr_of_tensor_of_sfasfb_ptrs, cute.make_layout((num_groups, 2), stride=(2, 1))
    )
    tensor_of_problem_sizes = cute.make_tensor(
        ptr_of_tensor_of_problem_sizes, cute.make_layout((num_groups, 4), stride=(4, 1))
    )
    tensormaps = cute.make_tensor(
        ptr_of_tensor_of_tensormap,
        cute.make_layout((num_groups, 4, 16), stride=(64, 16, 1)),
    )

    # Fake tensors for atom creation
    min_a_shape = (cutlass.Int32(64), cutlass.Int32(64), cutlass.Int32(64), cutlass.Int32(1))
    min_b_shape = (cutlass.Int32(64), cutlass.Int32(64), cutlass.Int32(64), cutlass.Int32(1))

    initial_a = cute.make_tensor(
        cute.make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16),
        cute.make_layout(
            (min_a_shape[0], cute.assume(min_a_shape[2], 32), min_a_shape[3]),
            stride=(cute.assume(min_a_shape[2], 32), 1, cute.assume(min_a_shape[0] * min_a_shape[2], 32)),
        ),
    )
    initial_b = cute.make_tensor(
        cute.make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16),
        cute.make_layout(
            (min_b_shape[1], cute.assume(min_b_shape[2], 32), min_b_shape[3]),
            stride=(cute.assume(min_b_shape[2], 32), 1, cute.assume(min_b_shape[1] * min_b_shape[2], 32)),
        ),
    )

    sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(initial_a.shape, sf_vec_size)
    sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(initial_b.shape, sf_vec_size)

    initial_sfa = cute.make_tensor(
        cute.make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=16), sfa_layout
    )
    initial_sfb = cute.make_tensor(
        cute.make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=16), sfb_layout
    )

    mma_op = tcgen05.MmaMXF4NVF4Op(
        sf_dtype,
        (mma_tiler_mnk[0], mma_tiler_mnk[1], mma_inst_shape_k),
        tcgen05.CtaGroup.ONE,
        tcgen05.OperandSource.SMEM,
    )
    tiled_mma = cute.make_tiled_mma(mma_op)

    cluster_layout_vmnk = cute.tiled_divide(
        cute.make_layout((1, 1, 1)),
        (tiled_mma.thr_id.shape,),
    )

    # SMEM layouts for atoms
    a_smem_layout_staged = sm100_utils.make_smem_layout_a(
        tiled_mma, mma_tiler_mnk, ab_dtype, num_ab_stage
    )
    b_smem_layout_staged = sm100_utils.make_smem_layout_b(
        tiled_mma, mma_tiler_mnk, ab_dtype, num_ab_stage
    )
    sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
        tiled_mma, mma_tiler_mnk, sf_vec_size, num_ab_stage
    )
    sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(
        tiled_mma, mma_tiler_mnk, sf_vec_size, num_ab_stage
    )

    a_smem_layout = cute.slice_(a_smem_layout_staged, (None, None, None, 0))
    b_smem_layout = cute.slice_(b_smem_layout_staged, (None, None, None, 0))
    sfa_smem_layout = cute.slice_(sfa_smem_layout_staged, (None, None, None, 0))
    sfb_smem_layout = cute.slice_(sfb_smem_layout_staged, (None, None, None, 0))

    tma_atom_a, _ = cute.nvgpu.make_tiled_tma_atom_A(
        cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
        initial_a,
        a_smem_layout,
        mma_tiler_mnk,
        tiled_mma,
        cluster_layout_vmnk.shape,
    )
    tma_atom_b, _ = cute.nvgpu.make_tiled_tma_atom_B(
        cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
        initial_b,
        b_smem_layout,
        mma_tiler_mnk,
        tiled_mma,
        cluster_layout_vmnk.shape,
    )
    tma_atom_sfa, _ = cute.nvgpu.make_tiled_tma_atom_A(
        cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
        initial_sfa,
        sfa_smem_layout,
        mma_tiler_mnk,
        tiled_mma,
        cluster_layout_vmnk.shape,
        internal_type=cutlass.Int16,
    )
    tma_atom_sfb, _ = cute.nvgpu.make_tiled_tma_atom_B(
        cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
        initial_sfb,
        sfb_smem_layout,
        mma_tiler_mnk,
        tiled_mma,
        cluster_layout_vmnk.shape,
        internal_type=cutlass.Int16,
    )

    init_tensormaps_kernel(
        tma_atom_a,
        tma_atom_b,
        tma_atom_sfa,
        tma_atom_sfb,
        tensor_of_abc_ptrs,
        tensor_of_sfasfb_ptrs,
        tensor_of_problem_sizes,
        tensormaps,
    ).launch(
        grid=(1, 1, num_groups),
        block=[threads_per_cta, 1, 1],
        cluster=(1, 1, 1),
    )
    return


@cute.jit
def gemm_jit(
    ptr_of_tensor_of_problem_sizes: cute.Pointer,
    ptr_of_tensor_of_abc_ptrs: cute.Pointer,
    ptr_of_tensor_of_sfasfb_ptrs: cute.Pointer,
    ptr_of_tensor_of_tensormap: cute.Pointer,
    ptr_of_tensor_of_cluster_meta: cute.Pointer,
    total_num_clusters: cutlass.Int32,
    num_groups: cutlass.Int32,
):
    tensor_of_abc_ptrs = cute.make_tensor(
        ptr_of_tensor_of_abc_ptrs, cute.make_layout((num_groups, 3), stride=(3, 1))
    )
    tensor_of_sfasfb_ptrs = cute.make_tensor(
        ptr_of_tensor_of_sfasfb_ptrs, cute.make_layout((num_groups, 2), stride=(2, 1))
    )
    tensor_of_problem_sizes = cute.make_tensor(
        ptr_of_tensor_of_problem_sizes, cute.make_layout((num_groups, 4), stride=(4, 1))
    )
    tensormaps = cute.make_tensor(
        ptr_of_tensor_of_tensormap, cute.make_layout((num_groups, 4, 16), stride=(64, 16, 1))
    )
    tensor_of_cluster_meta = cute.make_tensor(
        ptr_of_tensor_of_cluster_meta, cute.make_layout((total_num_clusters, 3), stride=(3, 1))
    )

    # Fake tensors for atom creation
    min_a_shape = (cutlass.Int32(64), cutlass.Int32(64), cutlass.Int32(64), cutlass.Int32(1))
    min_b_shape = (cutlass.Int32(64), cutlass.Int32(64), cutlass.Int32(64), cutlass.Int32(1))

    initial_a = cute.make_tensor(
        cute.make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16),
        cute.make_layout(
            (min_a_shape[0], cute.assume(min_a_shape[2], 32), min_a_shape[3]),
            stride=(cute.assume(min_a_shape[2], 32), 1, cute.assume(min_a_shape[0] * min_a_shape[2], 32)),
        ),
    )
    initial_b = cute.make_tensor(
        cute.make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16),
        cute.make_layout(
            (min_b_shape[1], cute.assume(min_b_shape[2], 32), min_b_shape[3]),
            stride=(cute.assume(min_b_shape[2], 32), 1, cute.assume(min_b_shape[1] * min_b_shape[2], 32)),
        ),
    )

    sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(initial_a.shape, sf_vec_size)
    sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(initial_b.shape, sf_vec_size)

    initial_sfa = cute.make_tensor(
        cute.make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=16), sfa_layout
    )
    initial_sfb = cute.make_tensor(
        cute.make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=16), sfb_layout
    )

    mma_op = tcgen05.MmaMXF4NVF4Op(
        sf_dtype,
        (mma_tiler_mnk[0], mma_tiler_mnk[1], mma_inst_shape_k),
        tcgen05.CtaGroup.ONE,
        tcgen05.OperandSource.SMEM,
    )
    tiled_mma = cute.make_tiled_mma(mma_op)

    cluster_layout_vmnk = cute.tiled_divide(
        cute.make_layout((1, 1, 1)),
        (tiled_mma.thr_id.shape,),
    )

    # SMEM layouts
    a_smem_layout_staged = sm100_utils.make_smem_layout_a(
        tiled_mma, mma_tiler_mnk, ab_dtype, num_ab_stage
    )
    b_smem_layout_staged = sm100_utils.make_smem_layout_b(
        tiled_mma, mma_tiler_mnk, ab_dtype, num_ab_stage
    )
    sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
        tiled_mma, mma_tiler_mnk, sf_vec_size, num_ab_stage
    )
    sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(
        tiled_mma, mma_tiler_mnk, sf_vec_size, num_ab_stage
    )

    a_smem_layout = cute.slice_(a_smem_layout_staged, (None, None, None, 0))
    b_smem_layout = cute.slice_(b_smem_layout_staged, (None, None, None, 0))
    sfa_smem_layout = cute.slice_(sfa_smem_layout_staged, (None, None, None, 0))
    sfb_smem_layout = cute.slice_(sfb_smem_layout_staged, (None, None, None, 0))

    tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
        cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
        initial_a,
        a_smem_layout,
        mma_tiler_mnk,
        tiled_mma,
        cluster_layout_vmnk.shape,
    )
    tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
        cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
        initial_b,
        b_smem_layout,
        mma_tiler_mnk,
        tiled_mma,
        cluster_layout_vmnk.shape,
    )
    tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A(
        cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
        initial_sfa,
        sfa_smem_layout,
        mma_tiler_mnk,
        tiled_mma,
        cluster_layout_vmnk.shape,
        internal_type=cutlass.Int16,
    )
    tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B(
        cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
        initial_sfb,
        sfb_smem_layout,
        mma_tiler_mnk,
        tiled_mma,
        cluster_layout_vmnk.shape,
        internal_type=cutlass.Int16,
    )

    atom_thr_size = cute.size(tiled_mma.thr_id.shape)
    a_copy_size = cute.size_in_bytes(ab_dtype, a_smem_layout)
    b_copy_size = cute.size_in_bytes(ab_dtype, b_smem_layout)
    sfa_copy_size = cute.size_in_bytes(sf_dtype, sfa_smem_layout)
    sfb_copy_size = cute.size_in_bytes(sf_dtype, sfb_smem_layout)
    num_tma_load_bytes = (a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size) * atom_thr_size

    gemm_kernel(
        tiled_mma,
        tma_atom_a,
        tma_tensor_a,
        tma_atom_b,
        tma_tensor_b,
        tma_atom_sfa,
        tma_tensor_sfa,
        tma_atom_sfb,
        tma_tensor_sfb,
        tensor_of_abc_ptrs,
        tensor_of_sfasfb_ptrs,
        tensormaps,
        tensor_of_problem_sizes,
        tensor_of_cluster_meta,
        a_smem_layout_staged,
        b_smem_layout_staged,
        sfa_smem_layout_staged,
        sfb_smem_layout_staged,
        num_tma_load_bytes,
    ).launch(
        grid=(1, 1, total_num_clusters),
        block=[threads_per_cta, 1, 1],
        cluster=(1, 1, 1),
    )
    return


# -----------------------------------------------------------------------------
# Compile caches
# -----------------------------------------------------------------------------
_compiled_init_cache: Dict[str, Any] = {}
_compiled_gemm_cache: Dict[str, Any] = {}

_runtime_cache: Dict[Any, Any] = {}


def compile_init(num_groups: int):
    key = f"ng={num_groups}"
    if key in _compiled_init_cache:
        return _compiled_init_cache[key]

    cute_ptr_ps = make_ptr(cutlass.Int32, 0, cute.AddressSpace.gmem, assumed_align=16)
    cute_ptr_abc = make_ptr(cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16)
    cute_ptr_sfs = make_ptr(cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16)
    cute_ptr_tm = make_ptr(cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16)

    ng = cutlass.Int32(num_groups)

    fn = cute.compile(
        init_tensormaps_jit,
        cute_ptr_ps,
        cute_ptr_abc,
        cute_ptr_sfs,
        cute_ptr_tm,
        ng,
    )
    _compiled_init_cache[key] = fn
    return fn


def compile_gemm(num_groups: int):
    key = f"ng={num_groups}"
    if key in _compiled_gemm_cache:
        return _compiled_gemm_cache[key]

    cute_ptr_ps = make_ptr(cutlass.Int32, 0, cute.AddressSpace.gmem, assumed_align=16)
    cute_ptr_abc = make_ptr(cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16)
    cute_ptr_sfs = make_ptr(cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16)
    cute_ptr_tm = make_ptr(cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16)
    cute_ptr_meta = make_ptr(cutlass.Int32, 0, cute.AddressSpace.gmem, assumed_align=16)

    total_clusters = cutlass.Int32(1)
    ng = cutlass.Int32(num_groups)

    fn = cute.compile(
        gemm_jit,
        cute_ptr_ps,
        cute_ptr_abc,
        cute_ptr_sfs,
        cute_ptr_tm,
        cute_ptr_meta,
        total_clusters,
        ng,
    )
    _compiled_gemm_cache[key] = fn
    return fn


# -----------------------------------------------------------------------------
# Entry point
# -----------------------------------------------------------------------------
def custom_kernel(data: input_t) -> output_t:
    abc_tensors, _, sfasfb_reordered_tensors, problem_sizes = data
    num_groups = len(problem_sizes)

    init_fn = compile_init(num_groups)
    gemm_fn = compile_gemm(num_groups)

    # Build host pointer tuples
    abc_ptrs = []
    sfs_ptrs = []
    for (a, b, c), (sfa_r, sfb_r), _sz in zip(abc_tensors, sfasfb_reordered_tensors, problem_sizes):
        abc_ptrs.append((a.data_ptr(), b.data_ptr(), c.data_ptr()))
        sfs_ptrs.append((sfa_r.data_ptr(), sfb_r.data_ptr()))

    # Runtime cache per (problem_sizes)
    ps_key = tuple((int(m), int(n), int(k), int(l)) for (m, n, k, l) in problem_sizes)
    cache_key = (num_groups, ps_key)

    if cache_key not in _runtime_cache:
        # problem_sizes tensor (device)
        tensor_of_problem_sizes = torch.tensor(problem_sizes, dtype=torch.int32, device="cuda")

        # cluster_meta (device)
        cluster_meta = []
        cta_m = mma_tiler_mnk[0]
        cta_n = mma_tiler_mnk[1]
        for g, (m, n, k, l) in enumerate(problem_sizes):
            tiles_m = ceil_div(m, cta_m)
            tiles_n = ceil_div(n, cta_n)
            for ty in range(tiles_n):
                for tx in range(tiles_m):
                    cluster_meta.append((g, tx, ty))
        total_num_clusters = len(cluster_meta)
        tensor_of_cluster_meta = torch.tensor(cluster_meta, dtype=torch.int32, device="cuda")

        # persistent device pointer arrays
        tensor_of_abc_ptrs = torch.empty((num_groups, 3), dtype=torch.int64, device="cuda")
        tensor_of_sfs_ptrs = torch.empty((num_groups, 2), dtype=torch.int64, device="cuda")

        # tensormaps per group
        tensor_of_tensormap = torch.empty((num_groups, 4, bytes_per_tensormap // 8), dtype=torch.int64, device="cuda")

        _runtime_cache[cache_key] = {
            "tensor_of_problem_sizes": tensor_of_problem_sizes,
            "tensor_of_cluster_meta": tensor_of_cluster_meta,
            "total_num_clusters": total_num_clusters,
            "tensor_of_abc_ptrs": tensor_of_abc_ptrs,
            "tensor_of_sfs_ptrs": tensor_of_sfs_ptrs,
            "tensor_of_tensormap": tensor_of_tensormap,
            "last_abc_ptrs": None,
            "last_sfs_ptrs": None,
        }

    st = _runtime_cache[cache_key]

    # Update pointer arrays if changed
    need_init = False
    if st["last_abc_ptrs"] != tuple(abc_ptrs) or st["last_sfs_ptrs"] != tuple(sfs_ptrs):
        # small host tensors -> copy into persistent device tensors
        cpu_abc = torch.tensor(abc_ptrs, dtype=torch.int64, device="cpu")
        cpu_sfs = torch.tensor(sfs_ptrs, dtype=torch.int64, device="cpu")
        st["tensor_of_abc_ptrs"].copy_(cpu_abc, non_blocking=False)
        st["tensor_of_sfs_ptrs"].copy_(cpu_sfs, non_blocking=False)
        st["last_abc_ptrs"] = tuple(abc_ptrs)
        st["last_sfs_ptrs"] = tuple(sfs_ptrs)
        need_init = True

    # CuTe pointers
    cute_ptr_ps = make_ptr(
        cutlass.Int32, st["tensor_of_problem_sizes"].data_ptr(), cute.AddressSpace.gmem, assumed_align=16
    )
    cute_ptr_abc = make_ptr(
        cutlass.Int64, st["tensor_of_abc_ptrs"].data_ptr(), cute.AddressSpace.gmem, assumed_align=16
    )
    cute_ptr_sfs = make_ptr(
        cutlass.Int64, st["tensor_of_sfs_ptrs"].data_ptr(), cute.AddressSpace.gmem, assumed_align=16
    )
    cute_ptr_tm = make_ptr(
        cutlass.Int64, st["tensor_of_tensormap"].data_ptr(), cute.AddressSpace.gmem, assumed_align=16
    )
    cute_ptr_meta = make_ptr(
        cutlass.Int32, st["tensor_of_cluster_meta"].data_ptr(), cute.AddressSpace.gmem, assumed_align=16
    )

    # Init tensormaps once per (problem_sizes, pointers)
    if need_init:
        init_fn(
            cute_ptr_ps,
            cute_ptr_abc,
            cute_ptr_sfs,
            cute_ptr_tm,
            num_groups,
        )

    # GEMM
    gemm_fn(
        cute_ptr_ps,
        cute_ptr_abc,
        cute_ptr_sfs,
        cute_ptr_tm,
        cute_ptr_meta,
        st["total_num_clusters"],
        num_groups,
    )

    return [abc_tensors[i][2] for i in range(num_groups)]