submission 81959
mdouglas · python · License unknown
Use it
Vendorable · source mirrored · license unknownView source →
No package. Vendor the mirrored source: 538 lines, June 9 Researcher Reciprocity License v1.0.
submission_cuda.py
curl "https://kernelindex.com/api/v1/implementations/kernelbot-nvfp4-gemv-81959?include=source"interfacepython
Compatibility
measured onNVIDIA B200
declared hardwareNVIDIA B200
architecturessm_100
dtypesfp8_e4m3, nvfp4
Benchmark evidence
1 measurement across 1 GPU, fastest first.
Operation / workload
Hardware
Latency
Rank
Observed
Reported · How evidence levels are derived →
Source and license
sourceavailable
revision digestsha256:bc14b6cc865b1818c09755b8ad5a55f488786d03c70e1b6017f847f32ddfaed9
license declaredunknown
license concludedunknown
authorsmdouglas
imported2026-08-15
Techniques
Extracted from the mirrored source by pattern, never inferred. Each row cites its line.
fp4
const uint8_t* __restrict__ a, // [M, K//2] packed FP4 (2 per byte)fp8
const __nv_fp8_e4m3* __restrict__ sfa, // [M, K//16, L] with strides (K_div_16, 1, M*K_div_16)shared-memory
extern __shared__ uint8_t smem[];vector-width = uint4
uint4 a_data;Kernel source
submission_cuda.py538 lines
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t
import torch
# CUDA kernel code for NVFP4 block-scaled GEMV
# Uses native Blackwell (sm_100a) hardware intrinsics for FP4/FP8 conversion
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cuda_fp4.h>
#include <cuda_fp8.h>
#include <ATen/cuda/Exceptions.h>
#include <ATen/cuda/CUDAContext.h>
// NVFP4 is 4-bit float (e2m1): 1 sign bit, 2 exponent bits, 1 mantissa bit
// Stored as 2 values per byte
// Scale factors are FP8 (e4m3) for every 16 FP4 values
// Warp reduction using shuffle - optimized for Blackwell
__device__ __forceinline__ float warp_reduce_sum(float val) {
#pragma unroll
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_down_sync(0xffffffff, val, offset);
}
return val;
}
// Batched kernel: processes all L batches in one launch for better efficiency.
// Each block handles multiple M rows, each warp handles one or more L batches.
// Inputs have native PyTorch strides from .permute() - K dimension has stride 1.
// Template parameter UseNestedLoops: true for large K (avoid div/mod), false for small K (less overhead)
// Template parameter LPerWarp: number of L batches each warp processes (1 for L=8, 2 for L=4)
template<bool UseNestedLoops, int LPerWarp>
__global__ void nvfp4_gemv_batched_kernel(
const uint8_t* __restrict__ a, // [M, K//2, L] with strides (K_half, 1, M*K_half)
const uint8_t* __restrict__ b, // [N, K//2, L] with strides (K_half, 1, N*K_half), N=128 padded
const __nv_fp8_e4m3* __restrict__ sfa, // [M, K//16, L] with strides (K_div_16, 1, M*K_div_16)
const __nv_fp8_e4m3* __restrict__ sfb, // [N, K//16, L] with strides (K_div_16, 1, N*K_div_16)
half* __restrict__ c, // [M, 1, L] output FP16
int M,
int K,
int L
) {
// Shared memory layout: B vectors [L, K/2], sfb [L, K/16]
extern __shared__ uint8_t smem[];
int K_half = K / 2;
int K_div_16 = K / 16;
uint8_t* sb = smem; // B vectors: L × K/2 bytes
__nv_fp8_e4m3* ssfb = reinterpret_cast<__nv_fp8_e4m3*>(sb + L * K_half); // Scale factors B: L × K/16
int tid = threadIdx.x;
int warp_id = threadIdx.x / 32;
int lane = threadIdx.x % 32;
const int N_padded = 128; // B is padded to 128 rows for torch._scaled_mm
// Cooperatively load all L B vectors into shared memory
// Template specialization: compile-time branch selection based on K size
// B original layout: b[n, k, l] at offset n*K_half + k + l*N_padded*K_half
// We only need n=0 (the actual vector, rest is padding)
if constexpr (UseNestedLoops) {
// Large K: nested loops to avoid expensive div/mod
for (int l = 0; l < L; l++) {
for (int k = tid; k < K_half; k += blockDim.x) {
sb[l * K_half + k] = b[k + l * N_padded * K_half];
}
}
} else {
// Small K: flat loop with div/mod has less overhead
for (int kl = tid; kl < K_half * L; kl += blockDim.x) {
int k = kl / L;
int l = kl % L;
sb[l * K_half + k] = b[k + l * N_padded * K_half];
}
}
// Cooperatively load scale factors for B
if constexpr (UseNestedLoops) {
// Large K: nested loops
for (int l = 0; l < L; l++) {
for (int k = tid; k < K_div_16; k += blockDim.x) {
ssfb[l * K_div_16 + k] = sfb[k + l * N_padded * K_div_16];
}
}
} else {
// Small K: flat loop
for (int kl = tid; kl < K_div_16 * L; kl += blockDim.x) {
int k = kl / L;
int l = kl % L;
ssfb[l * K_div_16 + k] = sfb[k + l * N_padded * K_div_16];
}
}
__syncthreads();
// Parallelize across L dimension with LPerWarp warps handling multiple L batches
const int WARPS_PER_BLOCK = blockDim.x / 32;
const int WARPS_PER_M_ROW = L / LPerWarp; // Fewer warps when each handles multiple L
const int M_ROWS_PER_BLOCK = WARPS_PER_BLOCK / WARPS_PER_M_ROW;
int m_base = blockIdx.x * M_ROWS_PER_BLOCK;
// Which M row and which L batch group?
int m_local = warp_id / WARPS_PER_M_ROW;
int l_group = warp_id % WARPS_PER_M_ROW;
int m = m_base + m_local;
if (m >= M) return;
// Each warp processes LPerWarp consecutive L batches
int l_base = l_group * LPerWarp;
// Process LPerWarp L batches per warp
for (int l_offset = 0; l_offset < LPerWarp; l_offset++) {
int l = l_base + l_offset;
if (l >= L) break;
float sum = 0.0f;
// K dimension has stride 1, enabling coalesced access
const int a_base = m * K_half + l * M * K_half;
const int sfa_base = m * K_div_16 + l * M * K_div_16;
const uint8_t* sb_row = &sb[l * K_half];
const __nv_fp8_e4m3* ssfb_row = &ssfb[l * K_div_16];
// Process 2 scale blocks per iteration for better ILP (16 bytes with uint4)
int num_scale_pairs = K_div_16 / 2;
for (int scale_pair = lane; scale_pair < num_scale_pairs; scale_pair += 32) {
int scale_block_0 = scale_pair * 2;
int scale_block_1 = scale_block_0 + 1;
// Load both scale factors
half scale_a0 = __ushort_as_half(__nv_cvt_fp8_to_halfraw(sfa[sfa_base + scale_block_0].__x, __NV_E4M3).x);
half scale_b0 = __ushort_as_half(__nv_cvt_fp8_to_halfraw(ssfb_row[scale_block_0].__x, __NV_E4M3).x);
half combined_scale0 = scale_a0 * scale_b0;
__half2 scale2_0 = __half2half2(combined_scale0);
half scale_a1 = __ushort_as_half(__nv_cvt_fp8_to_halfraw(sfa[sfa_base + scale_block_1].__x, __NV_E4M3).x);
half scale_b1 = __ushort_as_half(__nv_cvt_fp8_to_halfraw(ssfb_row[scale_block_1].__x, __NV_E4M3).x);
half combined_scale1 = scale_a1 * scale_b1;
__half2 scale2_1 = __half2half2(combined_scale1);
// Load 16 bytes at once using uint4
int k_byte_base = scale_block_0 * 8;
// A matrix: streaming access, use .cg cache hint (L2 only)
uint4 a_data;
const uint32_t* a_ptr = reinterpret_cast<const uint32_t*>(&a[a_base + k_byte_base]);
asm volatile("ld.global.cg.v4.u32 {%0,%1,%2,%3}, [%4];"
: "=r"(a_data.x), "=r"(a_data.y), "=r"(a_data.z), "=r"(a_data.w)
: "l"(a_ptr));
// B from shared memory: normal load
const uint4 b_data = *reinterpret_cast<const uint4*>(&sb_row[k_byte_base]);
const __nv_fp4x2_storage_t* a_fp4x2 = reinterpret_cast<const __nv_fp4x2_storage_t*>(&a_data);
const __nv_fp4x2_storage_t* b_fp4x2 = reinterpret_cast<const __nv_fp4x2_storage_t*>(&b_data);
// Process first 8 bytes with scale_0
__half2 local_sum_0 = __float2half2_rn(0.0f);
#pragma unroll
for (int i = 0; i < 8; i++) {
__half2 a_vals = __nv_cvt_fp4x2_to_halfraw2(a_fp4x2[i], __NV_E2M1);
__half2 b_vals = __nv_cvt_fp4x2_to_halfraw2(b_fp4x2[i], __NV_E2M1);
__half2 scaled = __hmul2(__hmul2(a_vals, b_vals), scale2_0);
local_sum_0 = __hadd2(local_sum_0, scaled);
}
sum += __half2float(__hadd(local_sum_0.x, local_sum_0.y));
// Process second 8 bytes with scale_1
__half2 local_sum_1 = __float2half2_rn(0.0f);
#pragma unroll
for (int i = 8; i < 16; i++) {
__half2 a_vals = __nv_cvt_fp4x2_to_halfraw2(a_fp4x2[i], __NV_E2M1);
__half2 b_vals = __nv_cvt_fp4x2_to_halfraw2(b_fp4x2[i], __NV_E2M1);
__half2 scaled = __hmul2(__hmul2(a_vals, b_vals), scale2_1);
local_sum_1 = __hadd2(local_sum_1, scaled);
}
sum += __half2float(__hadd(local_sum_1.x, local_sum_1.y));
}
sum = warp_reduce_sum(sum);
// c_ref has shape [M, 1, L] with strides (1, 1, M) from permute
// So c_ref[m, 0, l] is at linear offset: m + l*M
if (lane == 0) {
c[m + l * M] = __float2half(sum);
}
}
}
__global__ void nvfp4_gemv_kernel(
const uint8_t* __restrict__ a, // [M, K//2] packed FP4 (2 per byte)
const uint8_t* __restrict__ b, // [1, K//2] packed FP4
const __nv_fp8_e4m3* __restrict__ sfa, // [M, K//16] FP8 scale factors for A
const __nv_fp8_e4m3* __restrict__ sfb, // [1, K//16] FP8 scale factors for B
half* __restrict__ c, // [M, 1] output FP16
int M,
int K
) {
// 4 warps per M row - more iterations per thread for better ILP and latency hiding
// 8 M rows per block to maximize work per block
const int WARPS_PER_M_ROW = 4;
const int M_ROWS_PER_BLOCK = 8;
// Shared memory: B vector, scale factors B, and warp partial sums
extern __shared__ uint8_t smem[];
uint8_t* sb = smem;
__nv_fp8_e4m3* ssfb = reinterpret_cast<__nv_fp8_e4m3*>(sb + K/2);
float* warp_sums = reinterpret_cast<float*>(ssfb + K/16); // 32 warp partial sums
int tid = threadIdx.x;
int warp_id = tid / 32;
int lane = tid % 32;
int K_half = K / 2;
int K_div_16 = K / 16;
// Cooperatively load B vector into shared memory (vectorized)
int num_vec_loads = K_half / 16;
for (int i = tid; i < num_vec_loads; i += blockDim.x) {
reinterpret_cast<uint4*>(sb)[i] = reinterpret_cast<const uint4*>(b)[i];
}
int vec_bytes = num_vec_loads * 16;
for (int i = tid + vec_bytes; i < K_half; i += blockDim.x) {
sb[i] = b[i];
}
// Cooperatively load scale factors for B
for (int i = tid; i < K_div_16; i += blockDim.x) {
ssfb[i] = sfb[i];
}
__syncthreads();
// Each block processes 2 M rows
int m_base = blockIdx.x * M_ROWS_PER_BLOCK;
// Which M row and K chunk does this warp handle?
int m_local = warp_id / WARPS_PER_M_ROW;
int m = m_base + m_local;
if (m >= M) return;
int warp_in_m_group = warp_id % WARPS_PER_M_ROW;
// Process 2 scale blocks per iteration for better ILP
int scale_pairs_per_warp = (K_div_16 / WARPS_PER_M_ROW) / 2; // 1024 / 8 / 2 = 64 scale pairs per warp
int scale_pair_start = warp_in_m_group * scale_pairs_per_warp;
int scale_pair_end = scale_pair_start + scale_pairs_per_warp;
float sum = 0.0f;
// Loop over scale pairs - each iteration processes 16 bytes (2 scale blocks)
// With 8 warps: 64 scale pairs / 32 threads = 2 iterations per thread
for (int scale_pair = scale_pair_start + lane; scale_pair < scale_pair_end; scale_pair += 32) {
int scale_block_0 = scale_pair * 2;
int scale_block_1 = scale_block_0 + 1;
// Load both scale factors
half scale_a0 = __ushort_as_half(__nv_cvt_fp8_to_halfraw(sfa[m * K_div_16 + scale_block_0].__x, __NV_E4M3).x);
half scale_b0 = __ushort_as_half(__nv_cvt_fp8_to_halfraw(ssfb[scale_block_0].__x, __NV_E4M3).x);
half combined_scale0 = scale_a0 * scale_b0;
__half2 scale2_0 = __half2half2(combined_scale0);
half scale_a1 = __ushort_as_half(__nv_cvt_fp8_to_halfraw(sfa[m * K_div_16 + scale_block_1].__x, __NV_E4M3).x);
half scale_b1 = __ushort_as_half(__nv_cvt_fp8_to_halfraw(ssfb[scale_block_1].__x, __NV_E4M3).x);
half combined_scale1 = scale_a1 * scale_b1;
__half2 scale2_1 = __half2half2(combined_scale1);
int k_byte_base = scale_block_0 * 8;
// A matrix: streaming access, use .cg cache hint (L2 only)
uint4 a_data;
const uint32_t* a_ptr = reinterpret_cast<const uint32_t*>(&a[m * K_half + k_byte_base]);
asm volatile("ld.global.cg.v4.u32 {%0,%1,%2,%3}, [%4];"
: "=r"(a_data.x), "=r"(a_data.y), "=r"(a_data.z), "=r"(a_data.w)
: "l"(a_ptr));
// B from shared memory: normal load
const uint4 b_data = *reinterpret_cast<const uint4*>(&sb[k_byte_base]);
const __nv_fp4x2_storage_t* a_fp4x2 = reinterpret_cast<const __nv_fp4x2_storage_t*>(&a_data);
const __nv_fp4x2_storage_t* b_fp4x2 = reinterpret_cast<const __nv_fp4x2_storage_t*>(&b_data);
// Process first 8 bytes with scale_0 - accumulate in half2 first
__half2 local_sum_0 = __float2half2_rn(0.0f);
#pragma unroll
for (int i = 0; i < 8; i++) {
__half2 a_vals = __nv_cvt_fp4x2_to_halfraw2(a_fp4x2[i], __NV_E2M1);
__half2 b_vals = __nv_cvt_fp4x2_to_halfraw2(b_fp4x2[i], __NV_E2M1);
__half2 scaled = __hmul2(__hmul2(a_vals, b_vals), scale2_0);
local_sum_0 = __hadd2(local_sum_0, scaled);
}
sum += __half2float(__hadd(local_sum_0.x, local_sum_0.y));
// Process second 8 bytes with scale_1 - accumulate in half2 first
__half2 local_sum_1 = __float2half2_rn(0.0f);
#pragma unroll
for (int i = 8; i < 16; i++) {
__half2 a_vals = __nv_cvt_fp4x2_to_halfraw2(a_fp4x2[i], __NV_E2M1);
__half2 b_vals = __nv_cvt_fp4x2_to_halfraw2(b_fp4x2[i], __NV_E2M1);
__half2 scaled = __hmul2(__hmul2(a_vals, b_vals), scale2_1);
local_sum_1 = __hadd2(local_sum_1, scaled);
}
sum += __half2float(__hadd(local_sum_1.x, local_sum_1.y));
}
// Intra-warp reduction
sum = warp_reduce_sum(sum);
// Lane 0 of each warp writes its partial sum to shared memory
if (lane == 0) {
warp_sums[warp_id] = sum;
}
__syncthreads();
// Final reduction: each of the first M_ROWS_PER_BLOCK threads reduces one M row
if (tid < M_ROWS_PER_BLOCK) {
int m_write = m_base + tid;
if (m_write < M) {
// Reduce WARPS_PER_M_ROW partial sums for this M row
float final_sum = 0.0f;
int warp_start = tid * WARPS_PER_M_ROW;
#pragma unroll
for (int w = 0; w < WARPS_PER_M_ROW; w++) {
final_sum += warp_sums[warp_start + w];
}
c[m_write] = __float2half(final_sum);
}
}
}
void nvfp4_gemv_cuda(
torch::Tensor a,
torch::Tensor b,
torch::Tensor sfa,
torch::Tensor sfb,
torch::Tensor c,
int M,
int K
) {
// 4 warps per M row, 8 M rows per block
const int WARPS_PER_M_ROW = 4;
const int M_ROWS_PER_BLOCK = 8;
const int WARPS_PER_BLOCK = WARPS_PER_M_ROW * M_ROWS_PER_BLOCK; // 32
const int threads = WARPS_PER_BLOCK * 32; // 1024 threads
const int blocks = (M + M_ROWS_PER_BLOCK - 1) / M_ROWS_PER_BLOCK;
// Shared memory: B vector + sfb + warp partial sums
const int smem_size = K / 2 + K / 16 + WARPS_PER_BLOCK * sizeof(float);
// Get current CUDA stream from PyTorch
cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.device().index());
nvfp4_gemv_kernel<<<blocks, threads, smem_size, stream>>>(
a.data_ptr<uint8_t>(),
b.data_ptr<uint8_t>(),
reinterpret_cast<const __nv_fp8_e4m3*>(sfa.data_ptr<at::Float8_e4m3fn>()),
reinterpret_cast<const __nv_fp8_e4m3*>(sfb.data_ptr<at::Float8_e4m3fn>()),
reinterpret_cast<half*>(c.data_ptr<at::Half>()),
M, K
);
// Check for kernel launch errors
AT_CUDA_CHECK(cudaGetLastError());
}
void nvfp4_gemv_batched_cuda(
torch::Tensor a,
torch::Tensor b,
torch::Tensor sfa,
torch::Tensor sfb,
torch::Tensor c,
int M,
int K,
int L
) {
const int WARPS_PER_BLOCK = 32;
const int threads = WARPS_PER_BLOCK * 32; // 1024 threads
// Get current CUDA stream from PyTorch
cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.device().index());
const int K_half = K / 2;
const int K_div_16 = K / 16;
const int smem_size = L * K_half + L * K_div_16;
// For L=4 with small K: each warp handles 2 L batches for better ILP
// For L=8: each warp handles 1 L batch
if (L == 4 && K < 4096) {
// L=4, small K: 2 L batches per warp, 16 M rows per block
const int L_PER_WARP = 2;
const int WARPS_PER_M_ROW = L / L_PER_WARP; // 2
const int M_ROWS_PER_BLOCK = WARPS_PER_BLOCK / WARPS_PER_M_ROW; // 16
const int blocks = (M + M_ROWS_PER_BLOCK - 1) / M_ROWS_PER_BLOCK;
nvfp4_gemv_batched_kernel<false, 2><<<blocks, threads, smem_size, stream>>>(
a.data_ptr<uint8_t>(),
b.data_ptr<uint8_t>(),
reinterpret_cast<const __nv_fp8_e4m3*>(sfa.data_ptr<at::Float8_e4m3fn>()),
reinterpret_cast<const __nv_fp8_e4m3*>(sfb.data_ptr<at::Float8_e4m3fn>()),
reinterpret_cast<half*>(c.data_ptr<at::Half>()),
M, K, L
);
} else if (K >= 4096) {
// Large K: use nested loops, 1 L per warp
const int WARPS_PER_M_ROW = L;
const int M_ROWS_PER_BLOCK = WARPS_PER_BLOCK / WARPS_PER_M_ROW;
const int blocks = (M + M_ROWS_PER_BLOCK - 1) / M_ROWS_PER_BLOCK;
nvfp4_gemv_batched_kernel<true, 1><<<blocks, threads, smem_size, stream>>>(
a.data_ptr<uint8_t>(),
b.data_ptr<uint8_t>(),
reinterpret_cast<const __nv_fp8_e4m3*>(sfa.data_ptr<at::Float8_e4m3fn>()),
reinterpret_cast<const __nv_fp8_e4m3*>(sfb.data_ptr<at::Float8_e4m3fn>()),
reinterpret_cast<half*>(c.data_ptr<at::Half>()),
M, K, L
);
} else {
// Default: small K, 1 L per warp
const int WARPS_PER_M_ROW = L;
const int M_ROWS_PER_BLOCK = WARPS_PER_BLOCK / WARPS_PER_M_ROW;
const int blocks = (M + M_ROWS_PER_BLOCK - 1) / M_ROWS_PER_BLOCK;
nvfp4_gemv_batched_kernel<false, 1><<<blocks, threads, smem_size, stream>>>(
a.data_ptr<uint8_t>(),
b.data_ptr<uint8_t>(),
reinterpret_cast<const __nv_fp8_e4m3*>(sfa.data_ptr<at::Float8_e4m3fn>()),
reinterpret_cast<const __nv_fp8_e4m3*>(sfb.data_ptr<at::Float8_e4m3fn>()),
reinterpret_cast<half*>(c.data_ptr<at::Half>()),
M, K, L
);
}
// Check for kernel launch errors
AT_CUDA_CHECK(cudaGetLastError());
}
// Dispatch function - handles L=1 extraction and L=8 splitting in C++
void nvfp4_gemv_dispatch_cuda(
torch::Tensor a_ref, // [M, K//2, L] FP4
torch::Tensor b_ref, // [128, K//2, L] FP4
torch::Tensor sfa, // [M, K//16, L] FP8
torch::Tensor sfb, // [128, K//16, L] FP8
torch::Tensor c_ref // [M, 1, L] FP16
) {
int M = a_ref.size(0);
int K_half = a_ref.size(1);
int L = a_ref.size(2);
int K = K_half * 2;
if (L == 1) {
// L=1: Extract L=0 slice (all operations are views, no copying)
auto a_slice = a_ref.index({torch::indexing::Slice(), torch::indexing::Slice(), 0});
auto b_slice = b_ref.index({0, torch::indexing::Slice(), 0});
auto sfa_slice = sfa.index({torch::indexing::Slice(), torch::indexing::Slice(), 0});
auto sfb_slice = sfb.index({0, torch::indexing::Slice(), 0});
auto a_bytes = a_slice.view(torch::kUInt8).contiguous();
auto b_bytes = b_slice.view(torch::kUInt8).contiguous();
nvfp4_gemv_cuda(
a_bytes,
b_bytes,
sfa_slice.contiguous(),
sfb_slice.contiguous(),
c_ref,
M, K
);
} else {
// Other L: use batched kernel
auto a_bytes = a_ref.view(torch::kUInt8);
auto b_bytes = b_ref.view(torch::kUInt8);
nvfp4_gemv_batched_cuda(
a_bytes, b_bytes, sfa, sfb, c_ref,
M, K, L
);
}
}
"""
cpp_source = """
void nvfp4_gemv_dispatch_cuda(
torch::Tensor a,
torch::Tensor b,
torch::Tensor sfa,
torch::Tensor sfb,
torch::Tensor c
);
"""
# Compile the CUDA extension inline
nvfp4_gemv_module = load_inline(
name='nvfp4_gemv',
cpp_sources=[cpp_source],
cuda_sources=[cuda_source],
functions=['nvfp4_gemv_dispatch_cuda'],
verbose=True,
extra_cflags=[
'-O3',
'-std=c++17',
'-march=native',
],
extra_cuda_cflags=[
'-O3',
'--use_fast_math',
'--extra-device-vectorization',
'-arch=sm_100a',
'-std=c++17',
'-Xptxas', '-v',
'-U__CUDA_NO_HALF_OPERATORS__', # Enable half operators
'-U__CUDA_NO_HALF_CONVERSIONS__', # Enable half conversions
],
)
def custom_kernel(data: input_t) -> output_t:
"""
Custom CUDA implementation of NVFP4 block-scaled GEMV.
Dispatch logic now in C++ to minimize Python overhead.
"""
a_ref, b_ref, sfa, sfb, _, _, c_ref = data
# Single C++ call - dispatch logic handled in C++
nvfp4_gemv_module.nvfp4_gemv_dispatch_cuda(
a_ref,
b_ref,
sfa,
sfb,
c_ref
)
return c_ref
scrolls · 538 lines total
Source code from GPU Mode and the KernelBot dataset · June 9 Researcher Reciprocity License v1.0
Changes from previous submission
Against this author's previous submission submission 77328.
⋯ 26 unchanged lines}// Batched kernel: processes all L batches in one launch for better efficiency.- // Each block handles one M value, each warp within handles one L batch.+ // Each block handles multiple M rows, each warp handles one or more L batches.// Inputs have native PyTorch strides from .permute() - K dimension has stride 1.// Template parameter UseNestedLoops: true for large K (avoid div/mod), false for small K (less overhead)- template<bool UseNestedLoops>+ // Template parameter LPerWarp: number of L batches each warp processes (1 for L=8, 2 for L=4)+ template<bool UseNestedLoops, int LPerWarp>__global__ void nvfp4_gemv_batched_kernel(const uint8_t* __restrict__ a, // [M, K//2, L] with strides (K_half, 1, M*K_half)const uint8_t* __restrict__ b, // [N, K//2, L] with strides (K_half, 1, N*K_half), N=128 padded⋯ 57 unchanged lines__syncthreads();- // Parallelize across L dimension: each warp handles one (M, L) pair+ // Parallelize across L dimension with LPerWarp warps handling multiple L batchesconst int WARPS_PER_BLOCK = blockDim.x / 32;- const int WARPS_PER_M_ROW = L; // One warp per L batch+ const int WARPS_PER_M_ROW = L / LPerWarp; // Fewer warps when each handles multiple Lconst int M_ROWS_PER_BLOCK = WARPS_PER_BLOCK / WARPS_PER_M_ROW;int m_base = blockIdx.x * M_ROWS_PER_BLOCK;- // Which (M, L) pair does this warp handle?+ // Which M row and which L batch group?int m_local = warp_id / WARPS_PER_M_ROW;- int l = warp_id % WARPS_PER_M_ROW;+ int l_group = warp_id % WARPS_PER_M_ROW;int m = m_base + m_local;if (m >= M) return;- float sum = 0.0f;+ // Each warp processes LPerWarp consecutive L batches+ int l_base = l_group * LPerWarp;- // K dimension has stride 1, enabling coalesced access- const int a_base = m * K_half + l * M * K_half;- const int sfa_base = m * K_div_16 + l * M * K_div_16;- const uint8_t* sb_row = &sb[l * K_half];- const __nv_fp8_e4m3* ssfb_row = &ssfb[l * K_div_16];+ // Process LPerWarp L batches per warp+ for (int l_offset = 0; l_offset < LPerWarp; l_offset++) {+ int l = l_base + l_offset;+ if (l >= L) break;- // Process 2 scale blocks per iteration for better ILP (16 bytes with uint4)- int num_scale_pairs = K_div_16 / 2;- for (int scale_pair = lane; scale_pair < num_scale_pairs; scale_pair += 32) {- int scale_block_0 = scale_pair * 2;- int scale_block_1 = scale_block_0 + 1;+ float sum = 0.0f;- // Load both scale factors- half scale_a0 = __ushort_as_half(__nv_cvt_fp8_to_halfraw(sfa[sfa_base + scale_block_0].__x, __NV_E4M3).x);- half scale_b0 = __ushort_as_half(__nv_cvt_fp8_to_halfraw(ssfb_row[scale_block_0].__x, __NV_E4M3).x);- half combined_scale0 = scale_a0 * scale_b0;- __half2 scale2_0 = __half2half2(combined_scale0);+ // K dimension has stride 1, enabling coalesced access+ const int a_base = m * K_half + l * M * K_half;+ const int sfa_base = m * K_div_16 + l * M * K_div_16;+ const uint8_t* sb_row = &sb[l * K_half];+ const __nv_fp8_e4m3* ssfb_row = &ssfb[l * K_div_16];- half scale_a1 = __ushort_as_half(__nv_cvt_fp8_to_halfraw(sfa[sfa_base + scale_block_1].__x, __NV_E4M3).x);- half scale_b1 = __ushort_as_half(__nv_cvt_fp8_to_halfraw(ssfb_row[scale_block_1].__x, __NV_E4M3).x);- half combined_scale1 = scale_a1 * scale_b1;- __half2 scale2_1 = __half2half2(combined_scale1);+ // Process 2 scale blocks per iteration for better ILP (16 bytes with uint4)+ int num_scale_pairs = K_div_16 / 2;+ for (int scale_pair = lane; scale_pair < num_scale_pairs; scale_pair += 32) {+ int scale_block_0 = scale_pair * 2;+ int scale_block_1 = scale_block_0 + 1;- // Load 16 bytes at once using uint4- int k_byte_base = scale_block_0 * 8;- const uint4 a_data = *reinterpret_cast<const uint4*>(&a[a_base + k_byte_base]);- const uint4 b_data = *reinterpret_cast<const uint4*>(&sb_row[k_byte_base]);+ // Load both scale factors+ half scale_a0 = __ushort_as_half(__nv_cvt_fp8_to_halfraw(sfa[sfa_base + scale_block_0].__x, __NV_E4M3).x);+ half scale_b0 = __ushort_as_half(__nv_cvt_fp8_to_halfraw(ssfb_row[scale_block_0].__x, __NV_E4M3).x);+ half combined_scale0 = scale_a0 * scale_b0;+ __half2 scale2_0 = __half2half2(combined_scale0);- const __nv_fp4x2_storage_t* a_fp4x2 = reinterpret_cast<const __nv_fp4x2_storage_t*>(&a_data);- const __nv_fp4x2_storage_t* b_fp4x2 = reinterpret_cast<const __nv_fp4x2_storage_t*>(&b_data);+ half scale_a1 = __ushort_as_half(__nv_cvt_fp8_to_halfraw(sfa[sfa_base + scale_block_1].__x, __NV_E4M3).x);+ half scale_b1 = __ushort_as_half(__nv_cvt_fp8_to_halfraw(ssfb_row[scale_block_1].__x, __NV_E4M3).x);+ half combined_scale1 = scale_a1 * scale_b1;+ __half2 scale2_1 = __half2half2(combined_scale1);- // Process first 8 bytes with scale_0- __half2 local_sum_0 = __float2half2_rn(0.0f);- #pragma unroll- for (int i = 0; i < 8; i++) {- __half2 a_vals = __nv_cvt_fp4x2_to_halfraw2(a_fp4x2[i], __NV_E2M1);- __half2 b_vals = __nv_cvt_fp4x2_to_halfraw2(b_fp4x2[i], __NV_E2M1);- __half2 scaled = __hmul2(__hmul2(a_vals, b_vals), scale2_0);- local_sum_0 = __hadd2(local_sum_0, scaled);- }- sum += __half2float(__hadd(local_sum_0.x, local_sum_0.y));+ // Load 16 bytes at once using uint4+ int k_byte_base = scale_block_0 * 8;- // Process second 8 bytes with scale_1- __half2 local_sum_1 = __float2half2_rn(0.0f);- #pragma unroll- for (int i = 8; i < 16; i++) {- __half2 a_vals = __nv_cvt_fp4x2_to_halfraw2(a_fp4x2[i], __NV_E2M1);- __half2 b_vals = __nv_cvt_fp4x2_to_halfraw2(b_fp4x2[i], __NV_E2M1);- __half2 scaled = __hmul2(__hmul2(a_vals, b_vals), scale2_1);- local_sum_1 = __hadd2(local_sum_1, scaled);+ // A matrix: streaming access, use .cg cache hint (L2 only)+ uint4 a_data;+ const uint32_t* a_ptr = reinterpret_cast<const uint32_t*>(&a[a_base + k_byte_base]);+ asm volatile("ld.global.cg.v4.u32 {%0,%1,%2,%3}, [%4];"+ : "=r"(a_data.x), "=r"(a_data.y), "=r"(a_data.z), "=r"(a_data.w)+ : "l"(a_ptr));++ // B from shared memory: normal load+ const uint4 b_data = *reinterpret_cast<const uint4*>(&sb_row[k_byte_base]);++ const __nv_fp4x2_storage_t* a_fp4x2 = reinterpret_cast<const __nv_fp4x2_storage_t*>(&a_data);+ const __nv_fp4x2_storage_t* b_fp4x2 = reinterpret_cast<const __nv_fp4x2_storage_t*>(&b_data);++ // Process first 8 bytes with scale_0+ __half2 local_sum_0 = __float2half2_rn(0.0f);+ #pragma unroll+ for (int i = 0; i < 8; i++) {+ __half2 a_vals = __nv_cvt_fp4x2_to_halfraw2(a_fp4x2[i], __NV_E2M1);+ __half2 b_vals = __nv_cvt_fp4x2_to_halfraw2(b_fp4x2[i], __NV_E2M1);+ __half2 scaled = __hmul2(__hmul2(a_vals, b_vals), scale2_0);+ local_sum_0 = __hadd2(local_sum_0, scaled);+ }+ sum += __half2float(__hadd(local_sum_0.x, local_sum_0.y));++ // Process second 8 bytes with scale_1+ __half2 local_sum_1 = __float2half2_rn(0.0f);+ #pragma unroll+ for (int i = 8; i < 16; i++) {+ __half2 a_vals = __nv_cvt_fp4x2_to_halfraw2(a_fp4x2[i], __NV_E2M1);+ __half2 b_vals = __nv_cvt_fp4x2_to_halfraw2(b_fp4x2[i], __NV_E2M1);+ __half2 scaled = __hmul2(__hmul2(a_vals, b_vals), scale2_1);+ local_sum_1 = __hadd2(local_sum_1, scaled);+ }+ sum += __half2float(__hadd(local_sum_1.x, local_sum_1.y));}- sum += __half2float(__hadd(local_sum_1.x, local_sum_1.y));- }- sum = warp_reduce_sum(sum);+ sum = warp_reduce_sum(sum);- // c_ref has shape [M, 1, L] with strides (1, 1, M) from permute- // So c_ref[m, 0, l] is at linear offset: m + l*M- if (lane == 0) {- c[m + l * M] = __float2half(sum);+ // c_ref has shape [M, 1, L] with strides (1, 1, M) from permute+ // So c_ref[m, 0, l] is at linear offset: m + l*M+ if (lane == 0) {+ c[m + l * M] = __float2half(sum);+ }}}⋯ 9 unchanged lines// 4 warps per M row - more iterations per thread for better ILP and latency hiding// 8 M rows per block to maximize work per blockconst int WARPS_PER_M_ROW = 4;- const int WARPS_PER_BLOCK = blockDim.x / 32;const int M_ROWS_PER_BLOCK = 8;// Shared memory: B vector, scale factors B, and warp partial sumsextern __shared__ uint8_t smem[];uint8_t* sb = smem;__nv_fp8_e4m3* ssfb = reinterpret_cast<__nv_fp8_e4m3*>(sb + K/2);- float* warp_sums = reinterpret_cast<float*>(ssfb + K/16); // WARPS_PER_BLOCK partial sums+ float* warp_sums = reinterpret_cast<float*>(ssfb + K/16); // 32 warp partial sumsint tid = threadIdx.x;int warp_id = tid / 32;⋯ 56 unchanged linesint k_byte_base = scale_block_0 * 8;- // Load 16 bytes at once using uint4- const uint4 a_data = *reinterpret_cast<const uint4*>(&a[m * K_half + k_byte_base]);+ // A matrix: streaming access, use .cg cache hint (L2 only)+ uint4 a_data;+ const uint32_t* a_ptr = reinterpret_cast<const uint32_t*>(&a[m * K_half + k_byte_base]);+ asm volatile("ld.global.cg.v4.u32 {%0,%1,%2,%3}, [%4];"+ : "=r"(a_data.x), "=r"(a_data.y), "=r"(a_data.z), "=r"(a_data.w)+ : "l"(a_ptr));++ // B from shared memory: normal loadconst uint4 b_data = *reinterpret_cast<const uint4*>(&sb[k_byte_base]);const __nv_fp4x2_storage_t* a_fp4x2 = reinterpret_cast<const __nv_fp4x2_storage_t*>(&a_data);⋯ 93 unchanged linesint K,int L) {- // Parallelize across L: each warp handles one (M, L) pair- // Use 32 warps to process more M rows per block, reducing redundant B+sfb loadsconst int WARPS_PER_BLOCK = 32;- const int WARPS_PER_M_ROW = L;- const int M_ROWS_PER_BLOCK = WARPS_PER_BLOCK / WARPS_PER_M_ROW;const int threads = WARPS_PER_BLOCK * 32; // 1024 threads- const int blocks = (M + M_ROWS_PER_BLOCK - 1) / M_ROWS_PER_BLOCK;- // Shared memory: B vectors (L × K/2) and sfb (L × K/16)+ // Get current CUDA stream from PyTorch+ cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.device().index());+const int K_half = K / 2;const int K_div_16 = K / 16;const int smem_size = L * K_half + L * K_div_16;- // Get current CUDA stream from PyTorch- cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.device().index());+ // For L=4 with small K: each warp handles 2 L batches for better ILP+ // For L=8: each warp handles 1 L batch+ if (L == 4 && K < 4096) {+ // L=4, small K: 2 L batches per warp, 16 M rows per block+ const int L_PER_WARP = 2;+ const int WARPS_PER_M_ROW = L / L_PER_WARP; // 2+ const int M_ROWS_PER_BLOCK = WARPS_PER_BLOCK / WARPS_PER_M_ROW; // 16+ const int blocks = (M + M_ROWS_PER_BLOCK - 1) / M_ROWS_PER_BLOCK;- // Choose template specialization based on K size to avoid runtime branching- if (K >= 4096) {- // Large K: use nested loops to avoid div/mod- nvfp4_gemv_batched_kernel<true><<<blocks, threads, smem_size, stream>>>(+ nvfp4_gemv_batched_kernel<false, 2><<<blocks, threads, smem_size, stream>>>(a.data_ptr<uint8_t>(),b.data_ptr<uint8_t>(),reinterpret_cast<const __nv_fp8_e4m3*>(sfa.data_ptr<at::Float8_e4m3fn>()),⋯ 1 unchanged linesreinterpret_cast<half*>(c.data_ptr<at::Half>()),M, K, L);+ } else if (K >= 4096) {+ // Large K: use nested loops, 1 L per warp+ const int WARPS_PER_M_ROW = L;+ const int M_ROWS_PER_BLOCK = WARPS_PER_BLOCK / WARPS_PER_M_ROW;+ const int blocks = (M + M_ROWS_PER_BLOCK - 1) / M_ROWS_PER_BLOCK;++ nvfp4_gemv_batched_kernel<true, 1><<<blocks, threads, smem_size, stream>>>(+ a.data_ptr<uint8_t>(),+ b.data_ptr<uint8_t>(),+ reinterpret_cast<const __nv_fp8_e4m3*>(sfa.data_ptr<at::Float8_e4m3fn>()),+ reinterpret_cast<const __nv_fp8_e4m3*>(sfb.data_ptr<at::Float8_e4m3fn>()),+ reinterpret_cast<half*>(c.data_ptr<at::Half>()),+ M, K, L+ );} else {- // Small K: use flat loop with div/mod- nvfp4_gemv_batched_kernel<false><<<blocks, threads, smem_size, stream>>>(+ // Default: small K, 1 L per warp+ const int WARPS_PER_M_ROW = L;+ const int M_ROWS_PER_BLOCK = WARPS_PER_BLOCK / WARPS_PER_M_ROW;+ const int blocks = (M + M_ROWS_PER_BLOCK - 1) / M_ROWS_PER_BLOCK;++ nvfp4_gemv_batched_kernel<false, 1><<<blocks, threads, smem_size, stream>>>(a.data_ptr<uint8_t>(),b.data_ptr<uint8_t>(),reinterpret_cast<const __nv_fp8_e4m3*>(sfa.data_ptr<at::Float8_e4m3fn>()),⋯ 38 unchanged linesc_ref,M, K);- } else if (L == 8) {- // L=8: Split into two L=4 calls- auto a_bytes = a_ref.view(torch::kUInt8);- auto b_bytes = b_ref.view(torch::kUInt8);-- // PyTorch slicing preserves strides - no copying- nvfp4_gemv_batched_cuda(- a_bytes.index({torch::indexing::Slice(), torch::indexing::Slice(), torch::indexing::Slice(0,4)}),- b_bytes.index({torch::indexing::Slice(), torch::indexing::Slice(), torch::indexing::Slice(0,4)}),- sfa.index({torch::indexing::Slice(), torch::indexing::Slice(), torch::indexing::Slice(0,4)}),- sfb.index({torch::indexing::Slice(), torch::indexing::Slice(), torch::indexing::Slice(0,4)}),- c_ref.index({torch::indexing::Slice(), torch::indexing::Slice(), torch::indexing::Slice(0,4)}),- M, K, 4- );-- nvfp4_gemv_batched_cuda(- a_bytes.index({torch::indexing::Slice(), torch::indexing::Slice(), torch::indexing::Slice(4,8)}),- b_bytes.index({torch::indexing::Slice(), torch::indexing::Slice(), torch::indexing::Slice(4,8)}),- sfa.index({torch::indexing::Slice(), torch::indexing::Slice(), torch::indexing::Slice(4,8)}),- sfb.index({torch::indexing::Slice(), torch::indexing::Slice(), torch::indexing::Slice(4,8)}),- c_ref.index({torch::indexing::Slice(), torch::indexing::Slice(), torch::indexing::Slice(4,8)}),- M, K, 4- );} else {// Other L: use batched kernelauto a_bytes = a_ref.view(torch::kUInt8);⋯ 24 unchanged linescuda_sources=[cuda_source],functions=['nvfp4_gemv_dispatch_cuda'],verbose=True,+ extra_cflags=[+ '-O3',+ '-std=c++17',+ '-march=native',+ ],extra_cuda_cflags=['-O3','--use_fast_math',+ '--extra-device-vectorization','-arch=sm_100a',- '--std=c++17',+ '-std=c++17',+ '-Xptxas', '-v','-U__CUDA_NO_HALF_OPERATORS__', # Enable half operators'-U__CUDA_NO_HALF_CONVERSIONS__', # Enable half conversions],
scrolls · 321 diff lines total
Best evidence level for this revision: reported
JSON