submission 493184

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

# Inline asm imports
from cutlass._mlir.dialects import llvm
from cutlass.cutlass_dsl import T, dsl_user_op
from cutlass.cute.typing import Int32 as DSLInt32, Int64 as DSLInt64

import functools
from typing import Tuple, List

import torch
from task import input_t, output_t

# ============================================================
# v34: Block tile distribution + inline PTX epilogue
# Changes from v31:
#   1. Block tile distribution: consecutive tiles per CTA instead of
#      interleaved (stride-148). Reduces group changes from ~2-3/CTA
#      to ~0-1/CTA for G=8 (saves ~3-5µs per avoided group change).
#   2. Full tile epilogue uses inline PTX (from v31):
#      - tcgen05.ld.sync.aligned.32x32b.x128.b32 (T2R load)
#      - cvt.rn.f16x2.f32 (FP32 -> packed FP16)
#      - st.global.v4.b32 (vectorized 128-bit stores)
# ============================================================

bytes_per_tensormap = 128
num_tensormaps = 4
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 = 5
num_tmem_alloc_cols = 512

TMA_WARP = 0
MMA_WARP = 1
MAX_ACTIVE_CLUSTERS = 148


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


# ============================================================
# INLINE ASM HELPERS
# ============================================================

def _build_epilogue_asm():
    """Build the PTX asm string for the inline epilogue with M predication.

    Operands:
        $0 = tmem_addr (r, 32-bit) — TMEM column base address
        $1 = c_base (l, 64-bit) — C tensor base address (byte ptr)
        $2 = m_idx (r, 32-bit) — thread's global M-row index (coord_x*128 + tidx)
        $3 = n_start (r, 32-bit) — tile's N-column start (coord_y * 128)
        $4 = n_stride (r, 32-bit) — N dimension (row stride in elements)
        $5 = m_valid (r, 32-bit) — total M for this group (thread stores only if m_idx < m_valid)

    Computes: addr = c_base + (m_idx * n_stride + n_start) * 2
    T2R load (all 128 threads, sync) -> convert FP32->FP16 -> predicated vectorized stores
    """
    lines = []
    lines.append("{")

    # Declare registers
    lines.append("    .reg .b32 t<128>;")    # T2R output (128 FP32 values)
    lines.append("    .reg .b32 p<64>;")     # Packed FP16x2 values
    lines.append("    .reg .b64 addr;")
    lines.append("    .reg .b64 tmp64;")
    lines.append("    .reg .b64 tmp64b;")
    lines.append("    .reg .pred valid;")    # M predication

    # M predication: only store if this thread's global M-row < total M
    lines.append("    setp.lt.s32 valid, $2, $5;")  # m_idx < m_valid

    # Address computation: addr = c_base + (m_idx * n_stride + n_start) * 2
    lines.append("    mul.wide.s32 tmp64, $2, $4;")      # m_idx * n_stride -> 64-bit
    lines.append("    cvt.s64.s32 tmp64b, $3;")           # n_start -> 64-bit
    lines.append("    add.s64 tmp64, tmp64, tmp64b;")      # + n_start
    lines.append("    shl.b64 tmp64, tmp64, 1;")           # * 2 (FP16 = 2 bytes)
    lines.append("    add.u64 addr, $1, tmp64;")           # + c_base

    # T2R load: 128 FP32 values from TMEM (all threads must participate for sync)
    t_list = ", ".join(f"t{i}" for i in range(128))
    lines.append(f"    tcgen05.ld.sync.aligned.32x32b.x128.b32 {{{t_list}}}, [$0];")

    # Convert FP32 pairs -> packed FP16x2 (64 instructions)
    for i in range(64):
        lines.append(f"    cvt.rn.f16x2.f32 p{i}, t{2*i+1}, t{2*i};")

    # Predicated vectorized global stores: 16 x v4.b32 = 128 FP16 values
    for i in range(16):
        offset = i * 16
        p_base = i * 4
        p_args = ", ".join(f"p{p_base+j}" for j in range(4))
        lines.append(f"    @valid st.global.v4.b32 [addr+{offset}], {{{p_args}}};")

    lines.append("}")
    return "\n\t".join(lines)

