submission 142859

from torch._higher_order_ops.torchbind import call_torchbind_fake
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.torch as cutlass_torch
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.utils.blockscaled_layout as blockscaled_utils
from cutlass.cute.runtime import make_ptr

# --- CONFIGURATION ---
# Tile sizes for M, N, K dimensions
mma_tiler_mnk = (128, 128, 256)  
# Shape of the K dimension for the MMA instruction
mma_inst_shape_k = 64
# FP4 data type for A and B
ab_dtype = cutlass.Float4E2M1FN  
# FP8 data type for scale factors
sf_dtype = cutlass.Float8E4M3FN  
# FP16 output type
c_dtype = cutlass.Float16  
# Scale factor block size (16 elements share one scale)
sf_vec_size = 16  
# Number of threads per CUDA thread block
threads_per_cta = 128  

# OPTIMIZATION: Increased stages for B200
num_acc_stage = 2
num_ab_stage = 6
# Total number of columns in tmem
num_tmem_alloc_cols = 512

# Helper function for ceiling division
def ceil_div(a, b):
    return (a + b - 1) // b

# The CuTe reference implementation for NVFP4 block-scaled GEMM
@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],
):
    """
    GPU device kernel performing the batched GEMM computation.
    """
    warp_idx = cute.arch.warp_idx()
    warp_idx = cute.arch.make_warp_uniform(warp_idx)
    
    # CRITICAL FIX: Unpack thread_idx to get scalar tidx for get_slice later
    tidx_tuple = cute.arch.thread_idx()
    tidx, _, _ = tidx_tuple

    #
    # Setup cta/thread coordinates
    #
    # Coords inside cluster
    bidx, bidy, bidz = cute.arch.block_idx()

    # Coords outside cluster
    cta_coord = (bidx, bidy, bidz)
    mma_tile_coord_mnl = (
        cta_coord[0] // cute.size(tiled_mma.thr_id.shape),
        cta_coord[1],
        cta_coord[2],
    )
    
    #
    # Define shared storage for kernel
    #
    @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)
    
    # Shared Memory Allocations with Stage Dimension
    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,
    )

    #
    # Local_tile partition global tensors
    #
    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])

    #
    # Partition global tensor for TiledMMA
    #
    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)

    #
    # Partition for TMA load
    #
    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)

    #
    # Partition shared/tensor memory tensor for TiledMMA
    #
    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)

    #
    # Slice to per mma tile index
    #
    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])]

    # Extract values for DSL capture
    ab_mbar_ptr = storage.ab_mbar_ptr.data_ptr()
    acc_mbar_ptr = storage.acc_mbar_ptr.data_ptr()
    tmem_holding_buf = storage.tmem_holding_buf

    # --- OOB CHECK WRAPPING ONLY RUNTIME LOGIC ---
    if (mma_tile_coord_mnl[0] * mma_tiler_mnk[0] < mC_mnl.shape[0] and
        mma_tile_coord_mnl[1] * mma_tiler_mnk[1] < mC_mnl.shape[1]):

        #
        # Initialize mainloop pipelines
        #
        ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
        ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1)
        
        # TMA Pipeline (Producer = TMA Engine, Consumer = SM Threads)
        ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create(
            barrier_storage=ab_mbar_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=acc_mbar_ptr,
            num_stages=num_acc_stage,
            producer_group=ab_pipeline_producer_group,
            consumer_group=pipeline.CooperativeGroup(
                pipeline.Agent.Thread,
                threads_per_cta,
            ),
        ).make_participants()

        #
        # Alloc TMEM
        #
        tmem_alloc_barrier = pipeline.NamedBarrier(
            barrier_id=1,
            num_threads=threads_per_cta,
        )
        tmem = utils.TmemAllocator(
            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)

        #
        # Make SFA/SFB tmem tensor
        #
        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)

        #
        # Partition for S2T copy of SFA/SFB
        #
        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(tidx)
        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(tidx)
        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)

        # --- PIPELINE LOGIC (Multistage) ---
        if warp_idx == 0:
            # Wait for accumulator buffer empty
            acc_empty = acc_producer.acquire_and_advance()
            tiled_mma.set(tcgen05.Field.ACCUMULATE, False)

            # 1. PROLOGUE: Pre-fetch stages
            # We fill the circular buffer with the first (num_stages - 1) tiles
            for stage in range(min(num_ab_stage - 1, k_tile_cnt)):
                ab_empty = ab_producer.acquire_and_advance()
                
                # Issue TMA loads
                cute.copy(tma_atom_a, tAgA[(None, stage)], tAsA[(None, ab_empty.index)], tma_bar_ptr=ab_empty.barrier)
                cute.copy(tma_atom_b, tBgB[(None, stage)], tBsB[(None, ab_empty.index)], tma_bar_ptr=ab_empty.barrier)
                cute.copy(tma_atom_sfa, tAgSFA[(None, stage)], tAsSFA[(None, ab_empty.index)], tma_bar_ptr=ab_empty.barrier)
                cute.copy(tma_atom_sfb, tBgSFB[(None, stage)], tBsSFB[(None, ab_empty.index)], tma_bar_ptr=ab_empty.barrier)

            # 2. MAIN LOOP: Compute K, Issue K + Stages
            for k_tile in range(k_tile_cnt):
                
                # A. ISSUE NEXT LOAD (if available)
                next_k_load = k_tile + num_ab_stage - 1
                if next_k_load < k_tile_cnt:
                    ab_empty = ab_producer.acquire_and_advance()
                    cute.copy(tma_atom_a, tAgA[(None, next_k_load)], tAsA[(None, ab_empty.index)], tma_bar_ptr=ab_empty.barrier)
                    cute.copy(tma_atom_b, tBgB[(None, next_k_load)], tBsB[(None, ab_empty.index)], tma_bar_ptr=ab_empty.barrier)
                    cute.copy(tma_atom_sfa, tAgSFA[(None, next_k_load)], tAsSFA[(None, ab_empty.index)], tma_bar_ptr=ab_empty.barrier)
                    cute.copy(tma_atom_sfb, tBgSFB[(None, next_k_load)], tBsSFB[(None, ab_empty.index)], tma_bar_ptr=ab_empty.barrier)

                # B. WAIT FOR CURRENT DATA
                ab_full = ab_consumer.wait_and_advance()

                # C. S2T COPY (Shared -> TMEM)
                s2t_stage_coord = (None, None, None, None, ab_full.index)
                tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord]
                tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord]
                cute.copy(tiled_copy_s2t_sfa, tCsSFA_compact_s2t_staged, tCtSFA_compact_s2t)
                cute.copy(tiled_copy_s2t_sfb, tCsSFB_compact_s2t_staged, tCtSFB_compact_s2t)

                # D. GEMM (Math)
                # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB
                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)
                    
                    # Enable accumulate on tCtAcc after first kblock
                    tiled_mma.set(tcgen05.Field.ACCUMULATE, True)

                # E. RELEASE BUFFER
                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)
        # tidx is now a scalar (0..127) thanks to unpacking
        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, cutlass.Float32)
        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)
    return


