submission 116542

import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

# Kernel configuration parameters
sf_vec_size = 16

# Helper function for ceiling division
def ceil_div(a, b):
    return (a + b - 1) // b


# CUDA kernel source code for NVFP4 block-scaled GEMV
# Optimized with PTX intrinsics for FP4->FP16 conversion and warp-per-row design
cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cuda_fp8.h>

// Configuration
#define THREADS_PER_BLOCK 128
#define WARPS_PER_BLOCK (THREADS_PER_BLOCK / 32)

// Kernel parameters struct - passed by value (fits in registers)
struct KernelParams {
    const uint8_t* a;
    const uint8_t* b;
    const uint8_t* sfa;
    const uint8_t* sfb;
    __half* c;
    int m;
    int k;
    int l;
};

// Convert 2 packed FP8 E4M3 values to half2 using PTX
__device__ __forceinline__ half2 cvt_fp8x2_to_h2(uint16_t packed) {
    uint32_t result;
    asm volatile(
        "{\n\t"
        ".reg .b16 fp8x2;\n\t"
        "mov.b16 fp8x2, %1;\n\t"
        "cvt.rn.f16x2.e4m3x2 %0, fp8x2;\n\t"
        "}"
        : "=r"(result)
        : "h"(packed)
    );
    return *reinterpret_cast<half2*>(&result);
}

// Warp-level reduction using shuffle
__device__ __forceinline__ float warp_reduce_sum(float val) {
    val += __shfl_down_sync(0xffffffff, val, 16);
    val += __shfl_down_sync(0xffffffff, val, 8);
    val += __shfl_down_sync(0xffffffff, val, 4);
    val += __shfl_down_sync(0xffffffff, val, 2);
    val += __shfl_down_sync(0xffffffff, val, 1);
    return val;
}

// Convert 4 bytes (8 FP4 E2M1 values) to 4 half2 values using PTX
__device__ __forceinline__ void cvt_fp4x8_to_h2x4(uint32_t packed, half2& h0, half2& h1, half2& h2, half2& h3) {
    uint32_t f16x2_0, f16x2_1, f16x2_2, f16x2_3;

    asm volatile(
        "{\n\t"
        ".reg .b8 b0, b1, b2, b3;\n\t"
        "mov.b32 {b0, b1, b2, b3}, %4;\n\t"
        "cvt.rn.f16x2.e2m1x2 %0, b0;\n\t"
        "cvt.rn.f16x2.e2m1x2 %1, b1;\n\t"
        "cvt.rn.f16x2.e2m1x2 %2, b2;\n\t"
        "cvt.rn.f16x2.e2m1x2 %3, b3;\n\t"
        "}"
        : "=r"(f16x2_0), "=r"(f16x2_1), "=r"(f16x2_2), "=r"(f16x2_3)
        : "r"(packed)
    );

    h0 = *reinterpret_cast<half2*>(&f16x2_0);
    h1 = *reinterpret_cast<half2*>(&f16x2_1);
    h2 = *reinterpret_cast<half2*>(&f16x2_2);
    h3 = *reinterpret_cast<half2*>(&f16x2_3);
}


