submission 75402

- """
- Direct CUDA implementation of NVFP4 block-scaled GEMV kernel.
-  
- This implementation uses explicit CUDA code with shared memory and pipelining,
- replacing the CuTE high-level API.
- """
-  
  import torch
  from task import input_t, output_t
- import math
   
- # Kernel configuration
- MMA_TILE_M = 128
- MMA_TILE_N = 1
- MMA_TILE_K = 64
- THREADS_PER_BLOCK = 128
- SF_VEC_SIZE = 16  # 16 elements per scale factor
+ 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
+ mma_tiler_mnk = (128, 1, 64)  # 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
+ sf_vec_size = 16  # Scale factor block size (16 elements share one scale)
+ 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
   
   
- # CUDA kernel source code
- CUDA_KERNEL_SOURCE = """
- #include <cuda_fp16.h>
- #include <cuda_runtime.h>
- #include <mma.h>
- #include <math.h>
- #include <torch/extension.h>
+ # The CuTE kernel 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()
   
- using namespace nvcuda;
+     # 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])
+     gC_mnl = cute.local_tile(
+         mC_mnl, cute.slice_(mma_tiler_mnk, (None, None, 0)), (None, None, None)
+     )
   
- // Kernel constants
- #define MMA_TILE_M 128
- #define MMA_TILE_N 1
- #define MMA_TILE_K 64
- #define SF_VEC_SIZE 16
+     # Select output element corresponding to this thread and block indices
+     tCgC = gC_mnl[tidx, None, bidx, bidy, bidz]
+     tCgC = cute.make_tensor(tCgC.iterator, 1)
+     res = cute.zeros_like(tCgC, cutlass.Float32)
   
- // Convert FP4E2M1FN (packed, 2 per byte) to FP32
- // FP4E2M1FN format: 1 sign bit, 2 exponent bits, 1 mantissa bit
- // Bit layout: [sign(1)][exp(2)][mantissa(1)]
- // Exponent bias: 1 (per E2M1FN specification)
- __device__ __forceinline__ float fp4_to_fp32(uint8_t packed, int idx) {
-     uint8_t val = (idx == 0) ? (packed & 0x0F) : ((packed >> 4) & 0x0F);
+     # Get the number of k tiles (depth dimension) for the reduction loop
+     k_tile_cnt = gA_mkl.layout[3].shape
+     for k_tile in range(k_tile_cnt):
+         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]
   
-     int sign = (val >> 3) & 0x1;
-     int exp = (val >> 1) & 0x3;
-     int mantissa = val & 0x1;
+         # Load NVFP4 or 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()
   
-     float result;
-     if (exp == 0) {
-         // Subnormal: value = mantissa * 2^{-1} (only mantissa=1 yields 0.5)
-         result = mantissa ? 0.5f : 0.0f;
-     } else {
-         // Normalized: value = (1 + mantissa * 0.5) * 2^{exp - 1}
-         float base = mantissa ? 1.5f : 1.0f;
-         float scale = (exp == 1) ? 1.0f : (exp == 2) ? 2.0f : 4.0f;
-         result = base * scale;
-     }
+         # Convert loaded values to float32 for computation (FFMA)
+         a_val = a_val_nvfp4.to(cutlass.Float32)
+         b_val = b_val_nvfp4.to(cutlass.Float32)
+         sfa_val = sfa_val_fp8.to(cutlass.Float32)
+         sfb_val = sfb_val_fp8.to(cutlass.Float32)
   
-     return sign ? -result : result;
- }
+         # Create register memory tensors using make_rmem_tensor_like (available in CUTLASS 4.3.0+)
+         # This creates register memory tensors with the correct shape for the target dtype
+         tArA = cute.make_rmem_tensor_like(tAgA, cutlass.Float32)
+         tBrB = cute.make_rmem_tensor_like(tBgB, cutlass.Float32)
+         tArSFA = cute.make_rmem_tensor_like(tAgSFA, cutlass.Float32)
+         tBrSFB = cute.make_rmem_tensor_like(tBgSFB, cutlass.Float32)
+         
+         # Store the converted values to RMEM CuTe tensors
+         tArA.store(a_val)
+         tBrB.store(b_val)
+         tArSFA.store(sfa_val)
+         tBrSFB.store(sfb_val)
   
- // Convert FP8 to FP32 (simplified)
- __device__ __forceinline__ float fp8_to_fp32(uint8_t val) {
-     // Simplified FP8E4M3FN conversion
-     // In production, use proper conversion
-     int sign = (val >> 7) & 1;
-     int exp = (val >> 3) & 15;
-     int mantissa = val & 7;
-     
-     if (exp == 0) {
-         // Subnormal
-         float result = mantissa / 8.0f * powf(2.0f, -6.0f);
-         return sign ? -result : result;
-     } else {
-         float result = (1.0f + mantissa / 8.0f) * powf(2.0f, exp - 7.0f);
-         return sign ? -result : result;
-     }
- }
+         # Iterate over SF vector tiles and compute the scale&matmul accumulation
+         for i in cutlass.range_constexpr(mma_tiler_mnk[2]):
+             res += tArA[i] * tArSFA[i] * tBrB[i] * tBrSFB[i]
   
- // Convert FP32 to FP16
- __device__ __forceinline__ __half fp32_to_fp16(float val) {
-     return __float2half(val);
- }
+     # Store the final float16 result back to global memory
+     tCgC.store(res.to(cutlass.Float16))
+     return
   
- extern "C" __global__ void nvfp4_gemv_kernel_v1(
-     const uint8_t* __restrict__ A,           // [m, k, l] FP4 packed
-     const uint8_t* __restrict__ B,           // [128, k, l] FP4 packed  
-     const uint8_t* __restrict__ scale_A,     // Scale factors for A [32, 4, rest_m, 4, rest_k, l]
-     const uint8_t* __restrict__ scale_B,     // Scale factors for B [32, 4, rest_n, 4, rest_k, l]
-     __half* __restrict__ C,                  // [m, 1, l] FP16 output
-     int m, int k, int l,
-     int stride_a_m, int stride_a_k, int stride_a_l,
-     int stride_b_n, int stride_b_k, int stride_b_l,
-     int stride_c_m, int stride_c_l,
-     int rest_m, int rest_k                    // Scale factor layout dimensions
- ) {
-     int tid = threadIdx.x;
-     int bid = blockIdx.x;
-     int batch_idx = blockIdx.z;
   
-     int output_row = bid * blockDim.x + tid;
-     if (output_row >= m) {
-         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))
+     )
+     # 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)
   
-     float acc = 0.0f;
+     sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(b_tensor.shape, sf_vec_size)
+     sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout)
   
-     int num_k_tiles = (k + MMA_TILE_K - 1) / MMA_TILE_K;
+     # Compute grid dimensions
+     # Grid is (M_blocks, 1, L) where:
+     # - M_blocks = ceil(M / 128) to cover all output rows
+     # - L = batch size
+     grid = (
+         cute.ceil_div(c_tensor.shape[0], 128),
+         1,
+         c_tensor.shape[2],
+     )
   
-     int a_row_base = output_row * stride_a_m + batch_idx * stride_a_l;
-     int b_row_base = batch_idx * stride_b_l;  // Always take N=0 row for GEMV
+     # 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
   
-     int m_tile = output_row / 128;
-     int m_in_tile = output_row % 128;
-     int m_atom_32 = m_in_tile / 32;
-     int m_pos_32 = m_in_tile % 32;
   
-     for (int k_tile = 0; k_tile < num_k_tiles; ++k_tile) {
-         int k_start = k_tile * MMA_TILE_K;
-         int k_end = min(k_start + MMA_TILE_K, k);
+ # Global cache for compiled kernel
+ _compiled_kernel_cache = None
   
-         for (int i = 0; i < MMA_TILE_K; ++i) {
-             int k_idx = k_start + i;
-             if (k_idx >= k_end) {
-                 break;
-             }
   
-             int byte_idx = k_idx / 2;
-             int bit_idx = k_idx % 2;
+ # 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.
   
-             uint8_t a_packed = A[a_row_base + byte_idx * stride_a_k];
-             float a_val = fp4_to_fp32(a_packed, bit_idx);
+     Returns:
+         The compiled kernel function
+     """
+     global _compiled_kernel_cache
   
