submission 101652

#!POPCORN leaderboard nvfp4_gemv
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 utils
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.pipeline as pipeline
from cutlass.cute.nvgpu import cpasync
import torch

from cutlass.cute.tensor import TensorSSA

from cutlass import Float32, Float16, Int8, Int32
from cutlass.cutlass_dsl import T, dsl_user_op
from cutlass._mlir import ir
from cutlass._mlir.dialects import nvvm, arith, llvm, vector, builtin

# Convert 8 float8e4m3 values to 8 float16 values
@dsl_user_op
def cvt_f8e4m3x8_to_f16x8(src_vec8, *, loc=None, ip=None):
    # Split into two i32 values instead of using i64
    vec_i32x2_type = ir.VectorType.get([2], Int32.mlir_type, loc=loc)
    src_i32x2 = llvm.bitcast(vec_i32x2_type, src_vec8, loc=loc, ip=ip)
    src_lo = llvm.extractelement(src_i32x2, arith.constant(Int32.mlir_type, 0), loc=loc, ip=ip)
    src_hi = llvm.extractelement(src_i32x2, arith.constant(Int32.mlir_type, 1), loc=loc, ip=ip)
    
    # Process lower 4 bytes (4 fp8 values)
    rst_lo_i32x2 = llvm.inline_asm(
        llvm.StructType.get_literal([T.i32(), T.i32()]),
        [src_lo],
        """{\n\t
            .reg .b16 h0, h1;\n\t
            mov.b32 {h0, h1}, $2;\n\t
            cvt.rn.f16x2.e4m3x2 $0, h0;\n\t
            cvt.rn.f16x2.e4m3x2 $1, h1;\n\t
        }""",
        "=r,=r,r",
    )
    
    # Process upper 4 bytes (4 fp8 values)
    rst_hi_i32x2 = llvm.inline_asm(
        llvm.StructType.get_literal([T.i32(), T.i32()]),
        [src_hi],
        """{\n\t
            .reg .b16 h0, h1;\n\t
            mov.b32 {h0, h1}, $2;\n\t
            cvt.rn.f16x2.e4m3x2 $0, h0;\n\t
            cvt.rn.f16x2.e4m3x2 $1, h1;\n\t
        }""",
        "=r,=r,r",
    )
    
    res0 = llvm.extractvalue(T.i32(), rst_lo_i32x2, [0])
    res1 = llvm.extractvalue(T.i32(), rst_lo_i32x2, [1])
    res2 = llvm.extractvalue(T.i32(), rst_hi_i32x2, [0])
    res3 = llvm.extractvalue(T.i32(), rst_hi_i32x2, [1])
    
    vec_i32x4_type = ir.VectorType.get([4], Int32.mlir_type, loc=loc)
    vec_i32x4 = vector.from_elements(
        vec_i32x4_type, [res0, res1, res2, res3], loc=loc, ip=ip
    )
    vec_f16x8_type = ir.VectorType.get([8], Float16.mlir_type, loc=loc)
    vec_f16x8 = llvm.bitcast(vec_f16x8_type, vec_i32x4, loc=loc, ip=ip)
    return vec_f16x8

@dsl_user_op
def cvt_f8e4m3_f16_intrinsic(vec_f8e4m3, length, *, loc=None, ip=None):
    """
    Convert a vector of float8e4m3 to a vector of float16.
    :param vec_f8e4m3: The input vector of float8e4m3.
    :type vec_f8e4m3: 1D vector of float8e4m3
    :param length: The length of the input vector.
    :type length: int
    :return: The output 1D vector of float16 with the same length as the input vector.
    :rtype: 1D vector of float16
    """
    src_pos = 0
    vec_src_i8 = builtin.unrealized_conversion_cast(
        [ir.VectorType.get([length], Int8.mlir_type, loc=loc)],
        [vec_f8e4m3],
        loc=loc,
        ip=ip,
    )
    vec_i8x8_type = ir.VectorType.get([8], Int8.mlir_type, loc=loc)
    vec_dst_type = ir.VectorType.get([length], Float16.mlir_type, loc=loc)
    vec_dst = llvm.mlir_zero(vec_dst_type, loc=loc, ip=ip)

    num_vec8 = length // 8
    for _ in range(num_vec8):
        vec_f8e4m3x8 = vector.extract_strided_slice(
            vec_i8x8_type, vec_src_i8, [src_pos], [8], [1], loc=loc, ip=ip
        )
        vec_f16x8 = cvt_f8e4m3x8_to_f16x8(vec_f8e4m3x8, loc=loc, ip=ip)
        vec_dst = vector.insert_strided_slice(
            vec_f16x8, vec_dst, [src_pos], [1], loc=loc, ip=ip
        )
        src_pos += 8
        length -= 8

    return vec_dst


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

