submission 75389

"""
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

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>

using namespace nvcuda;

// Kernel constants
#define MMA_TILE_M 128
#define MMA_TILE_N 1
#define MMA_TILE_K 64
#define SF_VEC_SIZE 16

// 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);

    int sign = (val >> 3) & 0x1;
    int exp = (val >> 1) & 0x3;
    int mantissa = val & 0x1;

    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;
    }

    return sign ? -result : result;
}

// 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;
    }
}

// Convert FP32 to FP16
__device__ __forceinline__ __half fp32_to_fp16(float val) {
    return __float2half(val);
}

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;
    }

    float acc = 0.0f;

    int num_k_tiles = (k + MMA_TILE_K - 1) / MMA_TILE_K;

    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

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

        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;

            uint8_t a_packed = A[a_row_base + byte_idx * stride_a_k];
            float a_val = fp4_to_fp32(a_packed, bit_idx);

            uint8_t b_packed = B[b_row_base + byte_idx * stride_b_k];
            float b_val = fp4_to_fp32(b_packed, bit_idx);

            int k_tile_idx = k_idx / MMA_TILE_K;
            int k_block_16 = (k_idx % MMA_TILE_K) / SF_VEC_SIZE;

            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;

            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]);
            }

            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.
    
    Args:
        data: Tuple of (a, b, sfa_cpu, sfb_cpu, sfa_permuted, sfb_permuted, c)
    
    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
    )
    
    return c