-             uint8_t b_packed = B[b_row_base + byte_idx * stride_b_k];
-             float b_val = fp4_to_fp32(b_packed, bit_idx);
+     if _compiled_kernel_cache is not None:
+         return _compiled_kernel_cache
   
-             int k_tile_idx = k_idx / MMA_TILE_K;
-             int k_block_16 = (k_idx % MMA_TILE_K) / SF_VEC_SIZE;
+     # 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)
   
-             int scale_idx = m_pos_32 * 4 * rest_m * 4 * rest_k * l +
-                            m_atom_32 * rest_m * 4 * rest_k * l +
-                            m_tile * 4 * rest_k * l +
-                            k_block_16 * rest_k * l +
-                            k_tile_idx * l +
-                            batch_idx;
+     # Compile the kernel
+     _compiled_kernel_cache = cute.compile(
+         my_kernel, a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (0, 0, 0, 0)
+     )
   
-             float scale_a = 0.0f;
-             if (scale_idx < (32 * 4 * rest_m * 4 * rest_k * l)) {
-                 scale_a = fp8_to_fp32(scale_A[scale_idx]);
-             }
+     return _compiled_kernel_cache
   
-             int scale_b_idx = k_block_16 * rest_k * l +
-                               k_tile_idx * l +
-                               batch_idx;
   
