submission 81883

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, llvm

  

# Kernel configuration parameters

# M_tile=128 provides optimal balance: sufficient grid size with minimal atomic contention

# Testing showed M_tile=64 and M_tile=32 were slower due to increased atomic contention overhead

# Larger tiles reduce atomic operations and improve memory coalescing efficiency

mma_tiler_mnk = (128, 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.Float16 # FP16 output type

accum_dtype = cutlass.Float32 # Accumulator data type

sf_vec_size = 16 # Scale factor block size (16 elements share one scale)

# Thread count must match M_tile (128) due to tensor layout constraints

threads_per_cta = 128 # Number of threads per CUDA thread block

  

# Helper function for ceiling division

def ceil_div(a, b):

    return (a + b - 1) // b


# Atomic add for Float32 (required for parallel K tile accumulation)

@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()

    )

  

# The CuTe reference implementation for NVFP4 block-scaled GEMV

# Parallelized via extra blocks - each block handles one K tile

# Performance Analysis (from NSight Compute profiling):

# - Memory Throughput: 55-82% (good, but could be improved)

# - Compute Throughput: 47-69% (reasonable)

# - Main bottlenecks:

#   1. Low Occupancy (5-6%): Register pressure (82 regs/thread) limits concurrent warps

#   2. Small Grid: Grid size determined by problem dimensions, but underutilizes 148 SMs

#   3. Instruction Stalls: ~20 cycles "Stall No Instruction" - kernel is very short

# Optimizations applied:

# - Removed intermediate register tensors (tABrAB, tSFrSF) to reduce register pressure

# - Fused computation in loop to compute A*B and SFA*SFB on-the-fly

# - Direct store operations to minimize register usage

# Optimizations implemented:

# - Reduced register pressure: Removed intermediate tensors (tABrAB, tSFrSF), fused computation

# - Fused computation in loop: compute A*B and SFA*SFB on-the-fly to minimize register usage

# - Direct store operations to reduce register pressure

# Performance findings:

# - M_tile=128 provides best performance (tested 128, 64, 32 - 128 was fastest)

# - Smaller M_tile values increase atomic contention overhead more than they help occupancy

# - Register pressure optimizations remain beneficial regardless of tile size

@cute.kernel