# Pre-build the asm string (it's constant)
EPILOGUE_ASM = _build_epilogue_asm()


def get_tmem_raw_addr(tmem_ptr):
    """Extract raw TMEM column address (uint32) from CuTe TMEM pointer.
    Uses CuTe's native ptrtoint (not llvm.ptrtoint which requires LLVM ptr types).
    """
    return tmem_ptr.toint()


@dsl_user_op
def inline_epilogue(tmem_addr, c_base, m_idx, n_start, n_stride, m_valid, *, loc=None, ip=None):
    """Inline PTX epilogue with M predication.

    T2R load + FP32->FP16 convert + M-predicated vectorized store.
    Handles both full and partial tiles (only needs M predication since N is always 128-aligned).

    Args:
        tmem_addr: Int32 — raw TMEM column address
        c_base: Int64 — C tensor base byte address
        m_idx: Int32 — thread's global M-row index (coord_x * 128 + tidx)
        n_start: Int32 — tile's N-column start (coord_y * 128)
        n_stride: Int32 — N dimension (elements per row)
        m_valid: Int32 — total M for this group (store predicate: m_idx < m_valid)
    """
    llvm.inline_asm(
        res=None,
        operands_=[
            DSLInt32(tmem_addr).ir_value(loc=loc, ip=ip),
            DSLInt64(c_base).ir_value(loc=loc, ip=ip),
            DSLInt32(m_idx).ir_value(loc=loc, ip=ip),
            DSLInt32(n_start).ir_value(loc=loc, ip=ip),
            DSLInt32(n_stride).ir_value(loc=loc, ip=ip),
            DSLInt32(m_valid).ir_value(loc=loc, ip=ip),
        ],
        asm_string=EPILOGUE_ASM,
        constraints="r,l,r,r,r,r",
        has_side_effects=True,
        is_align_stack=False,
        asm_dialect=llvm.AsmDialect.AD_ATT,
        loc=loc,
        ip=ip,
    )