-             float scale_b = 0.0f;
-             if (scale_b_idx < (32 * 4 * 1 * 4 * rest_k * l)) {
-                 scale_b = fp8_to_fp32(scale_B[scale_b_idx]);
-             }
-  
-             acc += (a_val * scale_a) * (b_val * scale_b);
-         }
-     }
-  
-     int c_idx = output_row * stride_c_m + batch_idx * stride_c_l;
-     C[c_idx] = fp32_to_fp16(acc);
- }
-  
- // C++ wrapper function to launch the kernel
- // This function is defined in the CUDA source so it can use <<<>>> syntax
- // Accepts torch::Tensor and extracts pointers
- void launch_nvfp4_gemv_kernel(
-     torch::Tensor A,
-     torch::Tensor B,
-     torch::Tensor scale_A,
-     torch::Tensor scale_B,
-     torch::Tensor C,
-     int m, int k, int l,
-     int stride_a_m, int stride_a_k, int stride_a_l,
-     int stride_b_n, int stride_b_k, int stride_b_l,
-     int stride_c_m, int stride_c_l,
-     int rest_m, int rest_k,
-     int grid_x, int grid_y, int grid_z,
-     int block_x, int block_y, int block_z
- ) {
-     // Extract raw pointers from torch tensors
-     const uint8_t* A_ptr = static_cast<const uint8_t*>(A.data_ptr());
-     const uint8_t* B_ptr = static_cast<const uint8_t*>(B.data_ptr());
-     const uint8_t* scale_A_ptr = static_cast<const uint8_t*>(scale_A.data_ptr());
-     const uint8_t* scale_B_ptr = static_cast<const uint8_t*>(scale_B.data_ptr());
-     __half* C_ptr = static_cast<__half*>(C.data_ptr());
-     
-     dim3 grid(grid_x, grid_y, grid_z);
-     dim3 block(block_x, block_y, block_z);
-     nvfp4_gemv_kernel_v1<<<grid, block>>>(
-         A_ptr, B_ptr, scale_A_ptr, scale_B_ptr, C_ptr,
-         m, k, l,
-         stride_a_m, stride_a_k, stride_a_l,
-         stride_b_n, stride_b_k, stride_b_l,
-         stride_c_m, stride_c_l,
-         rest_m, rest_k
-     );
-     cudaError_t err = cudaGetLastError();
-     if (err != cudaSuccess) {
-         // Error handling - in production you might want to throw an exception
-     }
-     cudaDeviceSynchronize();
- }
- """
-  
-  
- # Compile and cache the kernel
- _kernel_cache = None
- _kernel_compile_failed = False
- _in_fallback = False  # Guard to prevent recursion
-  
- def get_compiled_kernel():
-     """Compile the CUDA kernel and return the function."""
-     global _kernel_cache, _kernel_compile_failed
-     
-     if _kernel_compile_failed:
-         return None
-     
-     if _kernel_cache is not None:
-         return _kernel_cache
-     
-     try:
-         from torch.utils.cpp_extension import load_inline
-         
-         # C++ wrapper code - use torch::Tensor for automatic conversion
-         # load_inline will automatically generate PYBIND11_MODULE when using functions=
-         cpp_wrapper = """
- #include <torch/extension.h>
- #include <cuda_runtime.h>
- #include <cuda_fp16.h>
-  
- void launch_nvfp4_gemv_kernel(
-     torch::Tensor A,
-     torch::Tensor B,
-     torch::Tensor scale_A,
-     torch::Tensor scale_B,
-     torch::Tensor C,
-     int m, int k, int l,
-     int stride_a_m, int stride_a_k, int stride_a_l,
-     int stride_b_n, int stride_b_k, int stride_b_l,
-     int stride_c_m, int stride_c_l,
-     int rest_m, int rest_k,
-     int grid_x, int grid_y, int grid_z,
-     int block_x, int block_y, int block_z
- );
- """
-         
-         # Define the kernel
-         # load_inline automatically generates PYBIND11_MODULE when using functions=
-         kernel_module = load_inline(
-             name='nvfp4_gemv_cuda',
-             cpp_sources=cpp_wrapper,
-             cuda_sources=CUDA_KERNEL_SOURCE,
-             functions=['launch_nvfp4_gemv_kernel'],  # Function name must match
-             verbose=False
-         )
-         
-         _kernel_cache = kernel_module.launch_nvfp4_gemv_kernel
-         return _kernel_cache
-     except Exception as e:
-         _kernel_compile_failed = True
-         print(f"Failed to compile CUDA kernel: {type(e).__name__}")
-         if hasattr(e, 'message'):
-             print(f"Error message: {e.message}")
-         # Don't fallback - raise error instead to avoid recursion issues
-         raise RuntimeError(f"CUDA kernel compilation failed: {e}") from e
-  
-  
  def custom_kernel(data: input_t) -> output_t:
      """
