submission 156282

import cuda.bindings.driver as cuda
import torch
from task import input_t, output_t

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

MMA_TILER_MNK = (128, 128, 256)
MMA_INST_SHAPE_K = 64
AB_DTYPE = cutlass.Float4E2M1FN
SF_DTYPE = cutlass.Float8E4M3FN
C_DTYPE = cutlass.Float16
ACC_DTYPE = cutlass.Float32
SF_VEC_SIZE = 16
THREADS_PER_CTA = 128
NUM_AB_STAGE = 5
NUM_ACC_STAGE = 2
NUM_TMEM_ALLOC_COLS = 512


@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,
    mC_mnl: 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,
    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()
    bidx, bidy, bidz = cute.arch.block_idx()
    
    cta_coord = (bidx, bidy, bidz)
    mma_tile_coord_mnl = (
        cta_coord[0] // cute.size(tiled_mma.thr_id.shape),
        cta_coord[1],
        cta_coord[2],
    )
    
    @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)
    
    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_producer, ab_consumer = pipeline.PipelineTmaUmma.create(
        barrier_storage=storage.ab_mbar_ptr.data_ptr(),
        num_stages=NUM_AB_STAGE,
        producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
        consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1),
        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()
    
    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))
    gC_mnl = cute.local_tile(mC_mnl, cute.slice_(MMA_TILER_MNK, (None, None, 0)), (None, None, None))
    
    k_tile_cnt = cute.size(gA_mkl, mode=[3])
    
    thr_mma = tiled_mma.get_slice(0)
    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)
    
    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)
    
    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(ACC_DTYPE)
    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)
    
    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 = tcgen05.get_s2t_smem_desc_tensor(
        tiled_copy_s2t_sfa, thr_copy_s2t_sfa.partition_S(tCsSFA_compact))
    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 = tcgen05.get_s2t_smem_desc_tensor(
        tiled_copy_s2t_sfb, thr_copy_s2t_sfb.partition_S(tCsSFB_compact))
    tCtSFB_compact_s2t = thr_copy_s2t_sfb.partition_D(tCtSFB_compact)
    
    tAgA = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
    tBgB = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
    tAgSFA = tAgSFA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
    tBgSFB = tBgSFB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
    
    if warp_idx == 0:
        acc_empty = acc_producer.acquire_and_advance()
        tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
        
        for k_tile in range(k_tile_cnt):
            ab_empty = ab_producer.acquire_and_advance()
            
            cute.copy(tma_atom_a, tAgA[(None, k_tile)], tAsA[(None, ab_empty.index)], tma_bar_ptr=ab_empty.barrier)
            cute.copy(tma_atom_b, tBgB[(None, k_tile)], tBsB[(None, ab_empty.index)], tma_bar_ptr=ab_empty.barrier)
            cute.copy(tma_atom_sfa, tAgSFA[(None, k_tile)], tAsSFA[(None, ab_empty.index)], tma_bar_ptr=ab_empty.barrier)
            cute.copy(tma_atom_sfb, tBgSFB[(None, k_tile)], tBsSFB[(None, ab_empty.index)], tma_bar_ptr=ab_empty.barrier)
            
            ab_full = ab_consumer.wait_and_advance()
            
            s2t_coord = (None, None, None, None, ab_full.index)
            cute.copy(tiled_copy_s2t_sfa, tCsSFA_compact_s2t[s2t_coord], tCtSFA_compact_s2t)
            cute.copy(tiled_copy_s2t_sfb, tCsSFB_compact_s2t[s2t_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_coord = (None, None, kblock_idx)
                
                tiled_mma.set(tcgen05.Field.SFA, tCtSFA[sf_coord].iterator)
                tiled_mma.set(tcgen05.Field.SFB, tCtSFB[sf_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()
    
    op = tcgen05.Ld32x32bOp(tcgen05.Repetition.x128, tcgen05.Pack.NONE)
    copy_atom_t2r = cute.make_copy_atom(op, ACC_DTYPE)
    tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tCtAcc)
    thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
    
    tTR_tAcc = thr_copy_t2r.partition_S(tCtAcc)
    tTR_gC = thr_copy_t2r.partition_D(tCgC)
    tTR_rAcc = cute.make_rmem_tensor(tTR_gC[None, None, None, None, 0, 0, 0].shape, ACC_DTYPE)
    tTR_rC = cute.make_rmem_tensor(tTR_gC[None, None, None, None, 0, 0, 0].shape, C_DTYPE)
    
    simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), C_DTYPE)
    tTR_gC = tTR_gC[(None, None, None, None, *mma_tile_coord_mnl)]
    
    acc_full = acc_consumer.wait_and_advance()
    cute.copy(tiled_copy_t2r, tTR_tAcc, tTR_rAcc)
    acc_vec = tTR_rAcc.load().to(C_DTYPE)
    tTR_rC.store(acc_vec)
    cute.copy(simt_atom, tTR_rC, tTR_gC)
    acc_full.release()
    
    cute.arch.barrier()
    tmem.free(acc_tmem_ptr)


@cute.jit
def launch_kernel(
    a_ptr: cute.Pointer,
    b_ptr: cute.Pointer,
    sfa_ptr: cute.Pointer,
    sfb_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))))
    b_tensor = cute.make_tensor(b_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)))
    
    sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(a_tensor.shape, SF_VEC_SIZE)
    sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)
    sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(b_tensor.shape, SF_VEC_SIZE)
    sfb_tensor = cute.make_tensor(sfb_ptr, 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,))
    atom_thr_size = cute.size(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)
    
    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),
        a_tensor, 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),
        b_tensor, 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),
        sfa_tensor, 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),
        sfb_tensor, sfb_smem_layout, MMA_TILER_MNK, tiled_mma, cluster_layout_vmnk.shape,
        internal_type=cutlass.Int16)
    
    num_tma_load_bytes = (
        cute.size_in_bytes(AB_DTYPE, a_smem_layout) +
        cute.size_in_bytes(AB_DTYPE, b_smem_layout) +
        cute.size_in_bytes(SF_DTYPE, sfa_smem_layout) +
        cute.size_in_bytes(SF_DTYPE, sfb_smem_layout)) * atom_thr_size
    
    grid = (cute.ceil_div(c_tensor.shape[0], MMA_TILER_MNK[0]),
            cute.ceil_div(c_tensor.shape[1], MMA_TILER_MNK[1]),
            c_tensor.shape[2])
    
    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,
        c_tensor, a_smem_layout_staged, b_smem_layout_staged,
        sfa_smem_layout_staged, sfb_smem_layout_staged, num_tma_load_bytes,
    ).launch(grid=grid, block=[THREADS_PER_CTA, 1, 1], cluster=(1, 1, 1))


_compiled_kernel_cache = None

def compile_kernel():
    global _compiled_kernel_cache
    if _compiled_kernel_cache is not None:
        return _compiled_kernel_cache
    
    a_ptr = make_ptr(AB_DTYPE, 0, cute.AddressSpace.gmem, assumed_align=16)
    b_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)
    sfb_ptr = make_ptr(SF_DTYPE, 0, cute.AddressSpace.gmem, assumed_align=32)
    
    _compiled_kernel_cache = cute.compile(launch_kernel, a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (0, 0, 0, 0))
    return _compiled_kernel_cache


def custom_kernel(data: input_t) -> output_t:
    a, b, _, _, sfa_permuted, sfb_permuted, c = data
    compiled_func = compile_kernel()
    
    m, k, l = a.shape
    n, _, _ = b.shape
    k = k * 2
    
    a_ptr = make_ptr(AB_DTYPE, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
    b_ptr = make_ptr(AB_DTYPE, b.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)
    sfb_ptr = make_ptr(SF_DTYPE, sfb_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32)
    
    compiled_func(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l))
    return c