// ============================================================================
// Tunable GEMV kernel with WARPS_PER_ROW and M_TILE parameters
// WARPS_PER_ROW: Number of warps cooperating on each row (1, 2, 4, or 8)
// M_TILE: Number of rows each warp-group processes
// Processes 16 bytes (32 FP4 = 2 scale factors) per thread per iteration
// Simplified: Assumes M % 128 == 0 and K % 64 == 0
// ============================================================================
template <int WARPS_PER_ROW, int M_TILE>
__global__ void nvfp4_gemv_tunable_kernel(KernelParams p) {
    constexpr int THREADS_PER_ROW = WARPS_PER_ROW * 32;
    constexpr int ROW_GROUPS_PER_BLOCK = THREADS_PER_BLOCK / THREADS_PER_ROW;
    constexpr int ROWS_PER_BLOCK = ROW_GROUPS_PER_BLOCK * M_TILE;

    int tidx = threadIdx.x;
    int row_group = tidx / THREADS_PER_ROW;      // Which row-group within this block
    int tid_in_row = tidx % THREADS_PER_ROW;     // Thread ID within the row-group

    int row_tile = blockIdx.x;
    int batch = blockIdx.y;

    // Each row-group handles M_TILE rows
    int row_start = row_tile * ROWS_PER_BLOCK + row_group * M_TILE;

    int k_half = p.k / 2;
    int sf_k = p.k / 16;

    // B pointers (shared across all rows)
    const uint4* b_row_vec4 = reinterpret_cast<const uint4*>(p.b + batch * 128 * k_half);
    const uint16_t* sfb_row_vec2 = reinterpret_cast<const uint16_t*>(p.sfb + batch * 128 * sf_k);

    // Half2 accumulators for M_TILE rows - keeps partial sums in FP16
    // Each half2 holds {sum_for_scale_group_0, sum_for_scale_group_1}
    half2 acc_h2[M_TILE];
    #pragma unroll
    for (int r = 0; r < M_TILE; r++) {
        acc_h2[r] = __float2half2_rn(0.0f);
    }

    // Number of 16-byte chunks
    int num_vec4 = k_half / 16;

    // Threads stride by THREADS_PER_ROW across the K dimension
    for (int vec4_idx = tid_in_row; vec4_idx < num_vec4; vec4_idx += THREADS_PER_ROW) {
        // === PHASE 1: Issue all loads first ===
        uint4 a_vec4[M_TILE];
        uint16_t sfa_packed[M_TILE];
        #pragma unroll
        for (int r = 0; r < M_TILE; r++) {
            int row = row_start + r;
            const uint4* a_row_vec4 = reinterpret_cast<const uint4*>(p.a + row * k_half + batch * p.m * k_half);
            const uint16_t* sfa_row_vec2 = reinterpret_cast<const uint16_t*>(p.sfa + row * sf_k + batch * p.m * sf_k);
            a_vec4[r] = a_row_vec4[vec4_idx];
            sfa_packed[r] = sfa_row_vec2[vec4_idx];
        }

        uint4 b_vec4 = b_row_vec4[vec4_idx];
        uint16_t sfb_packed = sfb_row_vec2[vec4_idx];

        // === PHASE 2: Convert B (once, reused for all rows) ===
        half2 b_h2[16];
        cvt_fp4x8_to_h2x4(b_vec4.x, b_h2[0], b_h2[1], b_h2[2], b_h2[3]);
        cvt_fp4x8_to_h2x4(b_vec4.y, b_h2[4], b_h2[5], b_h2[6], b_h2[7]);
        cvt_fp4x8_to_h2x4(b_vec4.z, b_h2[8], b_h2[9], b_h2[10], b_h2[11]);
        cvt_fp4x8_to_h2x4(b_vec4.w, b_h2[12], b_h2[13], b_h2[14], b_h2[15]);

        // Convert B scales to half2: {scale_b0, scale_b1}
        half2 scale_b_h2 = cvt_fp8x2_to_h2(sfb_packed);

        // === PHASE 3: Convert A and compute for each row ===
        #pragma unroll
        for (int r = 0; r < M_TILE; r++) {
            half2 a_h2[16];
            cvt_fp4x8_to_h2x4(a_vec4[r].x, a_h2[0], a_h2[1], a_h2[2], a_h2[3]);
            cvt_fp4x8_to_h2x4(a_vec4[r].y, a_h2[4], a_h2[5], a_h2[6], a_h2[7]);
            cvt_fp4x8_to_h2x4(a_vec4[r].z, a_h2[8], a_h2[9], a_h2[10], a_h2[11]);
            cvt_fp4x8_to_h2x4(a_vec4[r].w, a_h2[12], a_h2[13], a_h2[14], a_h2[15]);

            // Convert A scales to half2 and compute combined scales in FP16
            half2 scale_a_h2 = cvt_fp8x2_to_h2(sfa_packed[r]);
            half2 combined_scale_h2 = __hmul2(scale_a_h2, scale_b_h2);

            // Dot product for first 16 FP4 elements (scale group 0)
            half2 sum_h2_0 = __float2half2_rn(0.0f);
            #pragma unroll
            for (int i = 0; i < 8; i++) {
                sum_h2_0 = __hfma2(a_h2[i], b_h2[i], sum_h2_0);
            }
            // Horizontal sum for scale group 0
            __half dot0 = __hadd(sum_h2_0.x, sum_h2_0.y);

            // Dot product for second 16 FP4 elements (scale group 1)
            half2 sum_h2_1 = __float2half2_rn(0.0f);
            #pragma unroll
            for (int i = 8; i < 16; i++) {
                sum_h2_1 = __hfma2(a_h2[i], b_h2[i], sum_h2_1);
            }
            __half dot1 = __hadd(sum_h2_1.x, sum_h2_1.y);

            // Pack dots into half2 and accumulate with scaling - stays in FP16
            half2 dots_h2 = make_half2(dot0, dot1);
            acc_h2[r] = __hfma2(dots_h2, combined_scale_h2, acc_h2[r]);
        }
    }

    // Convert half2 accumulators to float for reduction
    // Sum both scale groups: acc_h2[r].x + acc_h2[r].y
    float acc[M_TILE];
    #pragma unroll
    for (int r = 0; r < M_TILE; r++) {
        acc[r] = __half2float(acc_h2[r].x) + __half2float(acc_h2[r].y);
    }

    // === Reduction across THREADS_PER_ROW threads ===
    if constexpr (WARPS_PER_ROW == 1) {
        // Single warp: just warp reduction
        #pragma unroll
        for (int r = 0; r < M_TILE; r++) {
            float result = warp_reduce_sum(acc[r]);
            if (tid_in_row == 0) {
                p.c[row_start + r + batch * p.m] = __float2half(result);
            }
        }
    } else {
        // Multiple warps: need shared memory reduction
        __shared__ float smem[THREADS_PER_BLOCK];

        #pragma unroll
        for (int r = 0; r < M_TILE; r++) {
            // First do warp-level reduction
            float warp_sum = warp_reduce_sum(acc[r]);

            // Lane 0 of each warp writes to shared memory
            int warp_in_row = tid_in_row / 32;
            int lane_id = tid_in_row % 32;
            if (lane_id == 0) {
                smem[row_group * WARPS_PER_ROW + warp_in_row] = warp_sum;
            }
            __syncthreads();

            // First thread of each row-group reduces across warps
            if (tid_in_row == 0) {
                float result = 0.0f;
                #pragma unroll
                for (int w = 0; w < WARPS_PER_ROW; w++) {
                    result += smem[row_group * WARPS_PER_ROW + w];
                }
                p.c[row_start + r + batch * p.m] = __float2half(result);
            }
            __syncthreads();
        }
    }
}