-     Direct CUDA implementation of NVFP4 block-scaled GEMV.
-     
+     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, sfa_permuted, sfb_permuted, c)
-     
+         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, kept for compatibility)
+             sfb_cpu: [1, k, l] - Scale factors in float8_e4m3fn (not used, kept for compatibility)
+             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
   
-     # Make sure scale tensors are contiguous in the MMA layout expected by the kernel
-     if not sfa_permuted.is_contiguous():
-         sfa_permuted = sfa_permuted.contiguous()
-     if not sfb_permuted.is_contiguous():
-         sfb_permuted = sfb_permuted.contiguous()
-     
-     # Get dimensions
-     m, k_packed, l = a.shape
-     k = k_packed * 2  # FP4 is packed, 2 elements per byte
-     
-     # Get the compiled kernel
-     # This will raise an error if compilation fails - no fallback to avoid recursion
-     kernel_func = get_compiled_kernel()
-     
-     if kernel_func is None:
-         raise RuntimeError("Kernel compilation failed and no fallback available")
-     
-     # Compute scale factor layout dimensions
-     rest_m = ceil_div(m, 128)
-     rest_k = ceil_div(k, 64)
-     
-     # Compute strides
-     stride_a_m = k_packed
-     stride_a_k = 1
-     stride_a_l = m * k_packed
-     
-     stride_b_n = k_packed
-     stride_b_k = 1
-     stride_b_l = 128 * k_packed
-     
-     stride_c_m = 1
-     stride_c_l = m
-     
-     # Launch kernel
-     grid_x = ceil_div(m, THREADS_PER_BLOCK)
-     grid_y = 1
-     grid_z = l
-     
-     # Pass torch tensors directly - the wrapper will extract pointers
-     kernel_func(
-         a,
-         b,
-         sfa_permuted,
-         sfb_permuted,
-         c,
-         m, k, l,
-         stride_a_m, stride_a_k, stride_a_l,
-         stride_b_n, stride_b_k, stride_b_l,
-         stride_c_m, stride_c_l,
-         rest_m, rest_k,
-         grid_x, grid_y, grid_z,
-         THREADS_PER_BLOCK, 1, 1
+     # 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 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.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