# Kernel configuration parameters
threads_per_cta = 128  # Number of threads per CUDA thread block
mma_tiler_mnk = (threads_per_cta, 1, 256)  # 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.Float32  # FP16 output type
sf_vec_size = 16  # Scale factor block size (16 elements share one scale)


# The CuTe reference implementation for NVFP4 block-scaled GEMV
@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
    bidx, bidy, bidz = cute.arch.block_idx()
    tidx, _, _ = cute.arch.thread_idx()

    problem_m = mA_mkl.shape[0]
    m_blocks_needed = cute.ceil_div(problem_m, mma_tiler_mnk[0])
    if bidx < m_blocks_needed:

        scale_mma_tiler_mnk = (mma_tiler_mnk[0], mma_tiler_mnk[1], mma_tiler_mnk[2]//sf_vec_size)
        # Extract the local tile for input matrix A (shape: [block_M, block_K, rest_M, rest_K, rest_L])
        # (32, 1, 256)
        gA_mkl = cute.local_tile(
            mA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None)
        )
        # Extract the local tile for scale factor tensor for A (same shape as gA_mkl)
        # (32, 1, 16)
        gSFA_mkl = cute.local_tile(
            mSFA_mkl, cute.slice_(scale_mma_tiler_mnk, (None, 0, None)), (None, None, None)
        )
        # Extract the local tile for input matrix B (shape: [block_N, block_K, rest_N, rest_K, rest_L])
        # (1, 1, 256)
        gB_nkl = cute.local_tile(
            mB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None)
        )
        # Extract the local tile for scale factor tensor for B (same shape as gB_nkl)
        # (1, 1, 16)
        gSFB_nkl = cute.local_tile(
            mSFB_nkl, cute.slice_(scale_mma_tiler_mnk, (0, None, None)), (None, None, None)
        )
        # Extract the local tile for output matrix C (shape: [block_M, block_N, rest_M, rest_N, rest_L])
        # (32, 1, 1)
        gC_mnl = cute.local_tile(
            mC_mnl, cute.slice_(mma_tiler_mnk, (None, None, 0)), (None, None, None)
        )

        # Select output element corresponding to this thread and block indices
        tCgC = gC_mnl[tidx, None, bidx, 0, bidz]
        tCgC = cute.make_tensor(tCgC.iterator, 1)
        res = cute.zeros_like(tCgC, cutlass.Float32)

        # Get the number of k tiles (depth dimension) for the reduction loop
        tAgA = gA_mkl[tidx, None, bidx, bidy, bidz]
        tBgB = gB_nkl[0, None, 0, bidy, bidz]
        tAgSFA = gSFA_mkl[tidx, None, bidx, bidy, bidz]
        tBgSFB = gSFB_nkl[0, None, 0, bidy, bidz]

        # Convert loaded values to float32 for computation (FFMA)
        a_val = tAgA.load().to(cutlass.Float16)
        b_val = tBgB.load().to(cutlass.Float16)
        #sfa_val = tAgSFA.load().to(cutlass.Float32)
        #sfb_val = tBgSFB.load().to(cutlass.Float32)
        sfa_val = cvt_f8e4m3_f16_intrinsic(tAgSFA.load(), sf_vec_size)
        sfb_val = cvt_f8e4m3_f16_intrinsic(tBgSFB.load(), sf_vec_size)

        mult = a_val * b_val #k values
        register_scale_raw = sfa_val * sfb_val #k//16 values
        register_scale = TensorSSA(register_scale_raw, sf_vec_size, cutlass.Float16)
        

        for block_idx in cutlass.range_constexpr(mma_tiler_mnk[2] // sf_vec_size):  # 32 iterations
            block_sum = cute.zeros_like(tCgC, cutlass.Float32)
            rng = block_idx*sf_vec_size
            for elem_idx in cutlass.range_constexpr(sf_vec_size):  # 16 iterations
                block_sum += mult[rng + elem_idx]
            res += block_sum * register_scale[block_idx]

        # Store the final float16 result back to global memory
        atomic_add_fp32(res[0], tCgC.iterator) 
    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,
):
    """
    Host-side JIT function to prepare tensors and launch GPU kernel.
    """
    m, _, k, l = problem_size
    # Create CuTe Tensor via pointer and 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)),
        ),
    )
    # We use n=128 to create the torch tensor to do fp4 computation via torch._scaled_mm
    # then copy torch tensor to cute tensor for cute customize kernel computation
    # therefore we need to ensure b_tensor has the right stride with this 128 padded size on n.
    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))
    )

    k_scaled = k // sf_vec_size
    shape_a_scales = (m, k_scaled, l)
    shape_b_scales = (b_tensor.shape[0], k_scaled, l)

    # This is the corrected layout for the pre-permuted, compact scale factors.
    # The M-mode retains the complex swizzle to match the permutation.
    # The K-mode is now a simple, contiguous layout for a tile of 4 elements.
    atom_shape_scaled = ((32, 4), 4)
    atom_stride_scaled = ((16, 4), 1)
    layout_scaled = cute.make_layout(atom_shape_scaled, stride=atom_stride_scaled)

    sfa_layout = cute.tile_to_shape(layout_scaled, shape_a_scales, (2, 1, 3))
    sfb_layout = cute.tile_to_shape(layout_scaled, shape_b_scales, (2, 1, 3))

    sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)
    sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout)

    # Compute grid dimensions
    # Grid is (M_blocks, 1, L) where:
    # - M_blocks = ceil(M / 128) to cover all output rows
    # - L = batch size
    
    m_blocks_needed = cute.ceil_div(m, mma_tiler_mnk[0])
    k_blocks_needed = cute.ceil_div(k, mma_tiler_mnk[2]) # k is in the Y-dim of the grid
    l_blocks_needed = l                                 # l is in the Z-dim of the grid

    sms = 148
    blocks_per_sm = 5
    WAVE_SIZE = sms * blocks_per_sm
    
    # Calculate total blocks and pad that, then resolve back to grid.x
    total_blocks_needed = m_blocks_needed * k_blocks_needed * l_blocks_needed
    padded_m_blocks = 0
    if total_blocks_needed > 0:
        padded_total_blocks = cute.ceil_div(total_blocks_needed, WAVE_SIZE) * WAVE_SIZE
        padded_m_blocks = cute.ceil_div(padded_total_blocks, k_blocks_needed * l_blocks_needed)

    grid = (
        padded_m_blocks,
        k_blocks_needed,
        l_blocks_needed,
    )

    # Launch the CUDA kernel
    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


