submission 115992

"""
FP4 Blockscaled GEMV - 16x16 configuration

Middle ground between v11 (8x32) and v13 (64x8):
- 16 rows per block
- 16 threads per row  
- 256 total threads
"""

import torch
from task import input_t, output_t

import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import make_ptr
import cutlass.utils.blockscaled_layout as blockscaled_utils
import cutlass.utils as cutlass_utils

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
ROWS_PER_BLOCK = 16
THREADS_PER_ROW = 16
THREADS = ROWS_PER_BLOCK * THREADS_PER_ROW  # 256
K_TILE = 64

mma_tiler_mnk = (ROWS_PER_BLOCK, 1, K_TILE)
ab_dtype = cutlass.Float4E2M1FN
sf_dtype = cutlass.Float8E4M3FN
c_dtype = cutlass.Float16
accum_dtype = cutlass.Float32
sf_vec_size = 16


@cute.kernel
def gemv_16x16(
    mA_mkl: cute.Tensor,
    mB_nkl: cute.Tensor,
    mSFA_mkl: cute.Tensor,
    mSFB_nkl: cute.Tensor,
    mC_mnl: cute.Tensor,
):
    """GEMV with 16x16 configuration."""
    bidx, bidy, bidz = cute.arch.block_idx()
    tid = cute.arch.thread_idx()[0]
    
    row_in_block = tid // THREADS_PER_ROW
    lane = tid % THREADS_PER_ROW
    
    global_row = bidx * ROWS_PER_BLOCK + row_in_block
    actual_m = mC_mnl.shape[0]
    
    gA = cute.local_tile(mA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None))
    gSFA = cute.local_tile(mSFA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None))
    gB = cute.local_tile(mB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None))
    gSFB = cute.local_tile(mSFB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None))
    gC = cute.local_tile(mC_mnl, cute.slice_(mma_tiler_mnk, (None, None, 0)), (None, None, None))
    
    smem_alloc = cutlass_utils.SmemAllocator()
    smem_partial = smem_alloc.allocate_tensor(accum_dtype, cute.make_layout((THREADS,)), 16)
    
    acc = accum_dtype(0.0)
    
    if global_row < actual_m:
        k_tiles = gA.layout[3].shape
        
        for kt in range(lane, k_tiles, THREADS_PER_ROW):
            tA = gA[row_in_block, None, bidx, kt, bidz]
            tSFA = gSFA[row_in_block, None, bidx, kt, bidz]
            tB = gB[0, None, bidy, kt, bidz]
            tSFB = gSFB[0, None, bidy, kt, bidz]
            
            rA = cute.make_rmem_tensor_like(tA, accum_dtype)
            rSFA = cute.make_rmem_tensor_like(tSFA, accum_dtype)
            rB = cute.make_rmem_tensor_like(tB, accum_dtype)
            rSFB = cute.make_rmem_tensor_like(tSFB, accum_dtype)
            
            rA.store(tA.load().to(accum_dtype))
            rSFA.store(tSFA.load().to(accum_dtype))
            rB.store(tB.load().to(accum_dtype))
            rSFB.store(tSFB.load().to(accum_dtype))
            
            for i in cutlass.range_constexpr(K_TILE):
                acc += rA[i] * rSFA[i] * rB[i] * rSFB[i]
    
    smem_partial[tid] = acc
    cute.arch.sync_threads()
    
    if global_row < actual_m and lane == 0:
        tC = gC[row_in_block, None, bidx, bidy, bidz]
        tC = cute.make_tensor(tC.iterator, 1)
        res = cute.zeros_like(tC, accum_dtype)
        
        base = row_in_block * THREADS_PER_ROW
        for i in cutlass.range_constexpr(THREADS_PER_ROW):
            res += smem_partial[base + i]
        
        tC.store(res.to(c_dtype))


@cute.jit  
def kernel_launcher(
    a_ptr: cute.Pointer,
    b_ptr: cute.Pointer,
    sfa_ptr: cute.Pointer,
    sfb_ptr: cute.Pointer,
    c_ptr: cute.Pointer,
    problem_size: tuple,
):
    m, _, 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)),
        ),
    )
    
    n_padded = 128
    b_tensor = cute.make_tensor(
        b_ptr,
        cute.make_layout(
            (n_padded, cute.assume(k, 32), l),
            stride=(cute.assume(k, 32), 1, cute.assume(n_padded * k, 32)),
        ),
    )
    
    c_tensor = cute.make_tensor(
        c_ptr,
        cute.make_layout((cute.assume(m, 32), 1, l), stride=(1, 1, m))
    )
    
    sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(a_tensor.shape, sf_vec_size)
    sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(b_tensor.shape, sf_vec_size)
    
    sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)
    sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout)
    
    m_blocks = cute.ceil_div(c_tensor.shape[0], ROWS_PER_BLOCK)
    grid = (m_blocks, 1, c_tensor.shape[2])
    
    gemv_16x16(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor).launch(
        grid=grid,
        block=[THREADS, 1, 1],
        cluster=(1, 1, 1),
    )


_cache = None

def compile_kernel():
    global _cache
    if _cache is not None:
        return _cache
    
    a_ptr = make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=128)
    b_ptr = make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=128)
    c_ptr = make_ptr(c_dtype, 0, cute.AddressSpace.gmem, assumed_align=128)
    sfa_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=128)
    sfb_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=128)
    
    _cache = cute.compile(kernel_launcher, a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (0, 0, 0, 0))
    return _cache


def custom_kernel(data: input_t) -> output_t:
    a, b, _, _, sfa_permuted, sfb_permuted, c = data
    func = compile_kernel()
    
    m, k, l = a.shape
    k = k * 2
    n = 1
    
    a_ptr = make_ptr(ab_dtype, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=128)
    b_ptr = make_ptr(ab_dtype, b.data_ptr(), cute.AddressSpace.gmem, assumed_align=128)
    c_ptr = make_ptr(c_dtype, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=128)
    sfa_ptr = make_ptr(sf_dtype, sfa_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=128)
    sfb_ptr = make_ptr(sf_dtype, sfb_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=128)
    
    func(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l))
    return c