@cute.jit
def my_kernel(
    a_ptr: cute.Pointer,
    b_ptr: cute.Pointer,
    sfa_ptr: cute.Pointer,
    sfb_ptr: cute.Pointer,
    c_ptr: cute.Pointer,
    problem_size: tuple,
    cluster_shape: tuple, # New argument
):
    m, n, k, l = problem_size

    # Setup attributes that depend on gemm inputs
    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)

    # DYNAMIC CLUSTER LAYOUT: Use the cluster_shape passed in
    cluster_layout_vmnk = cute.tiled_divide(cute.make_layout(cluster_shape), (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)

    # Setup TMAs
    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,
    )

    # Compute TMA load bytes
    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

    grid_m = cute.ceil_div(c_tensor.shape[0], mma_tiler_mnk[0])
    grid_n = cute.ceil_div(c_tensor.shape[1], mma_tiler_mnk[1])

    # Pad grid to be multiple of cluster dims
    grid_m = cute.ceil_div(grid_m, cluster_shape[0]) * cluster_shape[0]
    grid_n = cute.ceil_div(grid_n, cluster_shape[1]) * cluster_shape[1]

    grid = (grid_m, grid_n, 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=cluster_shape, # Use the dynamic cluster shape
    )
    return


_compiled_kernel_cache = {} # Dictionary to map cluster shape to compiled kernel
def compile_kernel(cluster_shape):
    global _compiled_kernel_cache
    if cluster_shape in _compiled_kernel_cache:
        return _compiled_kernel_cache[cluster_shape]

    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)

    # Compile for the specific cluster shape
    compiled_k = cute.compile(my_kernel, a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (0, 0, 0, 0), cluster_shape)
    _compiled_kernel_cache[cluster_shape] = compiled_k
    return compiled_k


def custom_kernel(data: input_t) -> output_t:
    a, b, _, _, sfa_permuted, sfb_permuted, c = data
    
    m, k, l = a.shape
    n, _, _ = b.shape
    k = k * 2 

    # DYNAMIC CLUSTER SELECTION LOGIC
    # For large M, use 2x1x1 cluster (Fast Path)
    # For small M, use 1x1x1 cluster (Safe Path)
    if m >= 256:
        cluster_shape = (2, 1, 1)
    else:
        cluster_shape = (1, 1, 1)

    compiled_func = compile_kernel(cluster_shape)

    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), cluster_shape)
    return c