# Global cache for compiled kernel
_compiled_kernel_cache = None


# This function is used to compile the kernel once and cache it and then allow users to
# run the kernel multiple times to get more accurate timing results.
def compile_kernel():
    """
    Compile the kernel once and cache it.
    This should be called before any timing measurements.

    Returns:
        The compiled kernel function
    """
    global _compiled_kernel_cache

    if _compiled_kernel_cache is not None:
        return _compiled_kernel_cache

    # Create CuTe pointers for A/B/C/SFA/SFB via torch tensor data pointer
    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 the kernel
    _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:
    """
    Execute the block-scaled GEMV kernel.

    This is the main entry point called by the evaluation framework.
    It converts PyTorch tensors to CuTe tensors, launches the kernel,
    and returns the result.

    Args:
        data: Tuple of (a, b, sfa_cpu, sfb_cpu, c) PyTorch tensors
            a: [m, k, l] - Input matrix in float4e2m1fn
            b: [1, k, l] - Input vector in float4e2m1fn
            sfa_cpu: [m, k, l] - Scale factors in float8_e4m3fn
            sfb_cpu: [1, k, l] - Scale factors in float8_e4m3fn
            sfa_permuted: [32, 4, rest_m, 4, rest_k, l] - Scale factors in float8_e4m3fn
            sfb_permuted: [32, 4, rest_n, 4, rest_k, l] - Scale factors in float8_e4m3fn
            c: [m, 1, l] - Output vector in float16

    Returns:
        Output tensor c with computed GEMV results
    """
    a, b, _, _, sfa_permuted, sfb_permuted, c = data

    # Ensure kernel is compiled (will use cached version if available)
    # To avoid the compilation overhead, we compile the kernel once and cache it.
    compiled_func = compile_kernel()

    # Get dimensions from MxKxL layout
    m, k, l = a.shape
    # Torch use e2m1_x2 data type, thus k is halved
    k = k * 2
    # GEMV N dimension is always 1
    n = 1

    c_accum = torch.zeros_like(c, dtype=torch.float32)
    # Create CuTe pointers for A/B/C/SFA/SFB via torch tensor data pointer
    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_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
    )

    # Execute the compiled kernel
    compiled_func(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l))

    return c_accum.to(torch.float16)