@cute.kernel
def kernel(
    tiled_mma: cute.TiledMma,
    tma_atom_a: cute.CopyAtom,
    mA_mkl: cute.Tensor,
    tma_atom_b: cute.CopyAtom,
    mB_nkl: cute.Tensor,
    tma_atom_sfa: cute.CopyAtom,
    mSFA_mkl: cute.Tensor,
    tma_atom_sfb: cute.CopyAtom,
    mSFB_nkl: cute.Tensor,
    tensor_of_abc_ptrs: cute.Tensor,
    tensor_of_sfasfb_ptrs: cute.Tensor,
    tensormaps: cute.Tensor,
    tensor_of_problem_sizes: cute.Tensor,
    a_smem_layout_staged: cute.ComposedLayout,
    b_smem_layout_staged: cute.ComposedLayout,
    sfa_smem_layout_staged: cute.Layout,
    sfb_smem_layout_staged: cute.Layout,
    cta_mn_list: List[Tuple[int, int]],
    num_tma_load_bytes: cutlass.Constexpr[int],
    total_tiles: cutlass.Int32,
):
    warp_idx = cute.arch.warp_idx()
    warp_idx = cute.arch.make_warp_uniform(warp_idx)
    tidx, _, _ = cute.arch.thread_idx()
    bidx, bidy, bidz = cute.arch.block_idx()
    _, _, gdimz = cute.arch.grid_dim()

    # ============================================================
    # ONE-TIME SETUP
    # ============================================================

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

    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

    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)

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

    # Tensormap init (once per CTA)
    tensormap_manager = utils.TensorMapManager(utils.TensorMapUpdateMode.GMEM, 128)
    tensormap_a_gmem_ptr = tensormap_manager.get_tensormap_ptr(tensormaps[(bidz, 0, None)].iterator)
    tensormap_b_gmem_ptr = tensormap_manager.get_tensormap_ptr(tensormaps[(bidz, 1, None)].iterator)
    tensormap_sfa_gmem_ptr = tensormap_manager.get_tensormap_ptr(tensormaps[(bidz, 2, None)].iterator)
    tensormap_sfb_gmem_ptr = tensormap_manager.get_tensormap_ptr(tensormaps[(bidz, 3, None)].iterator)

    if warp_idx == TMA_WARP:
        tensormap_manager.init_tensormap_from_atom(tma_atom_a, tensormap_a_gmem_ptr, 0)
        tensormap_manager.init_tensormap_from_atom(tma_atom_b, tensormap_b_gmem_ptr, 0)
        tensormap_manager.init_tensormap_from_atom(tma_atom_sfa, tensormap_sfa_gmem_ptr, 0)
        tensormap_manager.init_tensormap_from_atom(tma_atom_sfb, tensormap_sfb_gmem_ptr, 0)

    cute.arch.barrier()

    # Global tensor partitions
    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)

    # TMA partitions (full — sliced per-tile inside loop)
    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
    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 allocation (once)
    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_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 setup
    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)

    # Extract raw TMEM address for inline epilogue (once, outside tile loop)
    tmem_raw_addr = get_tmem_raw_addr(acc_tmem_ptr)

    # ============================================================
    # PERSISTENT TILE LOOP
    # ============================================================

    last_group_idx = cutlass.Int32(-1)

    # Block tile distribution: consecutive tiles per CTA → fewer group changes
    tiles_per_cta = cute.ceil_div(total_tiles, gdimz)

    for local_idx in cutlass.range(0, tiles_per_cta, 1, unroll=1):
        tile_linear_idx = bidz * tiles_per_cta + local_idx

        if tile_linear_idx < total_tiles:
            # --- Delinearize tile_linear_idx → (group_idx, coord_x, coord_y) ---
            group_idx = 0
            find = False
            coord_x = 0
            coord_y = 0
            cta_rest = tile_linear_idx
            for _, (cta_m, cta_n) in enumerate(cta_mn_list):
                if cta_rest >= (cta_m * cta_n):
                    group_idx += 1
                    cta_rest -= cta_m * cta_n
                else:
                    if not find:
                        coord_y = cta_rest // cta_m
                        coord_x = cta_rest % cta_m
                        cta_rest -= cta_m * cta_n
                        find = True

            m = tensor_of_problem_sizes[group_idx, 0]
            n = tensor_of_problem_sizes[group_idx, 1]
            k = tensor_of_problem_sizes[group_idx, 2]
            l = tensor_of_problem_sizes[group_idx, 3]
            k_tile_cnt = cute.ceil_div(k, mma_tiler_mnk[2])

            # --- Update tensormaps when group changes ---
            is_group_changed = group_idx != last_group_idx
            if is_group_changed:
                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)
                mA_mkl_layout = cute.make_layout(
                    (m, k, l), stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32)))
                mB_nkl_layout = cute.make_layout(
                    (n, k, l), stride=(cute.assume(k, 32), 1, cute.assume(n * k, 32)))
                atom_shape = ((32, 4), (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)

                if warp_idx == TMA_WARP:
                    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)

                cute.arch.barrier()

            # --- Slice TMA to tile coords ---
            tAgA_tile = tAgA[(None, coord_x, None, 0)]
            tBgB_tile = tBgB[(None, coord_y, None, 0)]
            tAgSFA_tile = tAgSFA[(None, coord_x, None, 0)]
            tBgSFB_tile = tBgSFB[(None, coord_y, None, 0)]

            # ============================================================
            # WARP-SPECIALIZED MAINLOOP
            # ============================================================

            if warp_idx == TMA_WARP:
                for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):
                    ab_empty = ab_producer.acquire_and_advance()
                    cute.copy(tma_atom_a, tAgA_tile[(None, k_tile)], tAsA[(None, ab_empty.index)],
                              tma_bar_ptr=ab_empty.barrier,
                              tma_desc_ptr=tensormap_manager.get_tensormap_ptr(tensormap_a_gmem_ptr, cute.AddressSpace.generic))
                    cute.copy(tma_atom_b, tBgB_tile[(None, k_tile)], tBsB[(None, ab_empty.index)],
                              tma_bar_ptr=ab_empty.barrier,
                              tma_desc_ptr=tensormap_manager.get_tensormap_ptr(tensormap_b_gmem_ptr, cute.AddressSpace.generic))
                    cute.copy(tma_atom_sfa, tAgSFA_tile[(None, k_tile)], tAsSFA[(None, ab_empty.index)],
                              tma_bar_ptr=ab_empty.barrier,
                              tma_desc_ptr=tensormap_manager.get_tensormap_ptr(tensormap_sfa_gmem_ptr, cute.AddressSpace.generic))
                    cute.copy(tma_atom_sfb, tBgSFB_tile[(None, k_tile)], tBsSFB[(None, ab_empty.index)],
                              tma_bar_ptr=ab_empty.barrier,
                              tma_desc_ptr=tensormap_manager.get_tensormap_ptr(tensormap_sfb_gmem_ptr, cute.AddressSpace.generic))

            if warp_idx == MMA_WARP:
                acc_empty = acc_producer.acquire_and_advance()
                tiled_mma.set(tcgen05.Field.ACCUMULATE, False)

                for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):
                    ab_full = ab_consumer.wait_and_advance()

                    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)

                    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 (all warps) — inline PTX for ALL tiles
            # ============================================================

            acc_full = acc_consumer.wait_and_advance()

            # Inline PTX epilogue handles both full and partial tiles via M predication
            c_base = tensor_of_abc_ptrs[group_idx, 2]
            m_idx = cutlass.Int32(coord_x) * 128 + tidx
            n_start = cutlass.Int32(coord_y) * 128
            inline_epilogue(tmem_raw_addr, c_base, m_idx, n_start, n, m)

            acc_full.release()

            last_group_idx = group_idx

    # ============================================================
    # CLEANUP
    # ============================================================
    tmem.relinquish_alloc_permit()
    cute.arch.barrier()
    tmem.free(acc_tmem_ptr)
    pass


