submission 105643

# !POPCORN leaderboard nvfp4_gemv
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

from cutlass import Float32
from cutlass.cutlass_dsl import T, dsl_user_op
from cutlass._mlir.dialects import nvvm

# -----------------------------------------------------------------------------
# Kernel configuration
# -----------------------------------------------------------------------------

mma_tiler_mnk = (128, 1, 64)  # (M_tile, N_tile, K_tile)
ab_dtype = cutlass.Float4E2M1FN   # FP4 A/B
sf_dtype = cutlass.Float8E4M3FN   # FP8 SFA/SFB
c_dtype = cutlass.Float16         # final output type (PyTorch side)
accum_dtype = cutlass.Float32     # accumulation / atomic buffer type
sf_vec_size = 16                  # 16 values share one FP8 scale
threads_per_cta = 128             # threads per block
# Number of contiguous K tiles each CTA processes before issuing one atomic add.
# Reduces atomic traffic versus one atomic per K tile.
k_tiles_per_block = 4


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


# -----------------------------------------------------------------------------
# Atomic add helper (FP32 in global memory)
# -----------------------------------------------------------------------------

@dsl_user_op
def atomic_add_fp32(x: float | Float32, ptr: cute.Pointer, *, loc=None, ip=None) -> None:
    nvvm.atomicrmw(
        res=T.f32(),
        op=nvvm.AtomicOpKind.FADD,
        ptr=ptr.llvm_ptr,
        a=Float32(x).ir_value(),
    )


# -----------------------------------------------------------------------------
# First version: parallelize along K via extra blocks + global FP32 atomics
# -----------------------------------------------------------------------------

@cute.kernel
def kernel(
    mA_mkl: cute.Tensor,   # [M, K, L] in FP4
    mB_nkl: cute.Tensor,   # [128, K, L] padded N, FP4
    mSFA_mkl: cute.Tensor, # scales for A, FP8
    mSFB_nkl: cute.Tensor, # scales for B, FP8
    mC_mnl: cute.Tensor,   # [M, 1, L] FP32 accumulation buffer
):
    # block / thread indices
    bidx, bidy, bidz = cute.arch.block_idx()   # bidy = K-tile index
    tidx, _, _ = cute.arch.thread_idx()

    # Tile views
    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)
    )

    # Each thread computes a partial dot product for one output row and one K-tile.
    # N dimension is 1, so index 0 there.
    tCgC = gC_mnl[tidx, None, bidx, 0, bidz]
    tCgC = cute.make_tensor(tCgC.iterator, 1)

    # Local accumulation in FP32
    res = cute.zeros_like(tCgC, accum_dtype)

    # Process several contiguous K tiles in one CTA to amortize global atomics.
    k_tile_base = bidy * k_tiles_per_block
    k_tile_count = gA_mkl.layout[3].shape  # total K tiles

    for kt in cutlass.range_constexpr(k_tiles_per_block):
        k_tile = k_tile_base + kt
        # if k_tile >= k_tile_count:
            # Happens only on the tail group when k_tiles_per_block
            # doesn't divide the total number of K tiles.
            # continue

        # Select K-fragment for this tile
        tAgA   = gA_mkl[tidx, None, bidx, k_tile, bidz]
        tBgB   = gB_nkl[0,   None, 0,    k_tile, bidz]
        tAgSFA = gSFA_mkl[tidx, None, bidx, k_tile, bidz]
        tBgSFB = gSFB_nkl[0,   None, 0,    k_tile, bidz]

        # Register-memory fragments (A/B in half, scales in float)
        tArA   = cute.make_rmem_tensor_like(tAgA,   cutlass.Float16)
        tBrB   = cute.make_rmem_tensor_like(tBgB,   cutlass.Float16)
        tArSFA = cute.make_rmem_tensor_like(tAgSFA, accum_dtype)
        tBrSFB = cute.make_rmem_tensor_like(tBgSFB, accum_dtype)

        # Pre-multiplied fragments
        tABrAB = cute.make_rmem_tensor_like(tAgA,   cutlass.Float16)
        tSFrSF = cute.make_rmem_tensor_like(tAgSFA, accum_dtype)

        # Load and convert
        a_val_nvfp4 = tAgA.load()
        b_val_nvfp4 = tBgB.load()
        sfa_val_fp8 = tAgSFA.load()
        sfb_val_fp8 = tBgSFB.load()

        a_val = a_val_nvfp4.to(cutlass.Float16)
        b_val = b_val_nvfp4.to(cutlass.Float16)
        sfa_val = sfa_val_fp8.to(accum_dtype)
        sfb_val = sfb_val_fp8.to(accum_dtype)

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

        tABrAB.store(tArA.load() * tBrB.load())
        tSFrSF.store(tArSFA.load() * tBrSFB.load())

        # Reduce over SF-vector (K-fragment) inside this tile
        for i in cutlass.range_constexpr(mma_tiler_mnk[2]):
            res += tABrAB[i] * tSFrSF[i]

    # Cross-CTA reduction: atomic add into global FP32 C buffer
    atomic_add_fp32(res[0], tCgC.iterator)
    return


# -----------------------------------------------------------------------------
# JIT wrapper: build tensors + launch kernel
# -----------------------------------------------------------------------------

@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,     # FP32 buffer
    problem_size: tuple,
):
    m, _, k, l = problem_size

    # A: [M, K, L]
    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: padded N=128 to match torch._scaled_mm layout
    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 accumulation: [M, 1, L] in FP32
    c_tensor = cute.make_tensor(
        c_ptr,
        cute.make_layout((cute.assume(m, 32), 1, l), stride=(1, 1, m)),
    )

    # Scale factor layouts (same as before)
    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: split M across blocks; split K into groups of k_tiles_per_block per block
    k_tiles = cute.ceil_div(a_tensor.shape[1], mma_tiler_mnk[2])
    grid = (
        cute.ceil_div(c_tensor.shape[0], mma_tiler_mnk[0]),  # M tiles
        cute.ceil_div(k_tiles, k_tiles_per_block),           # grouped K tiles
        c_tensor.shape[2],                                  # L (batch)
    )

    kernel(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor).launch(
        grid=grid,
        block=[threads_per_cta, 1, 1],
        cluster=(1, 1, 1),
    )
    return


# -----------------------------------------------------------------------------
# Compile-once helper
# -----------------------------------------------------------------------------

_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)
    # Note: c_ptr is FP32 here
    c_ptr = make_ptr(accum_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


# -----------------------------------------------------------------------------
# Entry point used by the evaluation framework
# -----------------------------------------------------------------------------

def custom_kernel(data: input_t) -> output_t:
    """
    data:
        a: [m, k/2, l] (NVFP4x2 storage)
        b: [1, k/2, l]
        sfa_permuted: [...]
        sfb_permuted: [...]
        c: [m, 1, l] in Float16 (output buffer)
    """
    a, b, _, _, sfa_permuted, sfb_permuted, c = data

    compiled_func = compile_kernel()

    # Shapes in GEMV convention
    m, k_half, l = a.shape
    k = k_half * 2            # e2m1_x2 packing
    n = 1

    # FP32 accumulation buffer for C
    c_accum = torch.zeros_like(c, dtype=torch.float32)

    # Pointers
    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(accum_dtype, c_accum.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)

    # Run kernel (accumulate into c_accum)
    compiled_func(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l))

    # Convert FP32 accumulation buffer to FP16 result
    c.copy_(c_accum.to(dtype=c.dtype))
    return c