submission 111213

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

# Kernel configuration parameters
# Based on https://veitner.bearblog.dev/nvfp4-gemv-improved/
# Best config: M=64, K=128, threads=64x16 -> 34.9µs
mma_tiler_mnk = (64, 1, 128)  # Tile sizes for M, N, K dimensions
ab_dtype = cutlass.Float4E2M1FN  # FP4 data type for A and B
sf_dtype = cutlass.Float8E4M3FN  # FP8 data type for scale factors
c_dtype = cutlass.Float16  # FP16 output type
accum_dtype = cutlass.Float32  # FP32 accumulation
sf_vec_size = 16  # Scale factor block size (16 elements share one scale)

# 2D thread block configuration for K-parallelization
# Best config: 64x16 = 1024 threads -> 34.9µs
threads_per_m = 64   # Threads along M dimension
threads_per_k = 16   # Threads along K dimension
threads_per_cta = threads_per_m * threads_per_k  # Total: 1024 threads


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


# K-parallel GEMV with shared memory reduction (no atomics)
# Based on Simon Veitner's blog: https://veitner.bearblog.dev/nvfp4-gemv-improved/
@cute.kernel
def kernel(
    mA_mkl: cute.Tensor,
    mB_nkl: cute.Tensor,
    mSFA_mkl: cute.Tensor,
    mSFB_nkl: cute.Tensor,
    mC_mnl: cute.Tensor,
):
    # Get CUDA block and thread indices (2D thread block)
    bidx, bidy, bidz = cute.arch.block_idx()
    tidx, tidy, _ = cute.arch.thread_idx()

    # Extract the local tile for input matrix A
    gA_mkl = cute.local_tile(
        mA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None)
    )
    gSFA_mkl = cute.local_tile(
        mSFA_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)
    )
    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)
    )

    # Output element for this thread
    tCgC = gC_mnl[tidx, None, bidx, bidy, bidz]
    tCgC = cute.make_tensor(tCgC.iterator, 1)
    res = cute.zeros_like(tCgC, accum_dtype)

    # Allocate shared memory for K-parallel reduction
    # Shape: (threads_per_m, threads_per_k) with K-major stride for better reduction access
    allocator = cutlass.utils.SmemAllocator()
    smem_layout = cute.make_layout((threads_per_m, threads_per_k), stride=(threads_per_k, 1))
    shared_res = allocator.allocate_tensor(element_type=cutlass.Float32, layout=smem_layout)

    # Get number of K tiles
    k_tile_cnt = gA_mkl.layout[3].shape
    
    # Strided loop over K tiles - each tidy handles different K tiles
    # unroll_full=True tells compiler to fully unroll this loop for better ILP
    for k_tile in range(tidy, k_tile_cnt, threads_per_k, unroll_full=True):
        tAgA = gA_mkl[tidx, None, bidx, k_tile, bidz]
        tBgB = gB_nkl[0, None, bidy, k_tile, bidz]
        tAgSFA = gSFA_mkl[tidx, None, bidx, k_tile, bidz]
        tBgSFB = gSFB_nkl[0, None, bidy, k_tile, bidz]

        # Register tensors for A, B values
        tArA = cute.make_rmem_tensor_like(tAgA, c_dtype)
        tBrB = cute.make_rmem_tensor_like(tBgB, c_dtype)
        # Pre-computed product tensor A*B
        tABrAB = cute.make_rmem_tensor_like(tAgA, c_dtype)
        # Register tensors for scale factors
        tArSFA = cute.make_rmem_tensor_like(tAgSFA, accum_dtype)
        tBrSFB = cute.make_rmem_tensor_like(tBgSFB, accum_dtype)
        # Pre-computed product tensor SFA*SFB
        tSFrSF = cute.make_rmem_tensor_like(tAgSFA, accum_dtype)

        # Load and convert values
        a_val = tAgA.load().to(c_dtype)
        b_val = tBgB.load().to(c_dtype)
        sfa_val = tAgSFA.load().to(accum_dtype)
        sfb_val = tBgSFB.load().to(accum_dtype)

        tArA.store(a_val)
        tBrB.store(b_val)
        tArSFA.store(sfa_val)
        tBrSFB.store(sfb_val)

        # Pre-compute products outside inner loop (optimization from blog)
        tABrAB.store(tArA.load() * tBrB.load())
        tSFrSF.store(tArSFA.load() * tBrSFB.load())

        # Inner loop with pre-computed products (fewer operations per iteration)
        for i in cutlass.range_constexpr(mma_tiler_mnk[2]):
            res += tABrAB[i] * tSFrSF[i]

    # Store partial result to shared memory
    shared_res[(tidx, tidy)] = res[0]
    cute.arch.sync_threads()

    # Reduction: only tidy=0 threads write final result
    if tidy == 0:
        out = cute.zeros_like(tCgC, accum_dtype)
        # Sum across all K threads
        for i in cutlass.range_constexpr(threads_per_k):
            out += shared_res[(tidx, i)]

        # Store final result as FP16
        tCgC.store(out.to(cutlass.Float16))
    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,
):
    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 = 128
    b_tensor = cute.make_tensor(
        b_ptr,
        cute.make_layout(
            (n_padded_128, cute.assume(k, 32), l),
            stride=(cute.assume(k, 32), 1, cute.assume(n_padded_128 * 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)
    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)

    # Grid dimensions (M blocks based on threads_per_m = 64)
    grid = (
        cute.ceil_div(c_tensor.shape[0], 64),
        1,
        c_tensor.shape[2],
    )

    # Launch with 2D thread block: threads_per_m x threads_per_k
    kernel(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor).launch(
        grid=grid,
        block=[threads_per_m, threads_per_k, 1],
        cluster=(1, 1, 1),
    )
    return


# Global cache for compiled kernel
_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(
        my_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
    k = k * 2
    n = 1

    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