@cute.jit
def my_kernel(
    ptr_of_tensor_of_problem_sizes: cute.Pointer,
    ptr_of_tensor_of_abc_ptrs: cute.Pointer,
    ptr_of_tensor_of_sfasfb_ptrs: cute.Pointer,
    ptr_of_tensor_of_tensormap: cute.Pointer,
    total_num_clusters: cutlass.Int32,
    problem_sizes: List[Tuple[int, int, int, int]],
    num_groups: cutlass.Int32,
):
    tensor_of_abc_ptrs = cute.make_tensor(
        ptr_of_tensor_of_abc_ptrs, cute.make_layout((num_groups, 3), stride=(3, 1)))
    tensor_of_sfasfb_ptrs = cute.make_tensor(
        ptr_of_tensor_of_sfasfb_ptrs, cute.make_layout((num_groups, 2), stride=(2, 1)))
    tensor_of_problem_sizes = cute.make_tensor(
        ptr_of_tensor_of_problem_sizes, cute.make_layout((num_groups, 4), stride=(4, 1)))
    tensor_of_tensormap = cute.make_tensor(
        ptr_of_tensor_of_tensormap, cute.make_layout((MAX_ACTIVE_CLUSTERS, 4, 16), stride=(64, 16, 1)))

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

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

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

    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

    cta_mn_list = []
    for group_idx, (m, n, k, l) in enumerate(problem_sizes):
        x, y = cute.ceil_div(problem_sizes[group_idx][:2], mma_tiler_mnk[0:2])
        cta_mn_list.append((x, y))

    grid = (1, 1, MAX_ACTIVE_CLUSTERS)

    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,
        tensor_of_tensormap, tensor_of_problem_sizes,
        a_smem_layout_staged, b_smem_layout_staged,
        sfa_smem_layout_staged, sfb_smem_layout_staged,
        cta_mn_list, num_tma_load_bytes,
        total_num_clusters,
    ).launch(grid=grid, block=[threads_per_cta, 1, 1], cluster=(1, 1, 1))
    return


_compiled_kernel_cache = {}
_buf_ps = None
_buf_abc = None
_buf_sf = None
_buf_tm = None
_cute_ptr_ps = None
_cute_ptr_abc = None
_cute_ptr_sf = None
_cute_ptr_tm = None