// Launcher
template <int WARPS_PER_ROW, int M_TILE>
void launch_tunable_kernel(KernelParams p) {
    dim3 block(THREADS_PER_BLOCK);
    constexpr int THREADS_PER_ROW = WARPS_PER_ROW * 32;
    constexpr int ROW_GROUPS_PER_BLOCK = THREADS_PER_BLOCK / THREADS_PER_ROW;
    constexpr int ROWS_PER_BLOCK = ROW_GROUPS_PER_BLOCK * M_TILE;
    dim3 grid((p.m + ROWS_PER_BLOCK - 1) / ROWS_PER_BLOCK, p.l);
    nvfp4_gemv_tunable_kernel<WARPS_PER_ROW, M_TILE><<<grid, block>>>(p);
}

void nvfp4_gemv_cuda(std::vector<torch::Tensor> data) {
    // Unpack: data = (a, b, sfa, sfb, _, _, c)
    KernelParams p;
    p.a = reinterpret_cast<const uint8_t*>(data[0].data_ptr());
    p.b = reinterpret_cast<const uint8_t*>(data[1].data_ptr());
    p.sfa = reinterpret_cast<const uint8_t*>(data[2].data_ptr());
    p.sfb = reinterpret_cast<const uint8_t*>(data[3].data_ptr());
    p.c = reinterpret_cast<__half*>(data[6].data_ptr());
    p.m = data[0].size(0);
    p.k = data[0].size(1) * 2;
    p.l = data[0].size(2);

    launch_tunable_kernel<1, 1>(p);
}
"""

cpp_source = "void nvfp4_gemv_cuda(std::vector<torch::Tensor> data);"

# Compile at module load time
print("Compiling CUDA kernel...")
_module = load_inline(
    name='nvfp4_gemv',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['nvfp4_gemv_cuda'],
    verbose=True,
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3', '--expt-relaxed-constexpr', '-gencode', 'arch=compute_100a,code=sm_100a']
)
_kernel = _module.nvfp4_gemv_cuda


def custom_kernel(data: input_t) -> output_t:
    _kernel(data)
    return data[6]