def kernel(

    mA_mkl: cute.Tensor,

    mB_nkl: cute.Tensor,

    mSFA_mkl: cute.Tensor,

    mSFB_nkl: cute.Tensor,

    mC_mnl: cute.Tensor,  # Now Float32 accumulation buffer

):

    # Get CUDA block and thread indices

    bidx, bidy, bidz = cute.arch.block_idx()

    tidx, _, _ = cute.arch.thread_idx()



    # Extract the local tile for input matrix A (shape: [block_M, block_K, rest_M, rest_K, rest_L])

    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)

    # Here, block_M = (32, 4); block_K = (16, 4)

    gSFA_mkl = cute.local_tile(

        mSFA_mkl, cute.slice_(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])

    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)

    gSFB_nkl = cute.local_tile(

        mSFB_nkl, cute.slice_(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])

    # Tile M and N dimensions, but not L (batch) dimension

    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

    # Use bidy=0 for N dimension since N=1 for GEMV

    # bidz correctly selects the batch dimension

    tCgC = gC_mnl[tidx, None, bidx, 0, bidz]

    tCgC = cute.make_tensor(tCgC.iterator, 1)

    res = cute.zeros_like(tCgC, accum_dtype)



    # Load tiles for this K iteration (bidy selects the K tile)

    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]



    # Create register memory tensors (minimize register pressure)

    tArA = cute.make_rmem_tensor_like(tAgA, c_dtype)

    tBrB = cute.make_rmem_tensor_like(tBgB, c_dtype)

    tArSFA = cute.make_rmem_tensor_like(tAgSFA, accum_dtype)

    tBrSFB = cute.make_rmem_tensor_like(tBgSFB, accum_dtype)



    # Load and convert NVFP4/FP8 values from global memory

    a_val_nvfp4 = tAgA.load()

    b_val_nvfp4 = tBgB.load()

    sfa_val_fp8 = tAgSFA.load()

    sfb_val_fp8 = tBgSFB.load()



    # Convert loaded values and store directly (fused operations reduce register pressure)

    tArA.store(a_val_nvfp4.to(c_dtype))

    tBrB.store(b_val_nvfp4.to(c_dtype))

    tArSFA.store(sfa_val_fp8.to(accum_dtype))

    tBrSFB.store(sfb_val_fp8.to(accum_dtype))



    # Fused computation: compute A*B and SFA*SFB on-the-fly to reduce register pressure

    # Iterate over SF vector tiles and compute the scale&matmul accumulation

    # This performs: res += (A * SFA) * (B * SFB) = (A * B) * (SFA * SFB)

    for i in cutlass.range_constexpr(mma_tiler_mnk[2]):

        # Fused multiply-add: compute A*B and SFA*SFB on-the-fly

        res += (tArA[i] * tBrB[i]) * (tArSFA[i] * tBrSFB[i])



    # Atomic add to Float32 buffer (all K tiles accumulate to same output)

    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.

    Optimized grid launch configuration for B200.

    """

    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)),

        ),

    )

    # B tensor has n=128 padded size for proper alignment

    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 is Float32 accumulation buffer for atomic adds

    c_tensor = cute.make_tensor(

        c_ptr, cute.make_layout((cute.assume(m, 32), 1, l), stride=(1, 1, m))

    )

    # Convert scale factor tensors to MMA layout

    # The layout matches Tensor Core requirements: (((32, 4), REST_M), ((SF_K, 4), REST_K), (1, REST_L))

    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)

    

    # Compute grid dimensions

    # Grid is (M_blocks, K_blocks, L) where:

    # - M_blocks = ceil(M / 128) to cover all output rows

    # - K_blocks = ceil(K / 256) - each block handles one K tile

    # - L = batch size

    grid = (

        cute.ceil_div(c_tensor.shape[0], mma_tiler_mnk[0]),

        cute.ceil_div(a_tensor.shape[1], mma_tiler_mnk[2]),

        c_tensor.shape[2],

    )

    

    # Launch the CUDA kernel with optimized block size

    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

    # C pointer is Float32 for atomic accumulation

    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(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)

    

    # 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 optimized 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, sfa_permuted, sfb_permuted, 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 (not used, we use permuted)

    sfb_cpu: [1, k, l] - Scale factors in float8_e4m3fn (not used, we use permuted)

    sfa_permuted: [32, 4, rest_m, 4, rest_k, l] - Scale factors in MMA layout

    sfb_permuted: [32, 4, rest_n, 4, rest_k, l] - Scale factors in MMA layout

    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

    

    # Create Float32 accumulation buffer for atomic adds

    # All K tiles will atomically accumulate to this buffer

    # CuTe expects stride (1, 1, m) for shape (m, 1, l)

    # This means batches are stored contiguously: batch 0 at [0:m], batch 1 at [m:2m], etc.

    # Create a contiguous tensor of size m*l and view it with the correct stride

    c_fp32_flat = torch.zeros((m * l,), dtype=torch.float32, device=c.device)

    # View as (m, 1, l) with stride (1, 1, m) - this matches CuTe's expected layout

    c_fp32 = torch.as_strided(c_fp32_flat, size=(m, 1, l), stride=(1, 1, m))

    

    # 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_fp32_ptr = make_ptr(accum_dtype, c_fp32.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 (writes to Float32 buffer via atomic adds)

    compiled_func(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_fp32_ptr, (m, n, k, l))

    # Ensure all atomic operations complete before copying

    torch.cuda.synchronize()

    # Convert Float32 accumulator to Float16 output

    # c_fp32 has stride (1, 1, m), but c might have different stride

    # So we need to copy element by element or reshape

    # Since c_fp32 has the correct layout, we can directly convert and copy

    c.copy_(c_fp32.to(torch.float16))

    

    return c