def compile_kernel(problem_sizes):
    global _compiled_kernel_cache
    cache_key = f"{len(problem_sizes)}"
    if cache_key in _compiled_kernel_cache:
        return _compiled_kernel_cache[cache_key]

    cute_ptr_of_tensor_of_problem_sizes = make_ptr(cutlass.Int32, 0, cute.AddressSpace.gmem, assumed_align=16)
    cute_ptr_of_tensor_of_abc_ptrs = make_ptr(cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16)
    cute_ptr_of_tensor_of_sfasfb_ptrs = make_ptr(cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16)
    cute_ptr_of_tensor_of_tensormap = make_ptr(cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16)
    total_num_clusters = cutlass.Int32(1)
    num_groups = cutlass.Int32(len(problem_sizes))

    import sys
    try:
        compiled_func = cute.compile(
            my_kernel,
            cute_ptr_of_tensor_of_problem_sizes,
            cute_ptr_of_tensor_of_abc_ptrs,
            cute_ptr_of_tensor_of_sfasfb_ptrs,
            cute_ptr_of_tensor_of_tensormap,
            total_num_clusters,
            problem_sizes,
            num_groups
        )
    except Exception as e:
        print(f"COMPILATION ERROR: {type(e).__name__}: {e}", file=sys.stderr)
        import traceback
        traceback.print_exc(file=sys.stderr)
        raise RuntimeError(f"Compilation failed: {e}") from None
    _compiled_kernel_cache[cache_key] = compiled_func
    return compiled_func


def custom_kernel(data: input_t) -> output_t:
    global _buf_ps, _buf_abc, _buf_sf, _buf_tm
    global _cute_ptr_ps, _cute_ptr_abc, _cute_ptr_sf, _cute_ptr_tm
    abc_tensors, _, sfasfb_reordered_tensors, problem_sizes = data
    compiled_func = compile_kernel(problem_sizes)

    num_groups = len(problem_sizes)
    total_num_clusters = 0
    for m, n, _, _ in problem_sizes:
        total_num_clusters += ((m + 127) // 128) * ((n + 127) // 128)

    if _buf_ps is None or _buf_ps.shape[0] < num_groups:
        _buf_ps = torch.empty((num_groups, 4), dtype=torch.int32, device='cuda')
        _buf_abc = torch.empty((num_groups, 3), dtype=torch.int64, device='cuda')
        _buf_sf = torch.empty((num_groups, 2), dtype=torch.int64, device='cuda')
        _cute_ptr_ps = make_ptr(cutlass.Int32, _buf_ps.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
        _cute_ptr_abc = make_ptr(cutlass.Int64, _buf_abc.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
        _cute_ptr_sf = make_ptr(cutlass.Int64, _buf_sf.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
    if _buf_tm is None:
        _buf_tm = torch.empty((MAX_ACTIVE_CLUSTERS, num_tensormaps, bytes_per_tensormap // 8), dtype=torch.int64, device='cuda')
        _cute_ptr_tm = make_ptr(cutlass.Int64, _buf_tm.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)

    ps_cpu = torch.tensor(problem_sizes, dtype=torch.int32)
    abc_cpu = torch.empty((num_groups, 3), dtype=torch.int64)
    sf_cpu = torch.empty((num_groups, 2), dtype=torch.int64)
    for i, ((a, b, c), (sfa, sfb)) in enumerate(
        zip(abc_tensors, sfasfb_reordered_tensors)):
        abc_cpu[i, 0] = a.data_ptr()
        abc_cpu[i, 1] = b.data_ptr()
        abc_cpu[i, 2] = c.data_ptr()
        sf_cpu[i, 0] = sfa.data_ptr()
        sf_cpu[i, 1] = sfb.data_ptr()
    _buf_ps[:num_groups].copy_(ps_cpu, non_blocking=True)
    _buf_abc[:num_groups].copy_(abc_cpu, non_blocking=True)
    _buf_sf[:num_groups].copy_(sf_cpu, non_blocking=True)

    compiled_func(_cute_ptr_ps, _cute_ptr_abc, _cute_ptr_sf, _cute_ptr_tm,
                  total_num_clusters, problem_sizes, num_groups)
    return [abc[2] for abc in abc_tensors]