use std::ffi::c_void;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use cudarc::driver::{LaunchConfig, PushKernelArg, sys::CUdeviceptr};
use onnx_runtime_ep_api::{
DeviceGraphResource, DevicePtr, DevicePtrMut, EpError, Kernel, KernelFactory, Result,
TensorMetadata, TensorMut, TensorView, WorkspaceRequirement, WorkspaceView,
};
use onnx_runtime_ir::{DataType, Node};
use crate::blas::{self, GemmDtype, GemmEpilogue, GemmEpilogueKind, GemmEx, GemmParams};
use crate::error::driver_err;
use crate::kernels::marlin_gemm;
use crate::runtime::{CudaRuntime, GraphDeviceAllocation, cuptr, raw_ptr};
const DEQUANT_MODULE: &str = "matmul_nbits_dequant_f32";
const DEQUANT_ENTRY: &str = "matmul_nbits_dequant_f32";
const GEMV_MODULE: &str = "matmul_nbits_gemv";
const GEMV_F32_ENTRY: &str = "matmul_nbits_gemv_f32";
const GEMV_INT8_F32_ENTRY: &str = "matmul_nbits_gemv_int8_f32";
const GEMV_INT4_F32_BLOCK128_ENTRY: &str = "matmul_nbits_gemv_int4_f32_block128";
const GEMV_INT8_F32_BLOCK128_ENTRY: &str = "matmul_nbits_gemv_int8_f32_block128";
const QUANTIZE_ACCURACY4_ENTRY: &str = "matmul_nbits_quantize_accuracy4_block32";
const GEMV_ACCURACY4_ENTRY: &str = "matmul_nbits_gemv_accuracy4_block32";
const GEMV_ACCURACY4_STAGE64_ENTRY: &str = "matmul_nbits_gemv_accuracy4_block32_stage64";
const QUANTIZE_ACCURACY4_BLOCKWISE_ENTRY: &str = "matmul_nbits_quantize_accuracy4_blockwise";
const GEMV_ACCURACY4_BLOCKWISE_ENTRY: &str = "matmul_nbits_gemv_accuracy4_blockwise";
const ACCURACY4_MODULE: &str = "matmul_nbits_accuracy4";
const ACCURACY4_ENTRY: &str = "matmul_nbits_accuracy4";
const BLOCK_THREADS: u32 = 256;
const GEMV_ACCURACY4_THREADS: u32 = 256;
const GEMV_ACCURACY4_COLUMNS_PER_BLOCK: usize = 8;
const GEMV_ACCURACY4_SHARED_BYTES: u32 = 32 * 32;
const GEMV_ACCURACY4_STAGE64_SHARED_BYTES: u32 = 64 * 32;
const GEMV_F16_MODULE: &str = "matmul_nbits_gemv_f16";
const GEMV_F16_ENTRY: &str = "matmul_nbits_gemv_f16";
const GEMV_INT8_F16_ENTRY: &str = "matmul_nbits_gemv_int8_f16";
const GEMV_INT8_F16_SPLITK_ENTRY: &str = "matmul_nbits_gemv_int8_f16_splitk";
const GEMV_INT8_F16_SPLITK: usize = 2;
const GEMM_F16_ENTRY: &str = "matmul_nbits_gemm_f16";
const GEMV_F16_GENERAL_BS_ENTRY: &str = "matmul_nbits_gemv_f16_general_bs";
const GEMV_F16_GENERAL_BS_SPLITK_ENTRY: &str = "matmul_nbits_gemv_f16_general_bs_splitk";
const GEMV_F16_GENERAL_BS_WIDE_ENTRY: &str = "matmul_nbits_gemv_f16_general_bs_wide";
const GEMV_F16_GENERAL_BS_SPLITK_WIDE_ENTRY: &str = "matmul_nbits_gemv_f16_general_bs_splitk_wide";
const GEMV_F16_GENERAL_BS_WIDE_MULTICOL_ENTRY: &str =
"matmul_nbits_gemv_f16_general_bs_wide_multicol";
const GEMV_F16_GENERAL_BS_WIDE_MULTICOL_FP16_ENTRY: &str =
"matmul_nbits_gemv_f16_general_bs_wide_multicol_fp16";
const GEMV_F16_GENERAL_BS_WIDE_MULTICOL_INTERLEAVED_ENTRY: &str =
"matmul_nbits_gemv_f16_general_bs_wide_multicol_interleaved";
const GEMV_F16_GENERAL_BS_SPLITK_WIDE_INTERLEAVED_ENTRY: &str =
"matmul_nbits_gemv_f16_general_bs_splitk_wide_interleaved";
const GEMV_F16_GENERAL_BS_SPLITK_WIDE_MULTICOL_ENTRY: &str =
"matmul_nbits_gemv_f16_general_bs_splitk_wide_multicol";
const GEMV_F16_GENERAL_BS_SPLITK_WIDE_MULTICOL_INTERLEAVED_ENTRY: &str =
"matmul_nbits_gemv_f16_general_bs_splitk_wide_multicol_interleaved";
const INTERLEAVE_INT4_ENTRY: &str = "matmul_nbits_interleave_int4";
const GEMV_F16_WIDE_MULTICOL_NC: usize = 4;
const GENERAL_BS_SPLITK: usize = 4;
const GENERAL_BS_SPLITK_MULTICOL: usize = 4;
const GEMM_F16_GENERAL_BS_ENTRY: &str = "matmul_nbits_gemm_f16_general_bs";
const GEMV_F16_SCALES_F16_ENTRY: &str = "matmul_nbits_gemv_f16_scales_f16";
const GEMV_F16_SCALES_F16_ZP_ENTRY: &str = "matmul_nbits_gemv_f16_scales_f16_zp";
const GEMV_F16_SCALES_F16_PIPE_ENTRY: &str = "matmul_nbits_gemv_f16_scales_f16_pipe";
const GEMV_F16_SCALES_F16_ZP_PIPE_ENTRY: &str = "matmul_nbits_gemv_f16_scales_f16_zp_pipe";
const GEMV_F16_SCALES_F16_SPLITK_ENTRY: &str = "matmul_nbits_gemv_f16_scales_f16_splitk";
const GEMV_F16_SCALES_F16_ZP_SPLITK_ENTRY: &str = "matmul_nbits_gemv_f16_scales_f16_zp_splitk";
const GEMV_F16_SCALES_F16_SPLITK_PF_ENTRY: &str = "matmul_nbits_gemv_f16_scales_f16_splitk_pf";
const GEMV_F16_SCALES_F16_ZP_SPLITK_PF_ENTRY: &str =
"matmul_nbits_gemv_f16_scales_f16_zp_splitk_pf";
const GEMV_F16_SCALES_F16_ZP_SPLITK: usize = 2;
const GEMV_F16_SCALES_F16_ZP_SPLITK8_ENTRY: &str = "matmul_nbits_gemv_f16_scales_f16_zp_splitk8";
const GEMV_F16_SCALES_F16_ZP_SPLITK8_PF_ENTRY: &str =
"matmul_nbits_gemv_f16_scales_f16_zp_splitk8_pf";
const GEMV_F16_SCALES_F16_ZP_SPLITK8: usize = 8;
const GEMV_F16_SCALES_F16_RMSNORM_ENTRY: &str = "matmul_nbits_gemv_f16_scales_f16_rmsnorm";
const GEMV_F16_SCALES_F16_RMSNORM_SPLITK_ENTRY: &str =
"matmul_nbits_gemv_f16_scales_f16_rmsnorm_splitk";
const GEMV_F16_SCALES_F16_RMSNORM_ZP_ENTRY: &str = "matmul_nbits_gemv_f16_scales_f16_rmsnorm_zp";
const GEMV_INT8_F16_SCALES_F16_RMSNORM_ENTRY: &str =
"matmul_nbits_gemv_int8_f16_scales_f16_rmsnorm";
const GEMV_INT8_F16_SCALES_F16_RMSNORM_ZP_ENTRY: &str =
"matmul_nbits_gemv_int8_f16_scales_f16_rmsnorm_zp";
const RMSNORM_PREFILL_ENTRY: &str = "matmul_nbits_rmsnorm_f16_warp_half4";
const RMSNORM_PREFILL_THREADS: u32 = 32;
const GEMV_F16_DOWN_ENTRY: &str = "matmul_nbits_gemv_f16_scales_f16_down";
const GEMV_F16_DOWN_C4_ENTRY: &str = "matmul_nbits_gemv_f16_scales_f16_down_c4";
const GEMV_F16_DOWN_C2_ENTRY: &str = "matmul_nbits_gemv_f16_scales_f16_down_c2";
const GEMM_F16_TILE: usize = 16;
const DECODE_GEMV_LOOP_MAX_M_DEFAULT: usize = 8;
const GEMV_F16_ACTIVATION_VECTOR_ELEMENTS: usize = 8;
fn decode_gemv_loop_rows_aligned(k: usize) -> bool {
k.is_multiple_of(GEMV_F16_ACTIVATION_VECTOR_ELEMENTS)
}
fn marlin_weight_inputs_are_constant(constant_inputs: &[bool], gate_up_swiglu: bool) -> bool {
if gate_up_swiglu {
constant_inputs.get(1) == Some(&true) && constant_inputs.get(3) == Some(&true)
} else {
constant_inputs.get(1) == Some(&true)
}
}
fn decode_gemv_loop_max_m() -> usize {
static CACHE: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
*CACHE.get_or_init(|| {
std::env::var("ONNX_GENAI_DECODE_GEMV_LOOP_MAX_M")
.ok()
.and_then(|value| value.trim().parse::<usize>().ok())
.filter(|&value| value >= 1)
.unwrap_or(DECODE_GEMV_LOOP_MAX_M_DEFAULT)
})
}
const GEMV_F16_SMALL_THREADS: u32 = 64;
const GEMV_F16_LARGE_THREADS: u32 = 256;
const GEMV_F16_SMALL_N_MAX: usize = 1152;
const GEMV_F16_DOWN_BLOCK_SIZE: usize = 32;
const GEMV_F16_DOWN_THREADS: u32 = 256;
const GEMV_F16_DOWN_COLUMNS_PER_BLOCK: usize = 8;
const GATE_UP_SWIGLU_ENTRY: &str = "matmul_nbits_gemv_f16_gate_up_swiglu";
const GATE_UP_DECOMPOSED_SWIGLU_ENTRY: &str = "matmul_nbits_gemv_f16_gate_up_decomposed_swiglu";
const GATE_UP_SWIGLU_RMSNORM_ENTRY: &str = "matmul_nbits_gemv_f16_gate_up_swiglu_rmsnorm";
const GATE_UP_DECOMPOSED_SWIGLU_RMSNORM_ENTRY: &str =
"matmul_nbits_gemv_f16_gate_up_decomposed_swiglu_rmsnorm";
const GATE_UP_SWIGLU_ZP_ENTRY: &str = "matmul_nbits_gemv_f16_gate_up_swiglu_zp";
const GATE_UP_DECOMPOSED_SWIGLU_ZP_ENTRY: &str =
"matmul_nbits_gemv_f16_gate_up_decomposed_swiglu_zp";
const GATE_UP_SWIGLU_RMSNORM_ZP_ENTRY: &str = "matmul_nbits_gemv_f16_gate_up_swiglu_rmsnorm_zp";
const GATE_UP_DECOMPOSED_SWIGLU_RMSNORM_ZP_ENTRY: &str =
"matmul_nbits_gemv_f16_gate_up_decomposed_swiglu_rmsnorm_zp";
const GATE_UP_SWIGLU_VEC_ENTRY: &str = "matmul_nbits_gemv_f16_gate_up_swiglu_vec";
const GATE_UP_DECOMPOSED_SWIGLU_VEC_ENTRY: &str =
"matmul_nbits_gemv_f16_gate_up_decomposed_swiglu_vec";
const GATE_UP_SWIGLU_RMSNORM_VEC_ENTRY: &str = "matmul_nbits_gemv_f16_gate_up_swiglu_rmsnorm_vec";
const GATE_UP_DECOMPOSED_SWIGLU_RMSNORM_VEC_ENTRY: &str =
"matmul_nbits_gemv_f16_gate_up_decomposed_swiglu_rmsnorm_vec";
const GATE_UP_SWIGLU_RMSNORM_VEC_OCC_ENTRY: &str =
"matmul_nbits_gemv_f16_gate_up_swiglu_rmsnorm_vec_occ";
const GATE_UP_DECOMPOSED_SWIGLU_RMSNORM_VEC_OCC_ENTRY: &str =
"matmul_nbits_gemv_f16_gate_up_decomposed_swiglu_rmsnorm_vec_occ";
const GATE_UP_SWIGLU_THREADS: u32 = 256;
const DEQUANT_SRC: &str = r#"
extern "C" __global__ void matmul_nbits_dequant_f32(
const unsigned char* packed,
const float* scales,
const unsigned char* zero_points,
const int* group_indices,
float* weight_kn,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int bits)
{
const long total = (long)k * n;
for (long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
idx < total; idx += (long)gridDim.x * blockDim.x) {
const int depth = (int)(idx / n);
const int output = (int)(idx % n);
const int block = depth / block_size;
const int within = depth - block * block_size;
const int bit_offset = within * bits;
const unsigned char byte =
packed[((long)output * k_blocks + block) * blob_size + bit_offset / 8];
const int mask = bits == 8 ? 255 : ((1 << bits) - 1);
const int quantized = (byte >> (bit_offset & 7)) & mask;
const int group = group_indices ? group_indices[depth] : block;
if (group < 0 || group >= k_blocks) {
weight_kn[idx] = 0.0f;
continue;
}
int zero_point = 1 << (bits - 1);
if (zero_points) {
const int zp_bit_offset = group * bits;
const unsigned char zp =
zero_points[(long)output * zp_row_bytes + zp_bit_offset / 8];
zero_point = (zp >> (zp_bit_offset & 7)) & mask;
}
weight_kn[idx] =
((float)quantized - (float)zero_point) * scales[(long)output * k_blocks + group];
}
}
"#;
const DEQUANT_F16_MODULE: &str = "matmul_nbits_dequant_f16";
const DEQUANT_F16_ENTRY: &str = "matmul_nbits_dequant_f16";
const DEQUANT_F16_SRC: &str = r#"
#include <cuda_fp16.h>
// grid.x covers K/8 eight-weight words, grid.y one output column each.
extern "C" __global__ void matmul_nbits_dequant_f16(
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
__half* __restrict__ weight_nk,
const int k,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int block_shift,
const int scales_fp16)
{
const int words = k >> 3;
const int w = (int)(blockIdx.x * blockDim.x + threadIdx.x);
if (w >= words) return;
const int out = (int)blockIdx.y;
const int depth0 = w << 3;
const int block = depth0 >> block_shift;
const long row_bytes = (long)k_blocks * blob_size;
const unsigned codes =
*reinterpret_cast<const unsigned*>(packed + (long)out * row_bytes + (long)w * 4);
const long scale_idx = (long)out * k_blocks + block;
const float scale = scales_fp16
? __half2float(reinterpret_cast<const __half*>(scales_raw)[scale_idx])
: reinterpret_cast<const float*>(scales_raw)[scale_idx];
float zero_point = 8.0f;
if (zero_points) {
const unsigned char byte = zero_points[(long)out * zp_row_bytes + (block >> 1)];
zero_point = (float)((block & 1) ? (byte >> 4) : (byte & 15));
}
__half2 out2[4];
#pragma unroll
for (int i = 0; i < 4; ++i) {
const float lo = (float)((codes >> (8 * i)) & 15) - zero_point;
const float hi = (float)((codes >> (8 * i + 4)) & 15) - zero_point;
out2[i] = __floats2half2_rn(lo * scale, hi * scale);
}
*reinterpret_cast<float4*>(weight_nk + (long)out * k + depth0) =
*reinterpret_cast<const float4*>(out2);
}
"#;
const GEMV_SRC: &str = r#"
__device__ __forceinline__ float warp_sum(float value)
{
for (int offset = 16; offset > 0; offset >>= 1) {
value += __shfl_down_sync(0xffffffffu, value, offset);
}
return value;
}
__device__ __forceinline__ float block_sum(float value)
{
__shared__ float warp_sums[32];
const int lane = threadIdx.x & 31;
const int warp = threadIdx.x >> 5;
value = warp_sum(value);
if (lane == 0) {
warp_sums[warp] = value;
}
__syncthreads();
value = threadIdx.x < ((blockDim.x + 31) >> 5) ? warp_sums[lane] : 0.0f;
return warp == 0 ? warp_sum(value) : 0.0f;
}
// Model-agnostic fp32-activation int4/int8 decode GEMV supporting any
// power-of-two block_size. The int4 path is bit-for-bit identical to the
// original (nibble unpack, symmetric default 8, per-block-nibble zero points);
// the int8 branch reads one byte per weight, uses a symmetric default of 128,
// and reads one whole-byte zero point per block. The tuned block-32 int8 entry
// (`matmul_nbits_gemv_int8_f32`) bakes in the block-32 geometry; block sizes
// other than 32 route here so the block index derives from the real block_size.
extern "C" __global__ void matmul_nbits_gemv_f32(
const float* activation,
const unsigned char* packed,
const float* scales,
const unsigned char* zero_points,
const float* bias,
float* output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int bits)
{
const int column = (int)blockIdx.x;
if (column >= n) {
return;
}
float value = 0.0f;
for (int depth = (int)threadIdx.x; depth < k; depth += (int)blockDim.x) {
const int block = depth / block_size;
const int within = depth - block * block_size;
const long blob_base = ((long)column * k_blocks + block) * blob_size;
int quantized;
int zero_point;
if (bits == 8) {
quantized = (int)packed[blob_base + within];
zero_point =
zero_points ? (int)zero_points[(long)column * k_blocks + block] : 128;
} else {
const unsigned char byte = packed[blob_base + within / 2];
quantized = (within & 1) ? (byte >> 4) : (byte & 15);
zero_point = 8;
if (zero_points) {
const unsigned char zp =
zero_points[(long)column * zp_row_bytes + block / 2];
zero_point = (block & 1) ? (zp >> 4) : (zp & 15);
}
}
value += activation[depth] * ((float)quantized - (float)zero_point)
* scales[(long)column * k_blocks + block];
}
value = block_sum(value);
if (threadIdx.x == 0) {
output[column] = value + (bias ? bias[column] : 0.0f);
}
}
extern "C" __global__ void matmul_nbits_gemv_int8_f32(
const float* activation,
const unsigned char* packed,
const float* scales,
const unsigned char* zero_points,
const float* bias,
float* output,
const int k,
const int n,
const int k_blocks)
{
const int column = (int)blockIdx.x;
if (column >= n) {
return;
}
float value = 0.0f;
for (int depth = (int)threadIdx.x; depth < k; depth += (int)blockDim.x) {
const int block = depth >> 5;
const int within = depth & 31;
const int quantized =
(int)packed[((long)column * k_blocks + block) * 32 + within];
const int zero_point =
zero_points ? (int)zero_points[(long)column * k_blocks + block] : 128;
value += activation[depth] * ((float)quantized - (float)zero_point)
* scales[(long)column * k_blocks + block];
}
value = block_sum(value);
if (threadIdx.x == 0) {
output[column] = value + (bias ? bias[column] : 0.0f);
}
}
// Structurally-selected int8 / block_size == 128 / asymmetric fp32-activation
// decode GEMV. This is a specialization of `matmul_nbits_gemv_f32` for the
// generic block-128 int8 loop that dominates Qwen3-0.6B decode. It is
// BIT-FOR-BIT IDENTICAL to that kernel: each thread walks the same depth stride
// (grid-stride by blockDim), evaluates the same per-element expression
// `activation * ((float)quantized - (float)zero_point) * scale`, and reduces in
// the same block_sum order — only the address arithmetic is cheaper.
//
// With block_size == 128 the block index is a shift (`depth >> 7`) instead of an
// integer divide, and the packed-byte address collapses: for blob_size == 128,
// (column*k_blocks + block)*128 + (depth - block*128) == column*k_blocks*128 + depth,
// so `within` (the modulo) disappears entirely. The `column * k_blocks` base is
// hoisted once into `col_kb` for the scale and zero-point rows, and the runtime
// `bits == 8` branch of the generic kernel is gone. The `zero_points ? : 128`
// selection is retained (uniform across the CTA, effectively free) so the result
// stays identical whether or not zero points are present; dispatch restricts
// this entry to the asymmetric case in practice.
extern "C" __global__ void matmul_nbits_gemv_int8_f32_block128(
const float* activation,
const unsigned char* packed,
const float* scales,
const unsigned char* zero_points,
const float* bias,
float* output,
const int k,
const int n,
const int k_blocks)
{
const int column = (int)blockIdx.x;
if (column >= n) {
return;
}
const long col_kb = (long)column * k_blocks;
const long packed_row = col_kb << 7; // k_blocks * 128 bytes per weight row
const float* col_scales = scales + col_kb;
const unsigned char* col_zp = zero_points ? zero_points + col_kb : (const unsigned char*)0;
float value = 0.0f;
for (int depth = (int)threadIdx.x; depth < k; depth += (int)blockDim.x) {
const int block = depth >> 7;
const int quantized = (int)packed[packed_row + depth];
const int zero_point = col_zp ? (int)col_zp[block] : 128;
value += activation[depth] * ((float)quantized - (float)zero_point)
* col_scales[block];
}
value = block_sum(value);
if (threadIdx.x == 0) {
output[column] = value + (bias ? bias[column] : 0.0f);
}
}
// Int4 counterpart of the block-128 int8 specialization above. Each warp owns
// 32 consecutive depths, hence 16 aligned packed bytes. Four lanes load one
// aligned 32-bit word each and shuffle it to the eight lanes consuming its
// nibbles. This removes duplicate scalar byte loads without changing any
// thread's depth sequence or fp32 arithmetic.
extern "C" __global__ void matmul_nbits_gemv_int4_f32_block128(
const float* activation,
const unsigned char* packed,
const float* scales,
const unsigned char* zero_points,
const float* bias,
float* output,
const int k,
const int n,
const int k_blocks)
{
const int column = (int)blockIdx.x;
if (column >= n) {
return;
}
const int lane = (int)threadIdx.x & 31;
const long col_kb = (long)column * k_blocks;
const long packed_row = col_kb << 6; // k_blocks * 64 bytes per weight row
const float* col_scales = scales + col_kb;
const unsigned char* col_zp =
zero_points + (long)column * ((k_blocks + 1) >> 1);
float value = 0.0f;
for (int depth = (int)threadIdx.x; depth < k; depth += (int)blockDim.x) {
const int block = depth >> 7;
const int warp_depth = depth - lane;
const unsigned int* packed_words =
(const unsigned int*)(packed + packed_row + (warp_depth >> 1));
unsigned int packed_word = lane < 4 ? packed_words[lane] : 0;
packed_word = __shfl_sync(__activemask(), packed_word, lane >> 3);
const int quantized = (int)((packed_word >> ((lane & 7) << 2)) & 15);
const unsigned char zp_byte = col_zp[block >> 1];
const int zero_point = (block & 1) ? (zp_byte >> 4) : (zp_byte & 15);
value += activation[depth] * ((float)quantized - (float)zero_point)
* col_scales[block];
}
value = block_sum(value);
if (threadIdx.x == 0) {
output[column] = value + (bias ? bias[column] : 0.0f);
}
}
// Per-K-block (block-32) int8 activation quantization. One warp (CUDA block) owns
// one K-block and emits that block's own int8 scale, matching ORT/MLAS CompInt8
// and the CPU native path. A single per-row scale is dominated by activation
// outliers and rounds small in-block magnitudes to zero, flipping argmaxes on
// outlier-heavy models (e.g. Phi-3.5); per-block scales track fp32 faithfully.
extern "C" __global__ void matmul_nbits_quantize_accuracy4_block32(
const float* activation,
signed char* quantized_activation,
float* activation_scale_out,
const int k,
const int padded_k)
{
(void)padded_k;
const int block = (int)blockIdx.x;
const int lane = (int)threadIdx.x;
const int depth = block * 32 + lane;
const float value = (depth < k) ? activation[depth] : 0.0f;
float max_abs = fabsf(value);
for (int offset = 16; offset > 0; offset >>= 1) {
max_abs = fmaxf(max_abs,
__shfl_down_sync(0xffffffffu, max_abs, offset));
}
max_abs = __shfl_sync(0xffffffffu, max_abs, 0);
const float activation_scale = max_abs == 0.0f ? 0.0f : max_abs / 127.0f;
const float inverse_scale =
activation_scale == 0.0f ? 0.0f : 1.0f / activation_scale;
if (lane == 0) {
activation_scale_out[block] = activation_scale;
}
int quantized = 0;
if (depth < k && activation_scale != 0.0f) {
quantized = (int)roundf(fminf(127.0f, fmaxf(-127.0f,
value * inverse_scale)));
}
quantized_activation[depth] = (signed char)quantized;
}
__device__ __forceinline__ int unpack_int4x4(unsigned int packed, int offset)
{
const int w0 = (int)((packed >> (offset + 0)) & 15u) - 8;
const int w1 = (int)((packed >> (offset + 4)) & 15u) - 8;
const int w2 = (int)((packed >> (offset + 8)) & 15u) - 8;
const int w3 = (int)((packed >> (offset + 12)) & 15u) - 8;
return (w0 & 255) | ((w1 & 255) << 8) | ((w2 & 255) << 16)
| ((w3 & 255) << 24);
}
__device__ __forceinline__ int dot_int8x4(int lhs, int rhs, int accumulator)
{
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 610
return __dp4a(lhs, rhs, accumulator);
#else
#pragma unroll
for (int byte = 0; byte < 4; ++byte) {
int lhs_value = ((unsigned int)lhs >> (byte * 8)) & 255u;
int rhs_value = ((unsigned int)rhs >> (byte * 8)) & 255u;
lhs_value = lhs_value >= 128 ? lhs_value - 256 : lhs_value;
rhs_value = rhs_value >= 128 ? rhs_value - 256 : rhs_value;
accumulator += lhs_value * rhs_value;
}
return accumulator;
#endif
}
extern "C" __global__ void matmul_nbits_gemv_accuracy4_block32(
const signed char* quantized_activation,
const float* activation_scale_ptr,
const unsigned char* packed,
const float* scales,
const float* bias,
float* output,
const int k,
const int n,
const int k_blocks)
{
extern __shared__ signed char activation_tile[];
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int column = (int)blockIdx.x * 8 + warp;
float value = 0.0f;
for (int tile_block = 0; tile_block < k_blocks; tile_block += 32) {
const int tile_blocks = min(32, k_blocks - tile_block);
const int tile_depths = tile_blocks * 32;
for (int depth = tid; depth < tile_depths; depth += (int)blockDim.x) {
activation_tile[depth] =
quantized_activation[tile_block * 32 + depth];
}
__syncthreads();
const int block = tile_block + lane;
if (column < n && block < k_blocks) {
const long packed_start = ((long)column * k_blocks + block) * 16;
const uint4 packed_weights =
*reinterpret_cast<const uint4*>(packed + packed_start);
const unsigned int words[4] = {
packed_weights.x, packed_weights.y, packed_weights.z, packed_weights.w
};
const signed char* activation_block = activation_tile + lane * 32;
int dot = 0;
#pragma unroll
for (int word = 0; word < 4; ++word) {
const int activation0 =
*reinterpret_cast<const int*>(activation_block + word * 8);
const int activation1 =
*reinterpret_cast<const int*>(activation_block + word * 8 + 4);
dot = dot_int8x4(activation0, unpack_int4x4(words[word], 0), dot);
dot = dot_int8x4(activation1, unpack_int4x4(words[word], 16), dot);
}
// Per-block int8 activation scale times the per-block weight scale.
const float block_scale = __fmul_rn(
activation_scale_ptr[block],
scales[(long)column * k_blocks + block]);
value = __fadd_rn(value, __fmul_rn((float)dot, block_scale));
}
__syncthreads();
}
value = warp_sum(value);
if (lane == 0 && column < n) {
output[column] = bias ? __fadd_rn(value, bias[column]) : value;
}
}
extern "C" __global__ void matmul_nbits_gemv_accuracy4_block32_stage64(
const signed char* quantized_activation,
const float* activation_scale_ptr,
const unsigned char* packed,
const float* scales,
const float* bias,
float* output,
const int k,
const int n,
const int k_blocks)
{
extern __shared__ signed char activation_tile[];
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int column = (int)blockIdx.x * 8 + warp;
float value = 0.0f;
for (int tile_block = 0; tile_block < k_blocks; tile_block += 64) {
const int tile_blocks = min(64, k_blocks - tile_block);
const int tile_depths = tile_blocks * 32;
for (int depth = tid; depth < tile_depths; depth += (int)blockDim.x) {
activation_tile[depth] =
quantized_activation[tile_block * 32 + depth];
}
__syncthreads();
for (int tile_offset = lane; tile_offset < tile_blocks; tile_offset += 32) {
const int block = tile_block + tile_offset;
if (column >= n) {
continue;
}
const long packed_start = ((long)column * k_blocks + block) * 16;
const uint4 packed_weights =
*reinterpret_cast<const uint4*>(packed + packed_start);
const unsigned int words[4] = {
packed_weights.x, packed_weights.y, packed_weights.z, packed_weights.w
};
const signed char* activation_block = activation_tile + tile_offset * 32;
int dot = 0;
#pragma unroll
for (int word = 0; word < 4; ++word) {
const int activation0 =
*reinterpret_cast<const int*>(activation_block + word * 8);
const int activation1 =
*reinterpret_cast<const int*>(activation_block + word * 8 + 4);
dot = dot_int8x4(activation0, unpack_int4x4(words[word], 0), dot);
dot = dot_int8x4(activation1, unpack_int4x4(words[word], 16), dot);
}
const float block_scale = __fmul_rn(
activation_scale_ptr[block],
scales[(long)column * k_blocks + block]);
value = __fadd_rn(value, __fmul_rn((float)dot, block_scale));
}
__syncthreads();
}
value = warp_sum(value);
if (lane == 0 && column < n) {
output[column] = bias ? __fadd_rn(value, bias[column]) : value;
}
}
// Per-K-block int8 activation quantization for ANY power-of-two block_size. One
// warp (CUDA block) owns one K-block and emits that block's own int8 scale,
// exactly matching the per-block quantization the tiled `matmul_nbits_accuracy4`
// reference performs inline (block_max / 127, symmetric round-to-nearest). The
// block-32 entry above hard-codes 32 lanes == 32 depths; this generalization
// strides each of the 32 lanes across the block so a single warp covers
// block_size (e.g. 128) depths. Padded tail depths (depth >= k in a partial
// final block) quantize to zero so they contribute nothing to the GEMV.
extern "C" __global__ void matmul_nbits_quantize_accuracy4_blockwise(
const float* activation,
signed char* quantized_activation,
float* activation_scale_out,
const int k,
const int block_size,
const int padded_k)
{
(void)padded_k;
const int block = (int)blockIdx.x;
const int lane = (int)threadIdx.x;
const int begin = block * block_size;
float max_abs = 0.0f;
for (int within = lane; within < block_size; within += 32) {
const int depth = begin + within;
const float value = (depth < k) ? activation[depth] : 0.0f;
max_abs = fmaxf(max_abs, fabsf(value));
}
for (int offset = 16; offset > 0; offset >>= 1) {
max_abs = fmaxf(max_abs,
__shfl_down_sync(0xffffffffu, max_abs, offset));
}
max_abs = __shfl_sync(0xffffffffu, max_abs, 0);
const float activation_scale = max_abs == 0.0f ? 0.0f : max_abs / 127.0f;
const float inverse_scale =
activation_scale == 0.0f ? 0.0f : 1.0f / activation_scale;
if (lane == 0) {
activation_scale_out[block] = activation_scale;
}
for (int within = lane; within < block_size; within += 32) {
const int depth = begin + within;
int quantized = 0;
if (depth < k && activation_scale != 0.0f) {
quantized = (int)roundf(fminf(127.0f, fmaxf(-127.0f,
activation[depth] * inverse_scale)));
}
quantized_activation[depth] = (signed char)quantized;
}
}
// General-block-size int4 accuracy_level=4 decode GEMV over the pre-quantized
// int8 activation. One warp reduces one output column; the 32 lanes cooperate
// across the block_size depths of each K-block. The per-block integer dot is
// computed as sum(qa * qw) - zero_point * sum(qa), which is exactly the tiled
// reference's sum(qa * (qw - zero_point)). Block-128 uses two packed dp4a
// instructions per lane (one for each integer sum) on sm_61+, while older
// architectures and other block sizes retain the scalar loop. Padded tail
// activations are zero, so packed tail lanes remain exact no-ops. The fp32 block
// products are then accumulated by lane 0 in ascending block order with the same
// __fmul_rn / __fadd_rn rounding the tiled kernel uses, so the result is bit-for-bit
// identical. Symmetric int4 uses the default zero point 8; asymmetric int4
// reads the packed per-block nibble zero point. The grid width (warps per CTA)
// is chosen host-side from the device multiprocessor count so the launch fills
// consumer and datacenter GPUs alike.
extern "C" __global__ void matmul_nbits_gemv_accuracy4_blockwise(
const signed char* quantized_activation,
const float* activation_scale_ptr,
const unsigned char* packed,
const float* scales,
const unsigned char* zero_points,
const float* bias,
float* output,
const int k,
const int n,
const int k_blocks,
const int block_size,
const int blob_size,
const int zp_row_bytes)
{
(void)k;
const int lane = (int)threadIdx.x & 31;
const int warp = (int)threadIdx.x >> 5;
const int column = (int)blockIdx.x * (int)(blockDim.x >> 5) + warp;
if (column >= n) {
return;
}
float value = 0.0f;
for (int block = 0; block < k_blocks; ++block) {
const int begin = block * block_size;
const long blob_base = ((long)column * k_blocks + block) * blob_size;
int weighted = 0;
int activation_sum = 0;
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 610
if (block_size == 128) {
const int within = lane * 4;
const int activation_pack = *reinterpret_cast<const int*>(
quantized_activation + begin + within);
const unsigned int packed_weights =
(unsigned int)*reinterpret_cast<const unsigned short*>(
packed + blob_base + (within >> 1));
const unsigned int even_weights = packed_weights & 0x0f0fu;
const unsigned int odd_weights = (packed_weights >> 4) & 0x0f0fu;
const unsigned int weight_pack =
__byte_perm(even_weights, odd_weights, 0x5140);
weighted = dot_int8x4(
activation_pack, (int)weight_pack, weighted);
activation_sum = dot_int8x4(
activation_pack, 0x01010101, activation_sum);
} else
#endif
{
for (int within = lane; within < block_size; within += 32) {
const int quantized_activation_value =
(int)quantized_activation[begin + within];
const unsigned char byte = packed[blob_base + (within >> 1)];
const int quantized_weight =
(within & 1) ? (byte >> 4) : (byte & 15);
weighted += quantized_activation_value * quantized_weight;
activation_sum += quantized_activation_value;
}
}
for (int offset = 16; offset > 0; offset >>= 1) {
weighted += __shfl_down_sync(0xffffffffu, weighted, offset);
activation_sum +=
__shfl_down_sync(0xffffffffu, activation_sum, offset);
}
if (lane == 0) {
int zero_point = 8;
if (zero_points) {
const unsigned char zp =
zero_points[(long)column * zp_row_bytes + (block >> 1)];
zero_point = (block & 1) ? (zp >> 4) : (zp & 15);
}
const int dot = weighted - zero_point * activation_sum;
const float combined_scale = __fmul_rn(
activation_scale_ptr[block],
scales[(long)column * k_blocks + block]);
value = __fadd_rn(value, __fmul_rn((float)dot, combined_scale));
}
}
if (lane == 0) {
output[column] = bias ? __fadd_rn(value, bias[column]) : value;
}
}
"#;
const ACCURACY4_SRC: &str = r#"
extern "C" __global__ void matmul_nbits_accuracy4(
const float* a,
const unsigned char* packed,
const float* scales,
const unsigned char* zero_points,
const float* bias,
float* y,
const int m,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes)
{
const long total = (long)m * n;
for (long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
idx < total; idx += (long)gridDim.x * blockDim.x) {
const int row = (int)(idx / n);
const int output = (int)(idx % n);
const float* activation = a + (long)row * k;
// Per-K-block int8 activation quantization (block scale == the weight
// block granularity), matching ORT/MLAS CompInt8 and the CPU native
// path. A single per-row scale is dominated by activation outliers and
// collapses small in-block magnitudes to zero, which flips argmaxes on
// outlier-heavy models (e.g. Phi-3.5); per-block scales avoid that.
float value = 0.0f;
for (int block = 0; block < k_blocks; ++block) {
const int begin = block * block_size;
const int end = min(begin + block_size, k);
float block_max = 0.0f;
for (int depth = begin; depth < end; ++depth) {
block_max = fmaxf(block_max, fabsf(activation[depth]));
}
if (block_max == 0.0f) {
continue;
}
const float activation_scale = block_max / 127.0f;
const float inverse_scale = 1.0f / activation_scale;
int zero_point = 8;
if (zero_points) {
const unsigned char zp =
zero_points[(long)output * zp_row_bytes + block / 2];
zero_point = (block & 1) ? (zp >> 4) : (zp & 15);
}
int dot = 0;
for (int depth = begin; depth < end; ++depth) {
int quantized_activation =
(int)roundf(fminf(127.0f, fmaxf(-127.0f,
activation[depth] * inverse_scale)));
const int within = depth - begin;
const unsigned char byte =
packed[((long)output * k_blocks + block) * blob_size + within / 2];
const int quantized_weight =
(within & 1) ? (byte >> 4) : (byte & 15);
dot += quantized_activation * (quantized_weight - zero_point);
}
const float combined_scale = __fmul_rn(
activation_scale,
scales[(long)output * k_blocks + block]);
value = __fadd_rn(value, __fmul_rn((float)dot, combined_scale));
}
y[idx] = bias ? __fadd_rn(value, bias[output]) : value;
}
}
"#;
const GEMV_F16_SRC: &str = r#"
#include <cuda_fp16.h>
#include <cuda_bf16.h>
// Narrowing store shared by the GEMV epilogues that can write their result
// straight into a bf16 consumer buffer instead of an fp16 staging buffer.
//
// `out_bf16` is a launch-uniform flag, so the branch costs one predicated
// select on the single lane that stores. The bf16 arm deliberately rounds
// fp32 -> fp16 -> bf16 rather than fp32 -> bf16 directly: the staging path it
// replaces rounds to fp16 in the GEMV and then casts that fp16 to bf16, and
// reproducing the double rounding is what keeps greedy decoding bit-identical.
__device__ __forceinline__ void matmul_nbits_store_narrowed(
void* __restrict__ output, const int index, const __half value,
const int out_bf16)
{
if (out_bf16) {
((__nv_bfloat16*)output)[index] = __float2bfloat16(__half2float(value));
} else {
((__half*)output)[index] = value;
}
}
__device__ __forceinline__ float warp_sum(float value)
{
for (int offset = 16; offset > 0; offset >>= 1) {
value += __shfl_down_sync(0xffffffffu, value, offset);
}
return value;
}
// Fp16 GEMV bias epilogue. The fp32 accumulator is always rounded to fp16 for
// the base output. When a bias is present:
// * `bias_post_round == 0` (native MatMulNBits bias): add in fp32 and round
// once — `fp16(acc + bias)` — matching an ORT-style fused epilogue.
// * `bias_post_round != 0` (a folded standalone `Add`): round the accumulator
// to fp16 first, then add the fp16 bias with a second fp16 round —
// `fp16(fp16(acc) + bias)` — reproducing the original two-op path so greedy
// tokens stay byte-identical.
__device__ __forceinline__ __half fold_bias_f16(
const float value,
const __half* __restrict__ bias,
const int column,
const int bias_post_round)
{
const __half rounded = __float2half(value);
if (!bias) {
return rounded;
}
const float b = __half2float(bias[column]);
if (bias_post_round) {
return __float2half(__half2float(rounded) + b);
}
return __float2half(value + b);
}
// One warp per output column. Block-32 INT8 stores one unsigned quantized
// weight byte per K element and one optional uint8 zero point per block.
extern "C" __global__ void matmul_nbits_gemv_int8_f16(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int scales_fp16,
const int bias_post_round)
{
// Mirrors the int4 `matmul_nbits_gemv_f16` work split: four adjacent lanes
// cooperate on one block-32 column, eight blocks are consumed per warp step.
// Each lane issues one aligned 8-byte packed-int8 load (uint2) and one 16-byte
// activation load (uint4), then a four-lane shuffle reduction reconstructs the
// block dot product before its scale is applied. This replaces the previous
// one-byte-per-lane scalar walk that only advanced 32 K per warp step.
const int lane = (int)threadIdx.x & 31;
const int warp = (int)threadIdx.x >> 5;
const int columns_per_block = (int)blockDim.x >> 5;
const int column = (int)blockIdx.x * columns_per_block + warp;
float value = 0.0f;
if (column < n) {
const int quarter = lane & 3;
for (int block_base = 0; block_base < k_blocks; block_base += 8) {
const int block = block_base + (lane >> 2);
float block_partial = 0.0f;
if (block < k_blocks) {
const int zero_point =
zero_points ? (int)zero_points[(long)column * k_blocks + block] : 128;
const int depth = block * 32 + quarter * 8;
const long packed_start =
((long)column * k_blocks + block) * 32 + quarter * 8;
if (depth + 8 <= k) {
const uint2 packed_word =
*reinterpret_cast<const uint2*>(packed + packed_start);
const unsigned char* bytes =
reinterpret_cast<const unsigned char*>(&packed_word);
const uint4 act = *reinterpret_cast<const uint4*>(activation + depth);
const __half* acth = reinterpret_cast<const __half*>(&act);
#pragma unroll
for (int i = 0; i < 8; ++i) {
block_partial += ((float)(int)bytes[i] - (float)zero_point)
* __half2float(acth[i]);
}
} else if (depth < k) {
const int valid = min(8, k - depth);
#pragma unroll
for (int i = 0; i < 8; ++i) {
if (i < valid) {
const int quantized = (int)packed[packed_start + i];
block_partial += ((float)quantized - (float)zero_point)
* __half2float(activation[depth + i]);
}
}
}
}
block_partial += __shfl_down_sync(0xffffffffu, block_partial, 2, 4);
block_partial += __shfl_down_sync(0xffffffffu, block_partial, 1, 4);
if (quarter == 0 && block < k_blocks) {
const float scale = scales_fp16
? __half2float(reinterpret_cast<const __half*>(scales_raw)
[(long)column * k_blocks + block])
: reinterpret_cast<const float*>(scales_raw)
[(long)column * k_blocks + block];
value += block_partial * scale;
}
}
}
value = warp_sum(value);
if (lane == 0 && column < n) {
output[column] = fold_bias_f16(value, bias, column, bias_post_round);
}
}
// Split-K standalone int8 GEMV: K_SPLIT warps cooperate on one output column,
// each reducing a strided subset of the 8-block (256-wide) K steps, then summing
// their fp32 partials through shared memory. The launch grid is K_SPLIT x larger
// than the single-warp kernel, which fills the SMs on the grid-starved
// (~0.48 waves/SM) Phi int8 down-projection decode GEMV. This kernel has no
// serial prologue, so the added grid parallelism directly hides the
// Long-Scoreboard latency (unlike the fused RMSNorm-prologue int8 kernel, whose
// serial full-vector prologue caps any split-K benefit). The fp32 partial sum is
// a new block-sum association, so results are near-equal (not byte-identical) to
// the single-warp kernel; asymmetric-zp parity is validated against a dequant
// reference to tolerance. Requires K % 256 == 0 (whole steps, no divergent tail)
// — the launch only routes here in that case.
extern "C" __global__ void matmul_nbits_gemv_int8_f16_splitk(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int scales_fp16,
const int bias_post_round)
{
constexpr int K_SPLIT = 2;
const int lane = (int)threadIdx.x & 31;
const int warp = (int)threadIdx.x >> 5;
const int warps_per_block = (int)blockDim.x >> 5;
const int cols_per_block = warps_per_block / K_SPLIT;
const int col_local = warp / K_SPLIT;
const int ks = warp % K_SPLIT;
const int column = (int)blockIdx.x * cols_per_block + col_local;
__shared__ float partials[8][K_SPLIT];
float value = 0.0f;
if (column < n) {
const int quarter = lane & 3;
for (int block_base = ks * 8; block_base < k_blocks;
block_base += K_SPLIT * 8) {
const int block = block_base + (lane >> 2);
float block_partial = 0.0f;
if (block < k_blocks) {
const int zero_point =
zero_points ? (int)zero_points[(long)column * k_blocks + block] : 128;
const int depth = block * 32 + quarter * 8;
const long packed_start =
((long)column * k_blocks + block) * 32 + quarter * 8;
if (depth + 8 <= k) {
const uint2 packed_word =
*reinterpret_cast<const uint2*>(packed + packed_start);
const unsigned char* bytes =
reinterpret_cast<const unsigned char*>(&packed_word);
const uint4 act = *reinterpret_cast<const uint4*>(activation + depth);
const __half* acth = reinterpret_cast<const __half*>(&act);
#pragma unroll
for (int i = 0; i < 8; ++i) {
block_partial += ((float)(int)bytes[i] - (float)zero_point)
* __half2float(acth[i]);
}
} else if (depth < k) {
const int valid = min(8, k - depth);
#pragma unroll
for (int i = 0; i < 8; ++i) {
if (i < valid) {
const int quantized = (int)packed[packed_start + i];
block_partial += ((float)quantized - (float)zero_point)
* __half2float(activation[depth + i]);
}
}
}
}
block_partial += __shfl_down_sync(0xffffffffu, block_partial, 2, 4);
block_partial += __shfl_down_sync(0xffffffffu, block_partial, 1, 4);
if (quarter == 0 && block < k_blocks) {
const float scale = scales_fp16
? __half2float(reinterpret_cast<const __half*>(scales_raw)
[(long)column * k_blocks + block])
: reinterpret_cast<const float*>(scales_raw)
[(long)column * k_blocks + block];
value += block_partial * scale;
}
}
}
value = warp_sum(value);
if (lane == 0) {
partials[col_local][ks] = (column < n) ? value : 0.0f;
}
__syncthreads();
if (ks == 0 && lane == 0 && column < n) {
float acc = 0.0f;
#pragma unroll
for (int s = 0; s < K_SPLIT; ++s) {
acc += partials[col_local][s];
}
output[column] = fold_bias_f16(acc, bias, column, bias_post_round);
}
}
// block-32 activation/weight tile, so each packed weight is reused by up to 16
// prompt rows and each activation by up to 16 output columns. It deliberately
// uses only ordinary shared memory, fp32 arithmetic, and __half conversion:
// no tensor-core, async-copy, or architecture-specific PTX requirement.
extern "C" __global__ void matmul_nbits_gemm_f16(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int m,
const int k,
const int n,
const int k_blocks,
const int bits,
const int scales_fp16,
const int bias_post_round,
const int bias_row_stride)
{
__shared__ float activation_tile[16][32];
__shared__ float weight_tile[32][16];
const int tid = (int)threadIdx.y * 16 + (int)threadIdx.x;
const int row = (int)blockIdx.y * 16 + (int)threadIdx.y;
const int column = (int)blockIdx.x * 16 + (int)threadIdx.x;
float value = 0.0f;
for (int block = 0; block < k_blocks; ++block) {
#pragma unroll
for (int load = tid; load < 16 * 32; load += 16 * 16) {
const int tile_row = load >> 5;
const int within = load & 31;
const int depth = block * 32 + within;
const int global_row = (int)blockIdx.y * 16 + tile_row;
activation_tile[tile_row][within] =
global_row < m && depth < k
? __half2float(activation[(long)global_row * k + depth])
: 0.0f;
}
#pragma unroll
for (int load = tid; load < 32 * 16; load += 16 * 16) {
const int tile_column = load >> 5;
const int within = load & 31;
const int global_column = (int)blockIdx.x * 16 + tile_column;
const int depth = block * 32 + within;
float weight = 0.0f;
if (global_column < n && depth < k) {
const long scale_index = (long)global_column * k_blocks + block;
const float scale = scales_fp16
? __half2float(
reinterpret_cast<const __half*>(scales_raw)[scale_index])
: reinterpret_cast<const float*>(scales_raw)[scale_index];
int quantized;
int zero_point;
if (bits == 8) {
quantized = (int)packed[scale_index * 32 + within];
zero_point = zero_points ? (int)zero_points[scale_index] : 128;
} else {
const unsigned char byte =
packed[scale_index * 16 + (within >> 1)];
quantized = (within & 1) ? (byte >> 4) : (byte & 15);
zero_point = 8;
if (zero_points) {
const int zp_row_bytes = (k_blocks + 1) >> 1;
const unsigned char zp =
zero_points[(long)global_column * zp_row_bytes + (block >> 1)];
zero_point = (block & 1) ? (zp >> 4) : (zp & 15);
}
}
weight = ((float)quantized - (float)zero_point) * scale;
}
weight_tile[within][tile_column] = weight;
}
__syncthreads();
if (row < m && column < n) {
#pragma unroll
for (int within = 0; within < 32; ++within) {
value += activation_tile[threadIdx.y][within]
* weight_tile[within][threadIdx.x];
}
}
__syncthreads();
}
if (row < m && column < n) {
// A folded residual epilogue binds a per-token residual (row stride N)
// into the bias slot; a genuine broadcast bias keeps stride 0.
const __half* row_bias = bias ? bias + (long)row * bias_row_stride : bias;
output[(long)row * n + column] =
fold_bias_f16(value, row_bias, column, bias_post_round);
}
}
__device__ __forceinline__ void int4x8_to_half2x4_sub(
const unsigned int packed,
__half2* values,
const unsigned int sub2)
{
unsigned int* h = reinterpret_cast<unsigned int*>(values);
constexpr unsigned int bottom_mask = 0x000f000f;
constexpr unsigned int top_mask = 0x00f000f0;
constexpr unsigned int fp16_magic = 0x64006400;
constexpr unsigned int lop3_lut = (0xf0 & 0xcc) | 0xaa;
const unsigned int top = packed >> 8;
asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
: "=r"(h[0])
: "r"(packed), "n"(bottom_mask), "n"(fp16_magic), "n"(lop3_lut));
asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
: "=r"(h[1])
: "r"(packed), "n"(top_mask), "n"(fp16_magic), "n"(lop3_lut));
asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
: "=r"(h[2])
: "r"(top), "n"(bottom_mask), "n"(fp16_magic), "n"(lop3_lut));
asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
: "=r"(h[3])
: "r"(top), "n"(top_mask), "n"(fp16_magic), "n"(lop3_lut));
constexpr unsigned int fp16_1024 = 0x64006400;
constexpr unsigned int fp16_one_sixteenth = 0x2c002c00;
constexpr unsigned int fp16_neg64 = 0xd400d400;
asm volatile("sub.f16x2 %0, %1, %2;\n"
: "=r"(h[0]) : "r"(h[0]), "r"(fp16_1024));
asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
: "=r"(h[1])
: "r"(h[1]), "r"(fp16_one_sixteenth), "r"(fp16_neg64));
asm volatile("sub.f16x2 %0, %1, %2;\n"
: "=r"(h[2]) : "r"(h[2]), "r"(fp16_1024));
asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
: "=r"(h[3])
: "r"(h[3]), "r"(fp16_one_sixteenth), "r"(fp16_neg64));
// Center each nibble by subtracting the block zero point. A symmetric int4
// weight uses the implicit `sub2 == 8` (fp16 0x48004800), which reproduces
// the previous fixed `- 8` byte-for-byte; an asymmetric weight passes its
// per-block zero point instead so the dequant is `(code - zp)`.
#pragma unroll
for (int i = 0; i < 4; ++i) {
asm volatile("sub.f16x2 %0, %1, %2;\n"
: "=r"(h[i]) : "r"(h[i]), "r"(sub2));
}
}
// Symmetric int4 dequant: `(code - 8)` in fp16, byte-identical to the historical
// hard-coded `- 8` path (the `sub2` register just carries fp16 8.0).
__device__ __forceinline__ void int4x8_to_half2x4(
const unsigned int packed,
__half2* values)
{
constexpr unsigned int fp16_eight = 0x48004800;
int4x8_to_half2x4_sub(packed, values, fp16_eight);
}
// Fused symmetric int4 dequant: `(code - 8)` in fp16, emitting FOUR fewer
// `f16x2` ALU ops per packed word than [`int4x8_to_half2x4`] by folding the
// symmetric zero point (`- 8`) into the magic-bias-removal constants instead of
// issuing it as a separate trailing `sub.f16x2` per element pair.
//
// BYTE-IDENTICAL to `int4x8_to_half2x4` (the `- 8` path): every intermediate is
// an exactly-representable fp16 integer, so folding the two constant
// subtractions into one changes no rounding:
// * bottom nibbles decode to `1024 + code` (exact in [1024, 1039]); the
// original does `(x - 1024) - 8`, this does `x - 1032` — `1032 = 0x6408` is
// exact and `(1024 + code) - 1032 = code - 8` with no intermediate rounding.
// * top nibbles decode to `1024 + 16*code`; the original fma `x*(1/16) - 64`
// yields the exact integer `code`, then `- 8`; this fuses to
// `x*(1/16) - 72` (`72 = 0xD480` exact). `x*(1/16)` is an exact power-of-two
// scale, so the single fma rounding lands on `code - 8` identically.
// This is a pure instruction-count reduction for the issue-bound symmetric
// paired gate/up decode GEMV — no reassociation, no math change.
__device__ __forceinline__ void int4x8_to_half2x4_sym8(
const unsigned int packed,
__half2* values)
{
unsigned int* h = reinterpret_cast<unsigned int*>(values);
constexpr unsigned int bottom_mask = 0x000f000f;
constexpr unsigned int top_mask = 0x00f000f0;
constexpr unsigned int fp16_magic = 0x64006400;
constexpr unsigned int lop3_lut = (0xf0 & 0xcc) | 0xaa;
const unsigned int top = packed >> 8;
asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
: "=r"(h[0])
: "r"(packed), "n"(bottom_mask), "n"(fp16_magic), "n"(lop3_lut));
asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
: "=r"(h[1])
: "r"(packed), "n"(top_mask), "n"(fp16_magic), "n"(lop3_lut));
asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
: "=r"(h[2])
: "r"(top), "n"(bottom_mask), "n"(fp16_magic), "n"(lop3_lut));
asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
: "=r"(h[3])
: "r"(top), "n"(top_mask), "n"(fp16_magic), "n"(lop3_lut));
// 1032 = 1024 (magic bias) + 8 (symmetric zero point); -72 = -64 - 8.
constexpr unsigned int fp16_1032 = 0x64086408;
constexpr unsigned int fp16_one_sixteenth = 0x2c002c00;
constexpr unsigned int fp16_neg72 = 0xd480d480;
asm volatile("sub.f16x2 %0, %1, %2;\n"
: "=r"(h[0]) : "r"(h[0]), "r"(fp16_1032));
asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
: "=r"(h[1])
: "r"(h[1]), "r"(fp16_one_sixteenth), "r"(fp16_neg72));
asm volatile("sub.f16x2 %0, %1, %2;\n"
: "=r"(h[2]) : "r"(h[2]), "r"(fp16_1032));
asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
: "=r"(h[3])
: "r"(h[3]), "r"(fp16_one_sixteenth), "r"(fp16_neg72));
}
// Pack a scalar block zero point (nibble in [0, 15]) into an fp16x2 subtrahend
// for [`int4x8_to_half2x4_sub`].
__device__ __forceinline__ unsigned int int4_zero_point_sub2(const int zero_point)
{
const __half zp = __float2half((float)zero_point);
const __half2 zp2 = __halves2half2(zp, zp);
return *reinterpret_cast<const unsigned int*>(&zp2);
}
// Load the block zero point for `column`/`block` from the packed nibble layout,
// or the symmetric default (8) when the weight carries no zero points.
__device__ __forceinline__ int int4_block_zero_point(
const unsigned char* __restrict__ zero_points,
const long column,
const int block,
const int zp_row_bytes)
{
if (!zero_points) {
return 8;
}
const unsigned char zp = zero_points[column * zp_row_bytes + (block >> 1)];
return (block & 1) ? (zp >> 4) : (zp & 15);
}
// Compile-time-specialized per-block subtrahend for the vectorized int4 GEMVs.
// `HasZp == false` (symmetric weights) folds to the constant fp16 `8.0`
// subtrahend with no memory traffic, so the compiler emits the exact
// pre-zero-point instruction stream; `HasZp == true` reads the per-block
// asymmetric zero point. Keying off the template parameter — never the runtime
// pointer — keeps the symmetric decode path byte-identical and register-light.
template <bool HasZp>
__device__ __forceinline__ unsigned int block_sub2(
const unsigned char* __restrict__ zero_points,
const long column,
const int block,
const int zp_row_bytes)
{
if (!HasZp) {
return 0x48004800u;
}
return int4_zero_point_sub2(
int4_block_zero_point(zero_points, column, block, zp_row_bytes));
}
// Scalar counterpart of [`block_sub2`] for the partial-block tail. `HasZp ==
// false` returns the symmetric default (8) with no load.
template <bool HasZp>
__device__ __forceinline__ int block_zp(
const unsigned char* __restrict__ zero_points,
const long column,
const int block,
const int zp_row_bytes)
{
if (!HasZp) {
return 8;
}
return int4_block_zero_point(zero_points, column, block, zp_row_bytes);
}
__device__ __forceinline__ float dot_int4x8_f16(
const unsigned int packed,
const __half* __restrict__ activation)
{
const uint4 a = *reinterpret_cast<const uint4*>(activation);
constexpr unsigned int low_halves = 0x5410;
constexpr unsigned int high_halves = 0x7632;
uint4 permuted;
asm volatile("prmt.b32 %0, %1, %2, %3;\n"
: "=r"(permuted.x) : "r"(a.x), "r"(a.z), "r"(low_halves));
asm volatile("prmt.b32 %0, %1, %2, %3;\n"
: "=r"(permuted.y) : "r"(a.x), "r"(a.z), "r"(high_halves));
asm volatile("prmt.b32 %0, %1, %2, %3;\n"
: "=r"(permuted.z) : "r"(a.y), "r"(a.w), "r"(low_halves));
asm volatile("prmt.b32 %0, %1, %2, %3;\n"
: "=r"(permuted.w) : "r"(a.y), "r"(a.w), "r"(high_halves));
__half2 q[4];
int4x8_to_half2x4(packed, q);
const float2 q04 = __half22float2(q[0]);
const float2 q15 = __half22float2(q[1]);
const float2 q26 = __half22float2(q[2]);
const float2 q37 = __half22float2(q[3]);
const float2 a04 = __half22float2(*reinterpret_cast<const __half2*>(&permuted.x));
const float2 a15 = __half22float2(*reinterpret_cast<const __half2*>(&permuted.y));
const float2 a26 = __half22float2(*reinterpret_cast<const __half2*>(&permuted.z));
const float2 a37 = __half22float2(*reinterpret_cast<const __half2*>(&permuted.w));
float dot = q04.x * a04.x;
dot += q15.x * a15.x;
dot += q26.x * a26.x;
dot += q37.x * a37.x;
dot += q04.y * a04.y;
dot += q15.y * a15.y;
dot += q26.y * a26.y;
dot += q37.y * a37.y;
return dot;
}
// Zero-point-aware [`dot_int4x8_f16`]. `sub2` is the fp16x2 subtrahend for this
// block (the packed zero point, or fp16 8.0 for symmetric weights). Because the
// centered code `(code - zp)` is an exact fp16 integer in [-15, 15], converting
// each `q` to float and accumulating the eight products in fp32 in ascending
// element order reproduces the scalar `(float)(code - zp) * __half2float(act)`
// path byte-for-byte — the LOP3 unpack only changes *how* the nibbles are
// decoded, not the arithmetic that follows. With `sub2 == 0x48004800` this is
// identical to [`dot_int4x8_f16`].
__device__ __forceinline__ float dot_int4x8_f16_sub(
const unsigned int packed,
const __half* __restrict__ activation,
const unsigned int sub2)
{
const uint4 a = *reinterpret_cast<const uint4*>(activation);
constexpr unsigned int low_halves = 0x5410;
constexpr unsigned int high_halves = 0x7632;
uint4 permuted;
asm volatile("prmt.b32 %0, %1, %2, %3;\n"
: "=r"(permuted.x) : "r"(a.x), "r"(a.z), "r"(low_halves));
asm volatile("prmt.b32 %0, %1, %2, %3;\n"
: "=r"(permuted.y) : "r"(a.x), "r"(a.z), "r"(high_halves));
asm volatile("prmt.b32 %0, %1, %2, %3;\n"
: "=r"(permuted.z) : "r"(a.y), "r"(a.w), "r"(low_halves));
asm volatile("prmt.b32 %0, %1, %2, %3;\n"
: "=r"(permuted.w) : "r"(a.y), "r"(a.w), "r"(high_halves));
__half2 q[4];
int4x8_to_half2x4_sub(packed, q, sub2);
const float2 q04 = __half22float2(q[0]);
const float2 q15 = __half22float2(q[1]);
const float2 q26 = __half22float2(q[2]);
const float2 q37 = __half22float2(q[3]);
const float2 a04 = __half22float2(*reinterpret_cast<const __half2*>(&permuted.x));
const float2 a15 = __half22float2(*reinterpret_cast<const __half2*>(&permuted.y));
const float2 a26 = __half22float2(*reinterpret_cast<const __half2*>(&permuted.z));
const float2 a37 = __half22float2(*reinterpret_cast<const __half2*>(&permuted.w));
float dot = q04.x * a04.x;
dot += q15.x * a15.x;
dot += q26.x * a26.x;
dot += q37.x * a37.x;
dot += q04.y * a04.y;
dot += q15.y * a15.y;
dot += q26.y * a26.y;
dot += q37.y * a37.y;
return dot;
}
// Split of `dot_int4x8_f16_sub` into an activation-decode half and a
// weight-dot half so the decoded activation can be REUSED across several output
// columns (column register-blocking). `decode_activation8` converts the eight
// contiguous fp16 activations at `activation` into fp32, laid out in the exact
// summation order `dot_int4x8_f16_sub` consumes them (a04.x, a15.x, a26.x,
// a37.x, a04.y, a15.y, a26.y, a37.y). `dot_int4x8_f16_sub_act` then reproduces
// the identical 8-term fp32 dot for one weight sub-word. Because the fp16->fp32
// conversions and the add order are unchanged, the result is BIT-IDENTICAL to
// `dot_int4x8_f16_sub`; the only difference is the activation is decoded once
// and shared by all columns instead of re-loaded per column.
__device__ __forceinline__ void decode_activation8(
const __half* __restrict__ activation,
float* __restrict__ a /* [8] */)
{
const uint4 av = *reinterpret_cast<const uint4*>(activation);
constexpr unsigned int low_halves = 0x5410;
constexpr unsigned int high_halves = 0x7632;
uint4 permuted;
asm volatile("prmt.b32 %0, %1, %2, %3;\n"
: "=r"(permuted.x) : "r"(av.x), "r"(av.z), "r"(low_halves));
asm volatile("prmt.b32 %0, %1, %2, %3;\n"
: "=r"(permuted.y) : "r"(av.x), "r"(av.z), "r"(high_halves));
asm volatile("prmt.b32 %0, %1, %2, %3;\n"
: "=r"(permuted.z) : "r"(av.y), "r"(av.w), "r"(low_halves));
asm volatile("prmt.b32 %0, %1, %2, %3;\n"
: "=r"(permuted.w) : "r"(av.y), "r"(av.w), "r"(high_halves));
const float2 a04 = __half22float2(*reinterpret_cast<const __half2*>(&permuted.x));
const float2 a15 = __half22float2(*reinterpret_cast<const __half2*>(&permuted.y));
const float2 a26 = __half22float2(*reinterpret_cast<const __half2*>(&permuted.z));
const float2 a37 = __half22float2(*reinterpret_cast<const __half2*>(&permuted.w));
a[0] = a04.x;
a[1] = a15.x;
a[2] = a26.x;
a[3] = a37.x;
a[4] = a04.y;
a[5] = a15.y;
a[6] = a26.y;
a[7] = a37.y;
}
__device__ __forceinline__ float dot_int4x8_f16_sub_act(
const unsigned int packed,
const float* __restrict__ a /* [8] */,
const unsigned int sub2)
{
__half2 q[4];
int4x8_to_half2x4_sub(packed, q, sub2);
const float2 q04 = __half22float2(q[0]);
const float2 q15 = __half22float2(q[1]);
const float2 q26 = __half22float2(q[2]);
const float2 q37 = __half22float2(q[3]);
float dot = q04.x * a[0];
dot += q15.x * a[1];
dot += q26.x * a[2];
dot += q37.x * a[3];
dot += q04.y * a[4];
dot += q15.y * a[5];
dot += q26.y * a[6];
dot += q37.y * a[7];
return dot;
}
// ---------------------------------------------------------------------------
// TRT-LLM-style interleaved + biased int4 -> fp16 decode (OPT-IN, symmetric).
//
// This is the runtime half of the `ONNX_GENAI_INTERLEAVE_DEQUANT` lever. It
// consumes weights that were **offline-interleaved** (nibbles of each 32-bit
// word rearranged from natural [e7 e6 e5 e4 | e3 e2 e1 e0] to the TRT-LLM
// [e7 e5 e3 e1 | e6 e4 e2 e0] order — even elements in the low four nibble
// slots, odd in the high four) by the host-side interleave pass. Given that
// layout, the SAME 4x LOP3 unpack that `int4x8_to_half2x4_sub` uses now emits
// the eight fp16 codes in NATURAL element order `{e0,e1},{e2,e3},{e4,e5},
// {e6,e7}`, so no `prmt.b32` activation reorder is needed downstream (the
// activation is consumed straight as contiguous __half2 pairs).
//
// It also folds the symmetric `-8` bias directly into the LOP3 magic-removal
// constants: the even lanes subtract 1032 (0x64086408 = 1024 + 8) and the odd
// lanes fma by 1/16 then subtract 72 (0xd480 = -(64 + 8)), so the converter
// yields `(code - 8)` in fp16 with NO trailing `sub.f16x2` loop. Total is
// 1 shift + 4 LOP3 + 2 sub.f16x2 + 2 fma.f16x2 = 9 instructions / 8 values.
//
// Correctness: for symmetric weights the previous path produces `(code - 8)` in
// fp16 via magic-removal followed by a `- 8` subtract; `(code - 8)` is an exact
// fp16 integer in [-8, 7], so this single-step form yields the byte-identical
// fp16 value. Because the eight products are then accumulated in fp32 in the
// same ascending element order (e0*a0, e1*a1, ..., e7*a7), the dot is
// BIT-IDENTICAL to `dot_int4x8_f16_sub` with `sub2 == fp16 8.0`.
__device__ __forceinline__ void int4x8_to_half2x4_interleaved_biased(
const unsigned int packed,
__half2* values)
{
unsigned int* h = reinterpret_cast<unsigned int*>(values);
constexpr unsigned int bottom_mask = 0x000f000f;
constexpr unsigned int top_mask = 0x00f000f0;
constexpr unsigned int fp16_magic = 0x64006400;
constexpr unsigned int lop3_lut = (0xf0 & 0xcc) | 0xaa;
const unsigned int top = packed >> 8;
asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
: "=r"(h[0])
: "r"(packed), "n"(bottom_mask), "n"(fp16_magic), "n"(lop3_lut));
asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
: "=r"(h[1])
: "r"(packed), "n"(top_mask), "n"(fp16_magic), "n"(lop3_lut));
asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
: "=r"(h[2])
: "r"(top), "n"(bottom_mask), "n"(fp16_magic), "n"(lop3_lut));
asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
: "=r"(h[3])
: "r"(top), "n"(top_mask), "n"(fp16_magic), "n"(lop3_lut));
// Fold the symmetric -8 bias into the magic-removal: even lanes hold
// (1024 + code) and subtract 1032; odd lanes hold (1024 + 16*code) and
// fma (*1/16 - 72), both yielding (code - 8) directly.
constexpr unsigned int fp16_1032 = 0x64086408;
constexpr unsigned int fp16_one_sixteenth = 0x2c002c00;
constexpr unsigned int fp16_neg72 = 0xd480d480;
asm volatile("sub.f16x2 %0, %1, %2;\n"
: "=r"(h[0]) : "r"(h[0]), "r"(fp16_1032));
asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
: "=r"(h[1])
: "r"(h[1]), "r"(fp16_one_sixteenth), "r"(fp16_neg72));
asm volatile("sub.f16x2 %0, %1, %2;\n"
: "=r"(h[2]) : "r"(h[2]), "r"(fp16_1032));
asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
: "=r"(h[3])
: "r"(h[3]), "r"(fp16_one_sixteenth), "r"(fp16_neg72));
}
// Natural-order activation decode for the interleaved dequant path: the eight
// contiguous fp16 activations are converted to fp32 in ascending order with no
// `prmt.b32` reorder (the interleaved converter already emits the weights in
// natural element order). `a[i] == float(activation[i])`, matching the layout
// `decode_activation8` produces for the non-interleaved path.
__device__ __forceinline__ void decode_activation8_natural(
const __half* __restrict__ activation,
float* __restrict__ a /* [8] */)
{
const uint4 av = *reinterpret_cast<const uint4*>(activation);
const float2 a01 = __half22float2(*reinterpret_cast<const __half2*>(&av.x));
const float2 a23 = __half22float2(*reinterpret_cast<const __half2*>(&av.y));
const float2 a45 = __half22float2(*reinterpret_cast<const __half2*>(&av.z));
const float2 a67 = __half22float2(*reinterpret_cast<const __half2*>(&av.w));
a[0] = a01.x;
a[1] = a01.y;
a[2] = a23.x;
a[3] = a23.y;
a[4] = a45.x;
a[5] = a45.y;
a[6] = a67.x;
a[7] = a67.y;
}
// Interleaved-weight sibling of `dot_int4x8_f16_sub_act`: the pre-decoded
// natural-order activation `a[8]` is dotted with the eight `(code - 8)` fp16
// weights from the biased interleaved converter, accumulated in fp32 in
// ascending element order. Bit-identical to `dot_int4x8_f16_sub_act` with
// `sub2 == fp16 8.0` on the non-interleaved layout of the same logical weights.
__device__ __forceinline__ float dot_int4x8_f16_interleaved_act(
const unsigned int packed,
const float* __restrict__ a /* [8] */)
{
__half2 q[4];
int4x8_to_half2x4_interleaved_biased(packed, q);
const float2 q01 = __half22float2(q[0]);
const float2 q23 = __half22float2(q[1]);
const float2 q45 = __half22float2(q[2]);
const float2 q67 = __half22float2(q[3]);
float dot = q01.x * a[0];
dot += q01.y * a[1];
dot += q23.x * a[2];
dot += q23.y * a[3];
dot += q45.x * a[4];
dot += q45.y * a[5];
dot += q67.x * a[6];
dot += q67.y * a[7];
return dot;
}
// Interleaved-weight sibling of `dot_int4x8_f16_sub` (scalar-tail helper): loads
// the eight contiguous fp16 activations itself, then dots as above.
__device__ __forceinline__ float dot_int4x8_f16_interleaved(
const unsigned int packed,
const __half* __restrict__ activation)
{
float a[8];
decode_activation8_natural(activation, a);
return dot_int4x8_f16_interleaved_act(packed, a);
}
// ---------------------------------------------------------------------------
// High-memory-level-parallelism (wide-load) int4 decode GEMV.
//
// The single-warp `..._general_bs` loop loads ONE 32-bit weight word (8 nibbles)
// per lane per step and immediately dequant+FMAs it — a dependent
// load->dequant->FMA chain that keeps ~1 weight load in flight per lane, so at
// M=1 (where the weight stream is the entire DRAM traffic) DRAM tops out ~19%
// while the SM pipe saturates on narrow-load issue. A head-to-head ncu of ORT's
// `MatMulFloatInt4Kernel` on the identical gate_up matrix showed it streams the
// same weights at 2.42 TB/s (50% DRAM) vs our 0.92 TB/s at the SAME grid/block/
// registers/occupancy — a pure memory-level-parallelism gap, not a math or
// tiling difference.
//
// This helper closes that gap: each lane owns 32 CONTIGUOUS nibbles per step and
// loads them with ONE 128-bit `uint4` weight load (4x fewer load instructions,
// 4x bytes/instruction), and the next step's `uint4` is issued BEFORE the
// current step's dequant/FMA so >=2 wide loads are always in flight to hide the
// ~10-cycle Long-Scoreboard global-load latency. It reuses the byte-tested LOP3
// dequant (`dot_int4x8_f16_sub`) on the four 8-nibble sub-words, so the
// per-element arithmetic is unchanged. NOT cp.async (pure issue overhead at M=1,
// proven regressing) and NOT a scalar `#pragma unroll` of the LOP3 body
// (register bloat -> occupancy cliff, proven regressing) — just wider
// synchronous vector loads at ~constant register footprint, the mechanism ORT
// uses.
//
// Numerics: the 32-wide lane interleave (vs the 8-wide single-warp interleave)
// regroups the fp32 partial sums, so the result is near-equal, NOT byte-
// identical, to `..._general_bs` — exactly the K-slice reassociation the split-K
// entries already ship by default. Each 32-nibble chunk lies inside ONE block
// (guarded to `block_size % 32 == 0`, i.e. glm block-128 and qwen block-32; the
// rare block-16 export falls back to the narrow kernel), so a single scale and
// zero point cover the chunk and the `uint4` load is 16-byte aligned.
//
// Returns this lane's fp32 partial (pre warp-reduction) for the weight column
// `column`, walking chunk starts `depth0, depth0+warp_stride, ...`.
// Read logical element `i` (0..7) from an INTERLEAVED 32-bit weight word. The
// offline interleave stores physical nibble slots as [e0,e2,e4,e6,e1,e3,e5,e7],
// so logical i maps to physical slot (i even ? i/2 : 4 + i/2). Used only by the
// interleaved GEMV's scalar tail; the main loop reads whole words via the LOP3
// converter.
__device__ __forceinline__ int interleaved_nibble(const unsigned int word, const int i)
{
const int slot = (i & 1) ? (4 + (i >> 1)) : (i >> 1);
return (int)((word >> (slot * 4)) & 15u);
}
__device__ __forceinline__ float gemv_int4_wide_lane_dot(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales,
const unsigned char* __restrict__ zero_points,
const int k,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const long column,
const int depth0,
const int warp_stride)
{
const long col_kb = column * (long)k_blocks;
float value = 0.0f;
int depth = depth0;
bool have = (depth + 32 <= k);
uint4 w;
if (have) {
const int block = depth / block_size;
const long blob_base = (col_kb + block) * (long)blob_size;
const int within = depth - block * block_size;
w = *reinterpret_cast<const uint4*>(packed + blob_base + (within >> 1));
}
while (have) {
const int ndepth = depth + warp_stride;
const bool have_next = (ndepth + 32 <= k);
uint4 wn;
if (have_next) {
// Issue the next wide weight load before consuming the current one,
// so two 128-bit loads are in flight across the dequant/FMA below.
const int nblock = ndepth / block_size;
const long nblob = (col_kb + nblock) * (long)blob_size;
const int nwithin = ndepth - nblock * block_size;
wn = *reinterpret_cast<const uint4*>(packed + nblob + (nwithin >> 1));
}
// `depth` is 32-aligned and `block_size % 32 == 0`, so `[depth, depth+32)`
// is inside a single block: one scale + one zero point cover the chunk.
const int block = depth / block_size;
float scale;
if (scales_fp16) {
scale = __half2float(reinterpret_cast<const __half*>(scales)[col_kb + block]);
} else {
scale = reinterpret_cast<const float*>(scales)[col_kb + block];
}
const unsigned int sub2 =
int4_zero_point_sub2(int4_block_zero_point(zero_points, column, block, zp_row_bytes));
value += scale * dot_int4x8_f16_sub(w.x, activation + depth, sub2);
value += scale * dot_int4x8_f16_sub(w.y, activation + depth + 8, sub2);
value += scale * dot_int4x8_f16_sub(w.z, activation + depth + 16, sub2);
value += scale * dot_int4x8_f16_sub(w.w, activation + depth + 24, sub2);
depth = ndepth;
w = wn;
have = have_next;
}
// Partial trailing chunk for this lane (`depth < k < depth + 32`): decode the
// valid 8-nibble sub-words, matching the narrow kernel's tail arithmetic.
if (depth < k) {
const int block = depth / block_size;
float scale;
if (scales_fp16) {
scale = __half2float(reinterpret_cast<const __half*>(scales)[col_kb + block]);
} else {
scale = reinterpret_cast<const float*>(scales)[col_kb + block];
}
const int zero_point = int4_block_zero_point(zero_points, column, block, zp_row_bytes);
const unsigned int sub2 = int4_zero_point_sub2(zero_point);
const long blob_base = (col_kb + block) * (long)blob_size;
for (int off = 0; off < 32 && depth + off < k; off += 8) {
const int d = depth + off;
const int within = d - block * block_size;
const int valid = min(8, k - d);
const unsigned int packed_word =
*reinterpret_cast<const unsigned int*>(packed + blob_base + (within >> 1));
if (valid == 8) {
value += scale * dot_int4x8_f16_sub(packed_word, activation + d, sub2);
} else {
float partial = 0.0f;
#pragma unroll
for (int i = 0; i < 8; ++i) {
if (i < valid) {
const int q = (int)((packed_word >> (i * 4)) & 15u) - zero_point;
partial += (float)q * __half2float(activation[d + i]);
}
}
value += partial * scale;
}
}
}
return value;
}
// Interleaved + biased (symmetric-only) sibling of `gemv_int4_wide_lane_dot`.
// Single-column wide-load lane dot that consumes offline-interleaved weights and
// folds the fixed symmetric `-8` bias inside the LOP3 converter (no per-block
// zero point, no `prmt.b32` activation reorder). The depth-2 software pipeline,
// the ascending sub-word order (w.x, w.y, w.z, w.w), and the fp32 accumulation
// order are byte-for-byte identical to `gemv_int4_wide_lane_dot` on symmetric
// weights, so each lane's fp32 partial is BIT-IDENTICAL. Used by the split-K
// wide interleaved kernel (the split-K partials therefore reduce to the same
// value as the non-interleaved split-K wide kernel).
__device__ __forceinline__ float gemv_int4_wide_lane_dot_interleaved(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales,
const int k,
const int block_size,
const int k_blocks,
const int blob_size,
const int scales_fp16,
const long column,
const int depth0,
const int warp_stride)
{
const long col_kb = column * (long)k_blocks;
float value = 0.0f;
int depth = depth0;
bool have = (depth + 32 <= k);
uint4 w;
if (have) {
const int block = depth / block_size;
const long blob_base = (col_kb + block) * (long)blob_size;
const int within = depth - block * block_size;
w = *reinterpret_cast<const uint4*>(packed + blob_base + (within >> 1));
}
while (have) {
const int ndepth = depth + warp_stride;
const bool have_next = (ndepth + 32 <= k);
uint4 wn;
if (have_next) {
const int nblock = ndepth / block_size;
const long nblob = (col_kb + nblock) * (long)blob_size;
const int nwithin = ndepth - nblock * block_size;
wn = *reinterpret_cast<const uint4*>(packed + nblob + (nwithin >> 1));
}
const int block = depth / block_size;
float scale;
if (scales_fp16) {
scale = __half2float(reinterpret_cast<const __half*>(scales)[col_kb + block]);
} else {
scale = reinterpret_cast<const float*>(scales)[col_kb + block];
}
value += scale * dot_int4x8_f16_interleaved(w.x, activation + depth);
value += scale * dot_int4x8_f16_interleaved(w.y, activation + depth + 8);
value += scale * dot_int4x8_f16_interleaved(w.z, activation + depth + 16);
value += scale * dot_int4x8_f16_interleaved(w.w, activation + depth + 24);
depth = ndepth;
w = wn;
have = have_next;
}
if (depth < k) {
const int block = depth / block_size;
float scale;
if (scales_fp16) {
scale = __half2float(reinterpret_cast<const __half*>(scales)[col_kb + block]);
} else {
scale = reinterpret_cast<const float*>(scales)[col_kb + block];
}
const long blob_base = (col_kb + block) * (long)blob_size;
for (int off = 0; off < 32 && depth + off < k; off += 8) {
const int d = depth + off;
const int within = d - block * block_size;
const int valid = min(8, k - d);
const unsigned int packed_word =
*reinterpret_cast<const unsigned int*>(packed + blob_base + (within >> 1));
if (valid == 8) {
value += scale * dot_int4x8_f16_interleaved(packed_word, activation + d);
} else {
float partial = 0.0f;
#pragma unroll
for (int i = 0; i < 8; ++i) {
if (i < valid) {
const int q = interleaved_nibble(packed_word, i) - 8;
partial += (float)q * __half2float(activation[d + i]);
}
}
value += partial * scale;
}
}
}
return value;
}
// Column register-blocked wide lane-dot: one warp accumulates WIDE_NC output
// columns at once. Each 8-element activation sub-word is decoded to fp32 ONCE
// (`decode_activation8`) and reused across all WIDE_NC columns, cutting the
// redundant activation L1 traffic (the head-to-head ncu limiter on the wide
// gate_up kernel was L1/TEX throughput, not DRAM) by ~WIDE_NC x, while the
// WIDE_NC independent 128-bit weight loads per chunk supply the memory-level
// parallelism that hides the Long-Scoreboard latency (replacing the depth-2
// software pipeline of the single-column `gemv_int4_wide_lane_dot`). The
// per-column `values[c] += scale * dot(sub-word)` sequence is byte-for-byte the
// same order as `gemv_int4_wide_lane_dot`, so each column's fp32 result is
// BIT-IDENTICAL to the single-column wide kernel (hence to the narrow kernel it
// already matches). `col_base` is this warp's first column; columns
// `col_base + c` beyond `n` are skipped.
#define WIDE_NC 4
__device__ __forceinline__ void gemv_int4_wide_lane_dot_multicol(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales,
const unsigned char* __restrict__ zero_points,
const int k,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const long col_base,
const int n,
const int depth0,
const int warp_stride,
float* __restrict__ values /* [WIDE_NC] */)
{
bool valid[WIDE_NC];
long col_kb[WIDE_NC];
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
values[c] = 0.0f;
const long column = col_base + c;
valid[c] = (column < n);
col_kb[c] = column * (long)k_blocks;
}
int depth = depth0;
while (depth + 32 <= k) {
const int block = depth / block_size;
const int within = depth - block * block_size;
// Issue all WIDE_NC independent 128-bit weight loads up front so they are
// in flight together (the load-level parallelism that hides load latency).
uint4 w[WIDE_NC];
float scale[WIDE_NC];
unsigned int sub2[WIDE_NC];
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
if (!valid[c]) {
continue;
}
const long base = (col_kb[c] + block) * (long)blob_size;
w[c] = *reinterpret_cast<const uint4*>(packed + base + (within >> 1));
if (scales_fp16) {
scale[c] = __half2float(reinterpret_cast<const __half*>(scales)[col_kb[c] + block]);
} else {
scale[c] = reinterpret_cast<const float*>(scales)[col_kb[c] + block];
}
sub2[c] = int4_zero_point_sub2(
int4_block_zero_point(zero_points, col_base + c, block, zp_row_bytes));
}
// Decode each 8-element activation sub-word once, reuse across columns.
// Sub-word order 0..3 preserves the ascending-K accumulation of the
// single-column kernel (w.x, w.y, w.z, w.w).
float a8[8];
#pragma unroll
for (int s = 0; s < 4; ++s) {
decode_activation8(activation + depth + s * 8, a8);
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
if (!valid[c]) {
continue;
}
const unsigned int word =
(s == 0) ? w[c].x : (s == 1) ? w[c].y : (s == 2) ? w[c].z : w[c].w;
values[c] += scale[c] * dot_int4x8_f16_sub_act(word, a8, sub2[c]);
}
}
depth += warp_stride;
}
// Partial trailing chunk (`depth < k < depth + 32`): replicate the
// single-column tail arithmetic per column.
if (depth < k) {
const int block = depth / block_size;
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
if (!valid[c]) {
continue;
}
float scale;
if (scales_fp16) {
scale = __half2float(reinterpret_cast<const __half*>(scales)[col_kb[c] + block]);
} else {
scale = reinterpret_cast<const float*>(scales)[col_kb[c] + block];
}
const int zero_point =
int4_block_zero_point(zero_points, col_base + c, block, zp_row_bytes);
const unsigned int sub2 = int4_zero_point_sub2(zero_point);
const long blob_base = (col_kb[c] + block) * (long)blob_size;
for (int off = 0; off < 32 && depth + off < k; off += 8) {
const int d = depth + off;
const int within = d - block * block_size;
const int valid_n = min(8, k - d);
const unsigned int packed_word =
*reinterpret_cast<const unsigned int*>(packed + blob_base + (within >> 1));
if (valid_n == 8) {
values[c] += scale * dot_int4x8_f16_sub(packed_word, activation + d, sub2);
} else {
float partial = 0.0f;
#pragma unroll
for (int i = 0; i < 8; ++i) {
if (i < valid_n) {
const int q = (int)((packed_word >> (i * 4)) & 15u) - zero_point;
partial += (float)q * __half2float(activation[d + i]);
}
}
values[c] += partial * scale;
}
}
}
}
}
// Interleaved + biased (symmetric-only) sibling of
// `gemv_int4_wide_lane_dot_multicol`. Consumes offline-interleaved weights and
// applies the fixed symmetric `-8` bias inside the LOP3 converter, so there is
// no per-block zero point and no `prmt.b32` activation reorder. Every column's
// fp32 accumulation order is unchanged (ascending element order, sub-word order
// w.x, w.y, w.z, w.w), so each column result is BIT-IDENTICAL to
// `gemv_int4_wide_lane_dot_multicol` on symmetric weights.
__device__ __forceinline__ void gemv_int4_wide_lane_dot_multicol_interleaved(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales,
const int k,
const int block_size,
const int k_blocks,
const int blob_size,
const int scales_fp16,
const long col_base,
const int n,
const int depth0,
const int warp_stride,
float* __restrict__ values /* [WIDE_NC] */)
{
bool valid[WIDE_NC];
long col_kb[WIDE_NC];
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
values[c] = 0.0f;
const long column = col_base + c;
valid[c] = (column < n);
col_kb[c] = column * (long)k_blocks;
}
int depth = depth0;
while (depth + 32 <= k) {
const int block = depth / block_size;
const int within = depth - block * block_size;
uint4 w[WIDE_NC];
float scale[WIDE_NC];
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
if (!valid[c]) {
continue;
}
const long base = (col_kb[c] + block) * (long)blob_size;
w[c] = *reinterpret_cast<const uint4*>(packed + base + (within >> 1));
if (scales_fp16) {
scale[c] = __half2float(reinterpret_cast<const __half*>(scales)[col_kb[c] + block]);
} else {
scale[c] = reinterpret_cast<const float*>(scales)[col_kb[c] + block];
}
}
float a8[8];
#pragma unroll
for (int s = 0; s < 4; ++s) {
decode_activation8_natural(activation + depth + s * 8, a8);
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
if (!valid[c]) {
continue;
}
const unsigned int word =
(s == 0) ? w[c].x : (s == 1) ? w[c].y : (s == 2) ? w[c].z : w[c].w;
values[c] += scale[c] * dot_int4x8_f16_interleaved_act(word, a8);
}
}
depth += warp_stride;
}
if (depth < k) {
const int block = depth / block_size;
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
if (!valid[c]) {
continue;
}
float scale;
if (scales_fp16) {
scale = __half2float(reinterpret_cast<const __half*>(scales)[col_kb[c] + block]);
} else {
scale = reinterpret_cast<const float*>(scales)[col_kb[c] + block];
}
const long blob_base = (col_kb[c] + block) * (long)blob_size;
for (int off = 0; off < 32 && depth + off < k; off += 8) {
const int d = depth + off;
const int within = d - block * block_size;
const int valid_n = min(8, k - d);
const unsigned int packed_word =
*reinterpret_cast<const unsigned int*>(packed + blob_base + (within >> 1));
if (valid_n == 8) {
values[c] += scale * dot_int4x8_f16_interleaved(packed_word, activation + d);
} else {
float partial = 0.0f;
#pragma unroll
for (int i = 0; i < 8; ++i) {
if (i < valid_n) {
const int q = interleaved_nibble(packed_word, i) - 8;
partial += (float)q * __half2float(activation[d + i]);
}
}
values[c] += partial * scale;
}
}
}
}
}
// Decode eight contiguous fp16 activations into four `__half2` lanes laid out in
// the SAME element pairing that `int4x8_to_half2x4_sub` produces for the weights
// (`ah[i] = (a_i, a_{i+4})`), so `__hfma2(q[i], ah[i], acc)` multiplies matching
// (weight, activation) pairs. This is the fp16 sibling of `decode_activation8`:
// it stops at the permute step and keeps the halves packed (NO fp16->fp32
// conversion), because the fp16-mixed kernel consumes them directly in half2
// fused multiply-adds.
__device__ __forceinline__ void decode_activation8_h2(
const __half* __restrict__ activation,
__half2* __restrict__ ah /* [4] */)
{
const uint4 av = *reinterpret_cast<const uint4*>(activation);
constexpr unsigned int low_halves = 0x5410;
constexpr unsigned int high_halves = 0x7632;
uint4 permuted;
asm volatile("prmt.b32 %0, %1, %2, %3;\n"
: "=r"(permuted.x) : "r"(av.x), "r"(av.z), "r"(low_halves));
asm volatile("prmt.b32 %0, %1, %2, %3;\n"
: "=r"(permuted.y) : "r"(av.x), "r"(av.z), "r"(high_halves));
asm volatile("prmt.b32 %0, %1, %2, %3;\n"
: "=r"(permuted.z) : "r"(av.y), "r"(av.w), "r"(low_halves));
asm volatile("prmt.b32 %0, %1, %2, %3;\n"
: "=r"(permuted.w) : "r"(av.y), "r"(av.w), "r"(high_halves));
ah[0] = *reinterpret_cast<const __half2*>(&permuted.x);
ah[1] = *reinterpret_cast<const __half2*>(&permuted.y);
ah[2] = *reinterpret_cast<const __half2*>(&permuted.z);
ah[3] = *reinterpret_cast<const __half2*>(&permuted.w);
}
// ---------------------------------------------------------------------------
// fp16 mixed-precision column register-blocked wide GEMV.
//
// Identical column-register-blocking + wide-load structure as
// `gemv_int4_wide_lane_dot_multicol`, but the inner multiply-accumulate runs in
// fp16 `__hfma2` (two fused MACs per instruction) instead of fp32 FFMA. This is
// the ONLY way to actually cut the dequant/MAC ALU that limits the multicol
// kernel (fp32 FFMA is already one fused op, so a "fp16-multiply then
// fp32-add-every-product" scheme is strictly MORE ops and cannot win). ORT's
// `MatMulFloatInt4Kernel` is fp16 for exactly this reason, so this is the
// fp16-vs-fp16 equal-conditions path.
//
// PRECISION CONTRACT (matches ORT's `MatMulFloat4BitsKernelM1` exactly): the
// per-lane K reduction runs entirely in fp16 __half2 accumulators — each 32-term
// chunk is summed in fp16, its per-block scale is folded in with __hfma2, and the
// result accumulates into a per-column fp16 running `total`. fp32 is used ONLY in
// the final cross-lane `warp_sum` (the 5-step shuffle). This is safe because the
// fp16 accumulation is a WIDE, SHALLOW tree: with 32 lanes striding K by 32, each
// lane folds only ~K/1024 chunks (≈4 for K=4096, ≈13 for K=13696), and inside a
// chunk the __half2 holds two 16-deep lanes — a total fp16 depth of tens, not
// thousands, so mantissa loss is negligible. (A NAIVE deep single-accumulator
// fp16 sum of all K *does* lose mantissa and flip tokens — that is the trap this
// wide-tree layout avoids, and why ORT accumulates in fp16 throughout.) Because
// the arithmetic mirrors ORT's, the f64-oracle error lands in ORT's own error
// class; it is NOT byte-identical to the fp32 path, so it ships gated on accuracy
// (error <= ORT vs the f64 oracle), not on bit-identity.
__device__ __forceinline__ void gemv_int4_fp16_lane_dot_multicol(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales,
const unsigned char* __restrict__ zero_points,
const int k,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const long col_base,
const int n,
const int depth0,
const int warp_stride,
float* __restrict__ values /* [WIDE_NC] */)
{
bool valid[WIDE_NC];
long col_kb[WIDE_NC];
// Per-column fp16 running total across the lane's chunks (ORT-style): the
// per-block scale is folded into this __half2 accumulator with __hfma2, so
// the entire per-lane K reduction stays in fp16 and only the final
// cross-lane `warp_sum` runs in fp32. Matching ORT's arithmetic exactly puts
// this kernel in the same error class as ORT's own int4 M=1 kernel. A single
// fp16 accumulator is a *wide, shallow* reduction tree — each lane folds only
// a handful of chunks (K / (32 lanes * 32) ≈ 4 for K=4096, ≈13 for K=13696),
// so almost no mantissa is lost (this is why full-fp16 accumulate is
// production-safe here, unlike a naive deep single-accumulator fp16 sum).
__half2 total[WIDE_NC];
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
values[c] = 0.0f;
total[c] = __float2half2_rn(0.0f);
const long column = col_base + c;
valid[c] = (column < n);
col_kb[c] = column * (long)k_blocks;
}
int depth = depth0;
while (depth + 32 <= k) {
const int block = depth / block_size;
const int within = depth - block * block_size;
uint4 w[WIDE_NC];
__half2 scale2[WIDE_NC];
unsigned int sub2[WIDE_NC];
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
if (!valid[c]) {
continue;
}
const long base = (col_kb[c] + block) * (long)blob_size;
w[c] = *reinterpret_cast<const uint4*>(packed + base + (within >> 1));
// Splat the per-block scale into both half lanes so it can be folded
// into the fp16 accumulator with a single __hfma2.
if (scales_fp16) {
scale2[c] = __half2half2(reinterpret_cast<const __half*>(scales)[col_kb[c] + block]);
} else {
scale2[c] = __float2half2_rn(reinterpret_cast<const float*>(scales)[col_kb[c] + block]);
}
sub2[c] = int4_zero_point_sub2(
int4_block_zero_point(zero_points, col_base + c, block, zp_row_bytes));
}
// fp16 accumulators for this chunk's 32 (unscaled) products, one per col.
__half2 acc[WIDE_NC];
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
acc[c] = __float2half2_rn(0.0f);
}
// Decode each 8-element activation sub-word to half2 once, reuse across
// columns (the multicol L1-traffic win), and fold it into every column's
// fp16 accumulator with __hfma2 (2 fused MACs/instruction).
#pragma unroll
for (int s = 0; s < 4; ++s) {
__half2 ah[4];
decode_activation8_h2(activation + depth + s * 8, ah);
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
if (!valid[c]) {
continue;
}
const unsigned int word =
(s == 0) ? w[c].x : (s == 1) ? w[c].y : (s == 2) ? w[c].z : w[c].w;
__half2 q[4];
int4x8_to_half2x4_sub(word, q, sub2[c]);
#pragma unroll
for (int i = 0; i < 4; ++i) {
acc[c] = __hfma2(q[i], ah[i], acc[c]);
}
}
}
// Scale this chunk's fp16 partial by the block scale and fold it into the
// fp16 running total (one __hfma2 per column) — fp32 is deferred to the
// final cross-lane reduction.
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
if (!valid[c]) {
continue;
}
total[c] = __hfma2(acc[c], scale2[c], total[c]);
}
depth += warp_stride;
}
// Collapse each column's fp16 running total to fp32 (the two half lanes are
// the two halves of the dot product). The tail below adds into this in fp32.
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
values[c] = __low2float(total[c]) + __high2float(total[c]);
}
// Partial trailing chunk: compute in fp32 exactly like the fp32 multicol tail
// so the K-tail stays precise (negligible perf, avoids fp16 edge cases).
if (depth < k) {
const int block = depth / block_size;
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
if (!valid[c]) {
continue;
}
float scale;
if (scales_fp16) {
scale = __half2float(reinterpret_cast<const __half*>(scales)[col_kb[c] + block]);
} else {
scale = reinterpret_cast<const float*>(scales)[col_kb[c] + block];
}
const int zero_point =
int4_block_zero_point(zero_points, col_base + c, block, zp_row_bytes);
const unsigned int sub2 = int4_zero_point_sub2(zero_point);
const long blob_base = (col_kb[c] + block) * (long)blob_size;
for (int off = 0; off < 32 && depth + off < k; off += 8) {
const int d = depth + off;
const int within = d - block * block_size;
const int valid_n = min(8, k - d);
const unsigned int packed_word =
*reinterpret_cast<const unsigned int*>(packed + blob_base + (within >> 1));
if (valid_n == 8) {
values[c] += scale * dot_int4x8_f16_sub(packed_word, activation + d, sub2);
} else {
float partial = 0.0f;
#pragma unroll
for (int i = 0; i < 8; ++i) {
if (i < valid_n) {
const int q = (int)((packed_word >> (i * 4)) & 15u) - zero_point;
partial += (float)q * __half2float(activation[d + i]);
}
}
values[c] += partial * scale;
}
}
}
}
}
__device__ __forceinline__ uint4 permute_activation_f16x8(
const __half* __restrict__ activation)
{
const uint4 a = *reinterpret_cast<const uint4*>(activation);
constexpr unsigned int low_halves = 0x5410;
constexpr unsigned int high_halves = 0x7632;
uint4 permuted;
asm volatile("prmt.b32 %0, %1, %2, %3;\n"
: "=r"(permuted.x) : "r"(a.x), "r"(a.z), "r"(low_halves));
asm volatile("prmt.b32 %0, %1, %2, %3;\n"
: "=r"(permuted.y) : "r"(a.x), "r"(a.z), "r"(high_halves));
asm volatile("prmt.b32 %0, %1, %2, %3;\n"
: "=r"(permuted.z) : "r"(a.y), "r"(a.w), "r"(low_halves));
asm volatile("prmt.b32 %0, %1, %2, %3;\n"
: "=r"(permuted.w) : "r"(a.y), "r"(a.w), "r"(high_halves));
return permuted;
}
__device__ __forceinline__ void accumulate_int4x8_f16_permuted(
const unsigned int packed,
const uint4& activation,
const __half scale,
__half2& sum0,
__half2& sum1,
__half2& sum2,
__half2& sum3)
{
__half2 q[4];
int4x8_to_half2x4(packed, q);
const __half2 scale2 = __halves2half2(scale, scale);
sum0 = __hfma2(
__hmul2(q[0], scale2),
*reinterpret_cast<const __half2*>(&activation.x),
sum0);
sum1 = __hfma2(
__hmul2(q[1], scale2),
*reinterpret_cast<const __half2*>(&activation.y),
sum1);
sum2 = __hfma2(
__hmul2(q[2], scale2),
*reinterpret_cast<const __half2*>(&activation.z),
sum2);
sum3 = __hfma2(
__hmul2(q[3], scale2),
*reinterpret_cast<const __half2*>(&activation.w),
sum3);
}
// Zero-point-aware [`accumulate_int4x8_f16_permuted`]: `sub2` centers each
// nibble by the block zero point (fp16 8.0 for symmetric weights, giving a
// byte-identical result). Used by the paired gate/up kernels, which permute the
// shared activation once and dequant each projection with its own zero point.
__device__ __forceinline__ void accumulate_int4x8_f16_permuted_zp(
const unsigned int packed,
const uint4& activation,
const __half scale,
const unsigned int sub2,
__half2& sum0,
__half2& sum1,
__half2& sum2,
__half2& sum3)
{
__half2 q[4];
int4x8_to_half2x4_sub(packed, q, sub2);
const __half2 scale2 = __halves2half2(scale, scale);
sum0 = __hfma2(__hmul2(q[0], scale2),
*reinterpret_cast<const __half2*>(&activation.x), sum0);
sum1 = __hfma2(__hmul2(q[1], scale2),
*reinterpret_cast<const __half2*>(&activation.y), sum1);
sum2 = __hfma2(__hmul2(q[2], scale2),
*reinterpret_cast<const __half2*>(&activation.z), sum2);
sum3 = __hfma2(__hmul2(q[3], scale2),
*reinterpret_cast<const __half2*>(&activation.w), sum3);
}
// Fused-symmetric [`accumulate_int4x8_f16_permuted_zp`]: identical multiply-add
// (same permuted activation, same `q * scale` fp16 FMA order into `sum0..3`) but
// dequants with [`int4x8_to_half2x4_sym8`], which folds the `- 8` symmetric zero
// point into the bias constants. The `q[i]` values are byte-identical to the
// `sub2 == 8` path, so the accumulation is bit-for-bit the same while issuing
// four fewer `f16x2` ops per word — the instruction-count win for the
// issue-bound paired gate/up decode GEMV.
__device__ __forceinline__ void accumulate_int4x8_f16_permuted_sym8(
const unsigned int packed,
const uint4& activation,
const __half scale,
__half2& sum0,
__half2& sum1,
__half2& sum2,
__half2& sum3)
{
__half2 q[4];
int4x8_to_half2x4_sym8(packed, q);
const __half2 scale2 = __halves2half2(scale, scale);
sum0 = __hfma2(__hmul2(q[0], scale2),
*reinterpret_cast<const __half2*>(&activation.x), sum0);
sum1 = __hfma2(__hmul2(q[1], scale2),
*reinterpret_cast<const __half2*>(&activation.y), sum1);
sum2 = __hfma2(__hmul2(q[2], scale2),
*reinterpret_cast<const __half2*>(&activation.z), sum2);
sum3 = __hfma2(__hmul2(q[3], scale2),
*reinterpret_cast<const __half2*>(&activation.w), sum3);
}
__device__ __forceinline__ void accumulate_int4x8_dot_f16(
const unsigned int packed,
const uint4& activation,
const __half2 scale2,
__half2& sum)
{
__half2 q[4];
int4x8_to_half2x4(packed, q);
sum = __hfma2(
__hmul2(q[0], scale2),
*reinterpret_cast<const __half2*>(&activation.x),
sum);
sum = __hfma2(
__hmul2(q[1], scale2),
*reinterpret_cast<const __half2*>(&activation.y),
sum);
sum = __hfma2(
__hmul2(q[2], scale2),
*reinterpret_cast<const __half2*>(&activation.z),
sum);
sum = __hfma2(
__hmul2(q[3], scale2),
*reinterpret_cast<const __half2*>(&activation.w),
sum);
}
__device__ __forceinline__ float dot_int4x32_f16_permuted_scaled(
const uint4& packed,
const uint4& activation0,
const uint4& activation1,
const uint4& activation2,
const uint4& activation3,
const __half scale)
{
const __half2 scale2 = __halves2half2(scale, scale);
__half2 sum0 = __float2half2_rn(0.0f);
__half2 sum1 = __float2half2_rn(0.0f);
__half2 sum2 = __float2half2_rn(0.0f);
__half2 sum3 = __float2half2_rn(0.0f);
accumulate_int4x8_dot_f16(packed.x, activation0, scale2, sum0);
accumulate_int4x8_dot_f16(packed.y, activation1, scale2, sum1);
accumulate_int4x8_dot_f16(packed.z, activation2, scale2, sum2);
accumulate_int4x8_dot_f16(packed.w, activation3, scale2, sum3);
const float2 value0 = __half22float2(sum0);
const float2 value1 = __half22float2(sum1);
const float2 value2 = __half22float2(sum2);
const float2 value3 = __half22float2(sum3);
float value = value0.x;
value += value1.x;
value += value2.x;
value += value3.x;
value += value0.y;
value += value1.y;
value += value2.y;
value += value3.y;
return value;
}
__device__ __forceinline__ void accumulate_int4x8_f16(
const unsigned int packed,
const __half* __restrict__ activation,
const __half scale,
__half2& sum0,
__half2& sum1,
__half2& sum2,
__half2& sum3)
{
const uint4 permuted = permute_activation_f16x8(activation);
accumulate_int4x8_f16_permuted(
packed, permuted, scale, sum0, sum1, sum2, sum3);
}
// Zero-point-aware variant of [`accumulate_int4x8_f16`]: `sub2` is the fp16x2
// subtrahend for this block (the packed zero point, or fp16 8.0 for symmetric
// weights). With the symmetric default this is byte-identical to the plain
// accumulate, so callers can route both symmetric and asymmetric weights here.
__device__ __forceinline__ void accumulate_int4x8_f16_zp(
const unsigned int packed,
const __half* __restrict__ activation,
const __half scale,
const unsigned int sub2,
__half2& sum0,
__half2& sum1,
__half2& sum2,
__half2& sum3)
{
const uint4 permuted = permute_activation_f16x8(activation);
__half2 q[4];
int4x8_to_half2x4_sub(packed, q, sub2);
const __half2 scale2 = __halves2half2(scale, scale);
sum0 = __hfma2(__hmul2(q[0], scale2),
*reinterpret_cast<const __half2*>(&permuted.x), sum0);
sum1 = __hfma2(__hmul2(q[1], scale2),
*reinterpret_cast<const __half2*>(&permuted.y), sum1);
sum2 = __hfma2(__hmul2(q[2], scale2),
*reinterpret_cast<const __half2*>(&permuted.z), sum2);
sum3 = __hfma2(__hmul2(q[3], scale2),
*reinterpret_cast<const __half2*>(&permuted.w), sum3);
}
// Debias-then-scale int4 dequant (fold-scale). Recovers the exact integer code
// in fp16 (the same 1024/64 lop3-bias removal as `int4x8_to_half2x4_sub`, which
// is fp16-exact for codes 0..15 so there is no catastrophic cancellation), then
// folds the per-block scale AND zero point into ONE `fma` per pair:
// `(code - zp) * scale = fma(code, scale2, neg_zp_scale2)`. This lets the MAC
// drop its separate `__hmul2(q, scale)`, removing 4 fp16x2 multiplies per 8
// weights on the ALU-co-bound (measured 65% pipe) M=1 decode GEMV.
__device__ __forceinline__ void int4x8_to_half2x4_scaledsub(
const unsigned int packed,
__half2* values,
const unsigned int scale2,
const unsigned int neg_zp_scale2)
{
unsigned int* h = reinterpret_cast<unsigned int*>(values);
constexpr unsigned int bottom_mask = 0x000f000f;
constexpr unsigned int top_mask = 0x00f000f0;
constexpr unsigned int fp16_magic = 0x64006400;
constexpr unsigned int lop3_lut = (0xf0 & 0xcc) | 0xaa;
const unsigned int top = packed >> 8;
asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
: "=r"(h[0])
: "r"(packed), "n"(bottom_mask), "n"(fp16_magic), "n"(lop3_lut));
asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
: "=r"(h[1])
: "r"(packed), "n"(top_mask), "n"(fp16_magic), "n"(lop3_lut));
asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
: "=r"(h[2])
: "r"(top), "n"(bottom_mask), "n"(fp16_magic), "n"(lop3_lut));
asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
: "=r"(h[3])
: "r"(top), "n"(top_mask), "n"(fp16_magic), "n"(lop3_lut));
constexpr unsigned int fp16_1024 = 0x64006400;
constexpr unsigned int fp16_one_sixteenth = 0x2c002c00;
constexpr unsigned int fp16_neg64 = 0xd400d400;
// Debias to the exact integer code first (fp16-exact for 0..15).
asm volatile("sub.f16x2 %0, %1, %2;\n"
: "=r"(h[0]) : "r"(h[0]), "r"(fp16_1024));
asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
: "=r"(h[1])
: "r"(h[1]), "r"(fp16_one_sixteenth), "r"(fp16_neg64));
asm volatile("sub.f16x2 %0, %1, %2;\n"
: "=r"(h[2]) : "r"(h[2]), "r"(fp16_1024));
asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
: "=r"(h[3])
: "r"(h[3]), "r"(fp16_one_sixteenth), "r"(fp16_neg64));
// Fold `(code - zp) * scale` into one fma per pair (replaces the standalone
// zp-subtract here and the `__hmul2(q, scale)` in the MAC).
#pragma unroll
for (int i = 0; i < 4; ++i) {
asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
: "=r"(h[i]) : "r"(h[i]), "r"(scale2), "r"(neg_zp_scale2));
}
}
// per CTA. Four adjacent lanes split each block-32 weight blob into aligned
// uint32 loads, so every warp issues contiguous 128-byte packed-weight
// transactions. Each lane also reads eight activations with one uint4 load.
// Register-only nibble conversion and four-lane shuffle reduction reconstruct
// each block dot product before applying its scale.
extern "C" __global__ void matmul_nbits_gemv_f16(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round)
{
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int columns_per_block = (int)blockDim.x >> 5;
const int column = (int)blockIdx.x * columns_per_block + warp;
float value = 0.0f;
if (column < n) {
const int quarter = lane & 3;
for (int block_base = 0; block_base < k_blocks; block_base += 8) {
const int block = block_base + (lane >> 2);
float block_partial = 0.0f;
if (block < k_blocks) {
const int depth = block * block_size + quarter * 8;
const long packed_start =
((long)column * k_blocks + block) * blob_size + quarter * 4;
const unsigned int packed_word =
*reinterpret_cast<const unsigned int*>(packed + packed_start);
int zero_point = 8;
if (zero_points) {
const unsigned char zp =
zero_points[(long)column * zp_row_bytes + block / 2];
zero_point = (block & 1) ? (zp >> 4) : (zp & 15);
}
if (depth + 8 <= k) {
if (zero_points) {
#pragma unroll
for (int i = 0; i < 8; ++i) {
const int q =
(int)((packed_word >> (i * 4)) & 15u) - zero_point;
block_partial +=
(float)q * __half2float(activation[depth + i]);
}
} else {
block_partial = dot_int4x8_f16(packed_word, activation + depth);
}
} else if (depth < k) {
const int valid = min(8, k - depth);
#pragma unroll
for (int i = 0; i < 8; ++i) {
if (i < valid) {
const int q =
(int)((packed_word >> (i * 4)) & 15u) - zero_point;
block_partial +=
(float)q * __half2float(activation[depth + i]);
}
}
}
}
block_partial += __shfl_down_sync(0xffffffffu, block_partial, 2, 4);
block_partial += __shfl_down_sync(0xffffffffu, block_partial, 1, 4);
if (quarter == 0 && block < k_blocks) {
float scale;
if (scales_fp16) {
scale = __half2float(
reinterpret_cast<const __half*>(scales)[(long)column * k_blocks + block]);
} else {
scale =
reinterpret_cast<const float*>(scales)[(long)column * k_blocks + block];
}
value += block_partial * scale;
}
}
}
value = warp_sum(value);
if (lane == 0 && column < n) {
output[column] = fold_bias_f16(value, bias, column, bias_post_round);
}
}
template <bool HasZp>
__device__ __forceinline__ void matmul_nbits_gemv_f16_scales_f16_tpl(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round)
{
(void)block_size;
(void)scales_fp16;
const __half* __restrict__ scales =
reinterpret_cast<const __half*>(scales_raw);
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int columns_per_block = (int)blockDim.x >> 5;
const int column_base = (int)blockIdx.x * columns_per_block;
const int column = column_base + warp;
__half2 sum0 = __float2half2_rn(0.0f);
__half2 sum1 = __float2half2_rn(0.0f);
__half2 sum2 = __float2half2_rn(0.0f);
__half2 sum3 = __float2half2_rn(0.0f);
float tail = 0.0f;
if (column < n) {
const int lane_depth = lane * 8;
const __half* activation_ptr = activation + lane_depth;
const unsigned char* packed_ptr =
packed + (long)column * k_blocks * blob_size + lane * 4;
const __half* scale_ptr =
scales + (long)column * k_blocks + (lane >> 2);
int depth_base = 0;
for (; depth_base + lane_depth + 8 <= k; depth_base += 256) {
const unsigned int packed_word =
*reinterpret_cast<const unsigned int*>(packed_ptr);
// block == depth/32; each lane's 8 nibbles all sit in one block.
const int block = (depth_base >> 5) + (lane >> 2);
const unsigned int sub2 =
block_sub2<HasZp>(zero_points, column, block, zp_row_bytes);
accumulate_int4x8_f16_zp(
packed_word,
activation_ptr,
*scale_ptr,
sub2,
sum0,
sum1,
sum2,
sum3);
activation_ptr += 256;
packed_ptr += 128;
scale_ptr += 8;
}
const int tail_depth = depth_base + lane_depth;
if (tail_depth < k) {
const unsigned int packed_word =
*reinterpret_cast<const unsigned int*>(packed_ptr);
const float scale = __half2float(*scale_ptr);
const int tail_block = (depth_base >> 5) + (lane >> 2);
const int zero_point =
block_zp<HasZp>(zero_points, column, tail_block, zp_row_bytes);
const int valid = min(8, k - tail_depth);
#pragma unroll
for (int i = 0; i < 8; ++i) {
if (i < valid) {
const int q = (int)((packed_word >> (i * 4)) & 15u) - zero_point;
tail += (float)q * __half2float(activation_ptr[i]) * scale;
}
}
}
}
const float2 value04 = __half22float2(sum0);
const float2 value15 = __half22float2(sum1);
const float2 value26 = __half22float2(sum2);
const float2 value37 = __half22float2(sum3);
float value = tail + value04.x;
value += value15.x;
value += value26.x;
value += value37.x;
value += value04.y;
value += value15.y;
value += value26.y;
value += value37.y;
value = warp_sum(value);
if (lane == 0 && column < n) {
output[column] = fold_bias_f16(value, bias, column, bias_post_round);
}
}
extern "C" __global__ void matmul_nbits_gemv_f16_scales_f16(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round)
{
matmul_nbits_gemv_f16_scales_f16_tpl<false>(activation, packed, scales_raw, zero_points, bias, output, k, n, block_size, k_blocks, blob_size, zp_row_bytes, scales_fp16, bias_post_round);
}
extern "C" __global__ void matmul_nbits_gemv_f16_scales_f16_zp(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round)
{
matmul_nbits_gemv_f16_scales_f16_tpl<true>(activation, packed, scales_raw, zero_points, bias, output, k, n, block_size, k_blocks, blob_size, zp_row_bytes, scales_fp16, bias_post_round);
}
// Prefetch-pipelined sibling of `matmul_nbits_gemv_f16_scales_f16_tpl`.
// The single-warp scales-fp16 GEMV is Long-Scoreboard bound: ncu on qwen2.5-14b
// q/o (N=5120) shows ~8.9 active warps/scheduler but only ~0.97 eligible, with
// ~75% of warp cycles stalled waiting for the one in-flight 32-bit weight load.
// The math is not the bottleneck (SM ~41%, DRAM ~24% of an H200's 4.8 TB/s).
//
// This variant keeps the EXACT same lane->nibble mapping (lane owns 8 contiguous
// nibbles at stride 256), the same fp16 `accumulate_int4x8_f16_zp` calls, and the
// same accumulation order — so its output is BYTE-IDENTICAL to the single-warp
// kernel. The only change is memory-level parallelism: a depth-PF register shift
// register holds PF prefetched weight words so PF independent global loads are in
// flight per lane, hiding the ~13-cycle load latency instead of stalling on it.
// Register-resident (manual rotation, no dynamic-indexed array -> no local spill).
template <bool HasZp>
__device__ __forceinline__ void matmul_nbits_gemv_f16_scales_f16_pipe_tpl(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round)
{
(void)block_size;
(void)scales_fp16;
const __half* __restrict__ scales =
reinterpret_cast<const __half*>(scales_raw);
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int columns_per_block = (int)blockDim.x >> 5;
const int column_base = (int)blockIdx.x * columns_per_block;
const int column = column_base + warp;
__half2 sum0 = __float2half2_rn(0.0f);
__half2 sum1 = __float2half2_rn(0.0f);
__half2 sum2 = __float2half2_rn(0.0f);
__half2 sum3 = __float2half2_rn(0.0f);
float tail = 0.0f;
if (column < n) {
const int lane_depth = lane * 8;
const __half* activation_ptr = activation + lane_depth;
const unsigned char* packed_ptr =
packed + (long)column * k_blocks * blob_size + lane * 4;
const __half* scale_ptr =
scales + (long)column * k_blocks + (lane >> 2);
// Number of full 8-nibble steps this lane walks (step s is valid iff
// s * 256 + lane_depth + 8 <= k), identical to the scalar loop bound.
const int nfull = (lane_depth + 8 <= k) ? ((k - lane_depth - 8) / 256 + 1) : 0;
// Load the k-th step's weight word (step stride is 128 packed bytes).
auto load_step = [&](int s) -> unsigned int {
return *reinterpret_cast<const unsigned int*>(packed_ptr + (long)s * 128);
};
// Prime the shift register with the first PF steps' words (independent
// loads -> they pipeline). Out-of-range slots are zero (never consumed).
// The array is `#pragma unroll`-rotated with a compile-time `PF`, so it
// stays fully in registers (no dynamic indexing -> no local spill).
constexpr int PF = 2; // prefetch depth (weight words in flight per lane)
unsigned int wbuf[PF];
#pragma unroll
for (int s = 0; s < PF; ++s) {
wbuf[s] = (s < nfull) ? load_step(s) : 0u;
}
int depth_base = 0;
for (int i = 0; i < nfull; ++i, depth_base += 256) {
const unsigned int packed_word = wbuf[0];
// Rotate the shift register and issue the load PF steps ahead BEFORE
// consuming the current word, so PF loads stay in flight.
const int pf = i + PF;
#pragma unroll
for (int s = 0; s < PF - 1; ++s) {
wbuf[s] = wbuf[s + 1];
}
wbuf[PF - 1] = (pf < nfull) ? load_step(pf) : 0u;
const int block = (depth_base >> 5) + (lane >> 2);
const unsigned int sub2 =
block_sub2<HasZp>(zero_points, column, block, zp_row_bytes);
accumulate_int4x8_f16_zp(
packed_word, activation_ptr, *scale_ptr, sub2, sum0, sum1, sum2, sum3);
activation_ptr += 256;
scale_ptr += 8;
}
const int tail_depth = depth_base + lane_depth;
if (tail_depth < k) {
const unsigned int packed_word =
*reinterpret_cast<const unsigned int*>(packed_ptr + (long)nfull * 128);
const float scale = __half2float(*scale_ptr);
const int tail_block = (depth_base >> 5) + (lane >> 2);
const int zero_point =
block_zp<HasZp>(zero_points, column, tail_block, zp_row_bytes);
const int valid = min(8, k - tail_depth);
#pragma unroll
for (int i = 0; i < 8; ++i) {
if (i < valid) {
const int q = (int)((packed_word >> (i * 4)) & 15u) - zero_point;
tail += (float)q * __half2float(activation_ptr[i]) * scale;
}
}
}
}
const float2 value04 = __half22float2(sum0);
const float2 value15 = __half22float2(sum1);
const float2 value26 = __half22float2(sum2);
const float2 value37 = __half22float2(sum3);
float value = tail + value04.x;
value += value15.x;
value += value26.x;
value += value37.x;
value += value04.y;
value += value15.y;
value += value26.y;
value += value37.y;
value = warp_sum(value);
if (lane == 0 && column < n) {
output[column] = fold_bias_f16(value, bias, column, bias_post_round);
}
}
extern "C" __global__ void matmul_nbits_gemv_f16_scales_f16_pipe(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round)
{
matmul_nbits_gemv_f16_scales_f16_pipe_tpl<false>(activation, packed, scales_raw, zero_points, bias, output, k, n, block_size, k_blocks, blob_size, zp_row_bytes, scales_fp16, bias_post_round);
}
extern "C" __global__ void matmul_nbits_gemv_f16_scales_f16_zp_pipe(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round)
{
matmul_nbits_gemv_f16_scales_f16_pipe_tpl<true>(activation, packed, scales_raw, zero_points, bias, output, k, n, block_size, k_blocks, blob_size, zp_row_bytes, scales_fp16, bias_post_round);
}
// each reducing a strided subset of the 256-wide K steps, then summing their
// fp32 partials through shared memory. The launch grid is K_SPLIT x larger than
// the single-warp `_zp` kernel, which fills the SMs on this grid-starved,
// latency-bound decode GEMV. The fp32 partial sum is a new block-sum
// association, so results are near-equal (not byte-identical) to the
// single-warp kernel.
template <bool HasZp, int K_SPLIT>
__device__ __forceinline__ void matmul_nbits_gemv_f16_scales_f16_splitk_tpl(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
void* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round,
const int out_bf16)
{
(void)block_size;
(void)scales_fp16;
const __half* __restrict__ scales =
reinterpret_cast<const __half*>(scales_raw);
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int warps_per_block = (int)blockDim.x >> 5;
const int cols_per_block = warps_per_block / K_SPLIT;
const int col_local = warp / K_SPLIT;
const int ks = warp % K_SPLIT;
const int column = (int)blockIdx.x * cols_per_block + col_local;
__shared__ float partials[8][K_SPLIT];
__half2 sum0 = __float2half2_rn(0.0f);
__half2 sum1 = __float2half2_rn(0.0f);
__half2 sum2 = __float2half2_rn(0.0f);
__half2 sum3 = __float2half2_rn(0.0f);
float tail = 0.0f;
if (column < n) {
const int lane_depth = lane * 8;
int depth_base = ks * 256;
const __half* activation_ptr = activation + depth_base + lane_depth;
const unsigned char* packed_ptr =
packed + (long)column * k_blocks * blob_size +
(long)(depth_base >> 5) * blob_size + lane * 4;
const __half* scale_ptr =
scales + (long)column * k_blocks + (depth_base >> 5) + (lane >> 2);
for (; depth_base < k; depth_base += K_SPLIT * 256) {
const int depth = depth_base + lane_depth;
if (depth >= k) {
activation_ptr += K_SPLIT * 256;
packed_ptr += K_SPLIT * 128;
scale_ptr += K_SPLIT * 8;
continue;
}
const unsigned int packed_word =
*reinterpret_cast<const unsigned int*>(packed_ptr);
const int block = (depth_base >> 5) + (lane >> 2);
const unsigned int sub2 =
block_sub2<HasZp>(zero_points, column, block, zp_row_bytes);
if (depth + 8 <= k) {
accumulate_int4x8_f16_zp(
packed_word, activation_ptr, *scale_ptr, sub2, sum0, sum1, sum2, sum3);
} else {
const float scale = __half2float(*scale_ptr);
const int zero_point =
block_zp<HasZp>(zero_points, column, block, zp_row_bytes);
const int valid = k - depth;
#pragma unroll
for (int i = 0; i < 8; ++i) {
if (i < valid) {
const int q =
(int)((packed_word >> (i * 4)) & 15u) - zero_point;
tail += (float)q * __half2float(activation_ptr[i]) * scale;
}
}
}
activation_ptr += K_SPLIT * 256;
packed_ptr += K_SPLIT * 128;
scale_ptr += K_SPLIT * 8;
}
}
const float2 value04 = __half22float2(sum0);
const float2 value15 = __half22float2(sum1);
const float2 value26 = __half22float2(sum2);
const float2 value37 = __half22float2(sum3);
float value = tail + value04.x;
value += value15.x;
value += value26.x;
value += value37.x;
value += value04.y;
value += value15.y;
value += value26.y;
value += value37.y;
value = warp_sum(value);
if (lane == 0) {
partials[col_local][ks] = (column < n) ? value : 0.0f;
}
__syncthreads();
if (ks == 0 && lane == 0 && column < n) {
float acc = 0.0f;
#pragma unroll
for (int s = 0; s < K_SPLIT; ++s) {
acc += partials[col_local][s];
}
matmul_nbits_store_narrowed(
output, column, fold_bias_f16(acc, bias, column, bias_post_round),
out_bf16);
}
}
extern "C" __global__ void matmul_nbits_gemv_f16_scales_f16_splitk(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
void* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round,
const int out_bf16)
{
matmul_nbits_gemv_f16_scales_f16_splitk_tpl<false, 2>(
activation, packed, scales_raw, zero_points, bias, output, k, n,
block_size, k_blocks, blob_size, zp_row_bytes, scales_fp16,
bias_post_round, out_bf16);
}
extern "C" __global__ void matmul_nbits_gemv_f16_scales_f16_zp_splitk(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
void* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round,
const int out_bf16)
{
matmul_nbits_gemv_f16_scales_f16_splitk_tpl<true, 2>(
activation, packed, scales_raw, zero_points, bias, output, k, n,
block_size, k_blocks, blob_size, zp_row_bytes, scales_fp16,
bias_post_round, out_bf16);
}
// Deep-split sibling of `matmul_nbits_gemv_f16_scales_f16_zp_splitk` for output
// widths so narrow that K_SPLIT=2 still leaves the device under one wave. The
// GQA k/v projection is the case: N=256 at 8 warps/CTA and K_SPLIT=2 covers 4
// columns per block, so the whole launch is 64 blocks on a 108-SM A100 and the
// kernel runs at latency rather than at bandwidth. Measured 28 GB/s there — the
// same K at N=4096 moved 16x the bytes in 21% LESS time, which no bandwidth
// explanation survives. K_SPLIT=8 makes each CTA cover one column, taking that
// launch to 256 blocks. Same accumulation as the K_SPLIT=2 entry, one more level
// of fp32 partial summing.
extern "C" __global__ void matmul_nbits_gemv_f16_scales_f16_zp_splitk8(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
void* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round,
const int out_bf16)
{
matmul_nbits_gemv_f16_scales_f16_splitk_tpl<true, 8>(
activation, packed, scales_raw, zero_points, bias, output, k, n,
block_size, k_blocks, blob_size, zp_row_bytes, scales_fp16,
bias_post_round, out_bf16);
}
// Prefetch-pipelined sibling of `matmul_nbits_gemv_f16_scales_f16_splitk_tpl`.
// The split-K kernel fills idle SMs and adds per-COLUMN memory-level parallelism
// (K_SPLIT warps per column), but each lane still walks its weight-word chain
// with a single 32-bit global load in flight -> a load->accumulate->load
// dependency that stalls on the ~Long-Scoreboard weight-load latency the pipe
// kernel diagnosed (DRAM ~24% of an H200's 4.8 TB/s, math not the bottleneck).
// This variant adds per-LANE MLP on top of split-K by holding PF prefetched
// weight words in a register-resident shift register (manual rotation, no
// dynamic-indexed array -> no local spill), so PF independent global loads stay
// in flight per lane while the fp16 dequant/FMA of the current word proceeds.
//
// The lane->nibble mapping, the `accumulate_int4x8_f16_zp` calls, the per-split
// accumulation order and the shared-mem cross-split reduction are UNCHANGED, so
// the output is bit-identical to `matmul_nbits_gemv_f16_scales_f16_splitk_tpl`
// (which is already near-equal, not byte-identical, to the single-warp pipe
// entry). Only the timing of the weight loads changes.
template <bool HasZp, int K_SPLIT>
__device__ __forceinline__ void matmul_nbits_gemv_f16_scales_f16_splitk_pf_tpl(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
void* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round,
const int out_bf16)
{
(void)block_size;
(void)scales_fp16;
constexpr int PF = 3; // weight words in flight per lane
const __half* __restrict__ scales =
reinterpret_cast<const __half*>(scales_raw);
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int warps_per_block = (int)blockDim.x >> 5;
const int cols_per_block = warps_per_block / K_SPLIT;
const int col_local = warp / K_SPLIT;
const int ks = warp % K_SPLIT;
const int column = (int)blockIdx.x * cols_per_block + col_local;
__shared__ float partials[8][K_SPLIT];
__half2 sum0 = __float2half2_rn(0.0f);
__half2 sum1 = __float2half2_rn(0.0f);
__half2 sum2 = __float2half2_rn(0.0f);
__half2 sum3 = __float2half2_rn(0.0f);
float tail = 0.0f;
if (column < n) {
const int lane_depth = lane * 8;
// Depth this (lane, split) touches at step i is
// `depth_start + i * K_SPLIT * 256`; the weight word stride is
// `K_SPLIT * 128` packed bytes (256 depth = 128 bytes for block-32 int4).
const int depth_start = ks * 256 + lane_depth;
const unsigned char* packed_base =
packed + (long)column * k_blocks * blob_size +
(long)((ks * 256) >> 5) * blob_size + lane * 4;
// Full 8-nibble steps this (lane, split) walks (step i valid iff
// `depth_start + i*K_SPLIT*256 + 8 <= k`), identical to the scalar bound.
const int nfull =
(depth_start + 8 <= k) ? ((k - depth_start - 8) / (K_SPLIT * 256) + 1) : 0;
auto load_step = [&](int s) -> unsigned int {
return *reinterpret_cast<const unsigned int*>(
packed_base + (long)s * (K_SPLIT * 128));
};
// Prime the shift register; out-of-range slots are zero (never consumed).
unsigned int wbuf[PF];
#pragma unroll
for (int s = 0; s < PF; ++s) {
wbuf[s] = (s < nfull) ? load_step(s) : 0u;
}
const __half* activation_ptr = activation + depth_start;
const __half* scale_ptr =
scales + (long)column * k_blocks + ((ks * 256) >> 5) + (lane >> 2);
int block = (int)(((ks * 256) >> 5) + (lane >> 2));
for (int i = 0; i < nfull; ++i) {
const unsigned int packed_word = wbuf[0];
// Rotate and issue the load PF steps ahead BEFORE consuming, so PF
// independent loads stay in flight.
const int pf = i + PF;
#pragma unroll
for (int s = 0; s < PF - 1; ++s) {
wbuf[s] = wbuf[s + 1];
}
wbuf[PF - 1] = (pf < nfull) ? load_step(pf) : 0u;
const unsigned int sub2 =
block_sub2<HasZp>(zero_points, column, block, zp_row_bytes);
accumulate_int4x8_f16_zp(
packed_word, activation_ptr, *scale_ptr, sub2, sum0, sum1, sum2, sum3);
activation_ptr += K_SPLIT * 256;
scale_ptr += K_SPLIT * 8;
block += K_SPLIT * 8;
}
const int tail_depth = depth_start + nfull * (K_SPLIT * 256);
if (tail_depth < k) {
const unsigned int packed_word = load_step(nfull);
const float scale = __half2float(*scale_ptr);
const int zero_point =
block_zp<HasZp>(zero_points, column, block, zp_row_bytes);
const int valid = k - tail_depth;
#pragma unroll
for (int i = 0; i < 8; ++i) {
if (i < valid) {
const int q =
(int)((packed_word >> (i * 4)) & 15u) - zero_point;
tail += (float)q * __half2float(activation_ptr[i]) * scale;
}
}
}
}
const float2 value04 = __half22float2(sum0);
const float2 value15 = __half22float2(sum1);
const float2 value26 = __half22float2(sum2);
const float2 value37 = __half22float2(sum3);
float value = tail + value04.x;
value += value15.x;
value += value26.x;
value += value37.x;
value += value04.y;
value += value15.y;
value += value26.y;
value += value37.y;
value = warp_sum(value);
if (lane == 0) {
partials[col_local][ks] = (column < n) ? value : 0.0f;
}
__syncthreads();
if (ks == 0 && lane == 0 && column < n) {
float acc = 0.0f;
#pragma unroll
for (int s = 0; s < K_SPLIT; ++s) {
acc += partials[col_local][s];
}
matmul_nbits_store_narrowed(
output, column, fold_bias_f16(acc, bias, column, bias_post_round),
out_bf16);
}
}
extern "C" __global__ void matmul_nbits_gemv_f16_scales_f16_splitk_pf(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
void* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round,
const int out_bf16)
{
matmul_nbits_gemv_f16_scales_f16_splitk_pf_tpl<false, 2>(
activation, packed, scales_raw, zero_points, bias, output, k, n,
block_size, k_blocks, blob_size, zp_row_bytes, scales_fp16,
bias_post_round, out_bf16);
}
extern "C" __global__ void matmul_nbits_gemv_f16_scales_f16_zp_splitk_pf(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
void* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round,
const int out_bf16)
{
matmul_nbits_gemv_f16_scales_f16_splitk_pf_tpl<true, 2>(
activation, packed, scales_raw, zero_points, bias, output, k, n,
block_size, k_blocks, blob_size, zp_row_bytes, scales_fp16,
bias_post_round, out_bf16);
}
// Prefetch-pipelined sibling of the K_SPLIT=8 deep-split GQA k/v entry
// (`matmul_nbits_gemv_f16_scales_f16_zp_splitk8`). Same PF register shift-
// register transform as the K_SPLIT=2 `_pf` entries, just a deeper split, so it
// stays bit-identical to the plain `_splitk8` entry while hiding the per-lane
// weight-load latency that persists even after the grid is deepened.
extern "C" __global__ void matmul_nbits_gemv_f16_scales_f16_zp_splitk8_pf(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
void* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round,
const int out_bf16)
{
matmul_nbits_gemv_f16_scales_f16_splitk_pf_tpl<true, 8>(
activation, packed, scales_raw, zero_points, bias, output, k, n,
block_size, k_blocks, blob_size, zp_row_bytes, scales_fp16,
bias_post_round, out_bf16);
}
// Half4 view matching `skip_rmsnorm_f16_warp_half4` so the fused prologue below
// reduces the activation with the exact same chunking and rounding.
union MatMulNBitsSkipHalf4 {
unsigned long long raw;
__half2 pair[2];
};
// Scalar RMS-norm gamma load matching `skip_rmsnorm_f16_warp_half4`: gamma is
// only ever a final multiplicand (never part of the fp32 variance
// accumulation), so an fp32 gamma is read at full precision while an fp16 gamma
// keeps the half round-trip. This lets decoders that export gamma in fp32 (e.g.
// Phi-4-mini) take the fused RMS-norm-prologue GEMV path bit-identically to the
// standalone norm + GEMV pair.
__device__ __forceinline__ float load_rmsnorm_gamma(
const void* __restrict__ gamma,
const int gamma_is_half,
const int index)
{
return gamma_is_half
? __half2float(reinterpret_cast<const __half*>(gamma)[index])
: reinterpret_cast<const float*>(gamma)[index];
}
// General fp16/fp16-scales GEMV with a fused RMS-normalization prologue. The
// preceding GEMV's residual epilogue already produced the byte-identical
// residual sum that `SkipSimplifiedLayerNormalization` would emit as its
// residual output, so this kernel only has to (1) reduce that sum exactly as
// `skip_rmsnorm_f16_warp_half4` does, (2) write the normalized activation into
// shared memory with the same rounding, and (3) run the standard `scales_f16`
// int4 dot over that staged, normalized activation. Every arithmetic step
// mirrors the standalone norm + GEMV pair, so tokens stay bit-for-bit identical
// while the separate normalization kernel is removed from the decode graph.
template <bool HasZp, bool SplitK = false>
__device__ __forceinline__ void matmul_nbits_gemv_f16_scales_f16_rmsnorm_tpl(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const void* __restrict__ gamma,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int bias_post_round,
const int gamma_is_half,
const float epsilon)
{
// Normalized activation, staged 16-byte aligned so the dot below can reuse
// the `scales_f16` `uint4` activation loads unchanged.
extern __shared__ __align__(16) __half staged_activation[];
__shared__ float shared_inv_std;
__shared__ float partials[8][2];
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
// --- RMS reduction, byte-identical to `skip_rmsnorm_f16_warp_half4`. ---
if (warp == 0) {
const int chunks_per_lane = k / (32 * 4);
const unsigned long long* activation4 =
reinterpret_cast<const unsigned long long*>(activation);
float ss0 = 0.0f;
float ss1 = 0.0f;
float ss2 = 0.0f;
float ss3 = 0.0f;
for (int item = 0; item < chunks_per_lane; ++item) {
const int chunk = lane + item * 32;
MatMulNBitsSkipHalf4 residual;
residual.raw = activation4[chunk];
const float2 rounded0 = __half22float2(residual.pair[0]);
const float2 rounded1 = __half22float2(residual.pair[1]);
ss0 += rounded0.x * rounded0.x;
ss1 += rounded0.y * rounded0.y;
ss2 += rounded1.x * rounded1.x;
ss3 += rounded1.y * rounded1.y;
}
float ss = (ss0 + ss1) + (ss2 + ss3);
for (int off = 16; off > 0; off >>= 1) {
ss += __shfl_down_sync(0xffffffffu, ss, off);
}
if (lane == 0) {
shared_inv_std = 1.0f / sqrtf(ss / (float)k + epsilon);
}
}
__syncthreads();
const float inv_std = shared_inv_std;
// --- Normalized activation, matching the norm kernel's rounded output. ---
for (int j = tid; j < k; j += (int)blockDim.x) {
const float residual = __half2float(activation[j]);
const float scale = load_rmsnorm_gamma(gamma, gamma_is_half, j);
staged_activation[j] = __float2half((residual * inv_std) * scale);
}
__syncthreads();
// --- Standard `scales_f16` int4 dot over the staged, normalized input. ---
const int warps_per_block = (int)blockDim.x >> 5;
const int columns_per_block = SplitK ? warps_per_block / 2 : warps_per_block;
const int column_local = SplitK ? warp / 2 : warp;
const int k_split = SplitK ? warp & 1 : 0;
const int column = (int)blockIdx.x * columns_per_block + column_local;
const __half* __restrict__ scales =
reinterpret_cast<const __half*>(scales_raw);
__half2 sum0 = __float2half2_rn(0.0f);
__half2 sum1 = __float2half2_rn(0.0f);
__half2 sum2 = __float2half2_rn(0.0f);
__half2 sum3 = __float2half2_rn(0.0f);
float tail = 0.0f;
if (column < n) {
const int lane_depth = lane * 8;
int depth_base = k_split * 256;
const __half* activation_ptr = staged_activation + depth_base + lane_depth;
const unsigned char* packed_ptr =
packed + (long)column * k_blocks * blob_size
+ (long)(depth_base >> 5) * blob_size + lane * 4;
const __half* scale_ptr =
scales + (long)column * k_blocks + (depth_base >> 5) + (lane >> 2);
const int depth_step = SplitK ? 512 : 256;
for (; depth_base + lane_depth + 8 <= k; depth_base += depth_step) {
const unsigned int packed_word =
*reinterpret_cast<const unsigned int*>(packed_ptr);
const int block = (depth_base >> 5) + (lane >> 2);
const unsigned int sub2 =
block_sub2<HasZp>(zero_points, column, block, zp_row_bytes);
accumulate_int4x8_f16_zp(
packed_word,
activation_ptr,
*scale_ptr,
sub2,
sum0,
sum1,
sum2,
sum3);
activation_ptr += depth_step;
packed_ptr += depth_step / 2;
scale_ptr += depth_step / 32;
}
const int tail_depth = depth_base + lane_depth;
if (tail_depth < k) {
const unsigned int packed_word =
*reinterpret_cast<const unsigned int*>(packed_ptr);
const float scale = __half2float(*scale_ptr);
const int tail_block = (depth_base >> 5) + (lane >> 2);
const int zero_point =
block_zp<HasZp>(zero_points, column, tail_block, zp_row_bytes);
const int valid = min(8, k - tail_depth);
#pragma unroll
for (int i = 0; i < 8; ++i) {
if (i < valid) {
const int q = (int)((packed_word >> (i * 4)) & 15u) - zero_point;
tail += (float)q * __half2float(activation_ptr[i]) * scale;
}
}
}
}
const float2 value04 = __half22float2(sum0);
const float2 value15 = __half22float2(sum1);
const float2 value26 = __half22float2(sum2);
const float2 value37 = __half22float2(sum3);
float value = tail + value04.x;
value += value15.x;
value += value26.x;
value += value37.x;
value += value04.y;
value += value15.y;
value += value26.y;
value += value37.y;
value = warp_sum(value);
if constexpr (SplitK) {
if (lane == 0) {
partials[column_local][k_split] = column < n ? value : 0.0f;
}
__syncthreads();
if (k_split == 0 && lane == 0 && column < n) {
output[column] = fold_bias_f16(
partials[column_local][0] + partials[column_local][1],
bias,
column,
bias_post_round);
}
} else if (lane == 0 && column < n) {
output[column] = fold_bias_f16(value, bias, column, bias_post_round);
}
}
extern "C" __global__ void matmul_nbits_gemv_f16_scales_f16_rmsnorm(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const void* __restrict__ gamma,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int bias_post_round,
const int gamma_is_half,
const float epsilon)
{
matmul_nbits_gemv_f16_scales_f16_rmsnorm_tpl<false>(activation, packed, scales_raw, zero_points, gamma, bias, output, k, n, k_blocks, blob_size, zp_row_bytes, bias_post_round, gamma_is_half, epsilon);
}
extern "C" __global__ void matmul_nbits_gemv_f16_scales_f16_rmsnorm_splitk(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const void* __restrict__ gamma,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int bias_post_round,
const int gamma_is_half,
const float epsilon)
{
matmul_nbits_gemv_f16_scales_f16_rmsnorm_tpl<false, true>(activation, packed, scales_raw, zero_points, gamma, bias, output, k, n, k_blocks, blob_size, zp_row_bytes, bias_post_round, gamma_is_half, epsilon);
}
extern "C" __global__ void matmul_nbits_gemv_f16_scales_f16_rmsnorm_zp(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const void* __restrict__ gamma,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int bias_post_round,
const int gamma_is_half,
const float epsilon)
{
matmul_nbits_gemv_f16_scales_f16_rmsnorm_tpl<true>(activation, packed, scales_raw, zero_points, gamma, bias, output, k, n, k_blocks, blob_size, zp_row_bytes, bias_post_round, gamma_is_half, epsilon);
}
// Compile-time-specialized per-block int8 zero point. `HasZp == false`
// (symmetric int8) folds to the constant 128 with no load — mirroring the int4
// `block_zp` helper — so a future symmetric-int8 model keeps the constant
// subtrahend and never pays the per-block occupancy cost the int4 path shed.
template <bool HasZp>
__device__ __forceinline__ int block_zp_int8(
const unsigned char* __restrict__ zero_points,
const long column,
const int block,
const int k_blocks)
{
if (!HasZp) {
return 128;
}
return (int)zero_points[column * k_blocks + block];
}
// INT8 sibling of `matmul_nbits_gemv_f16_scales_f16_rmsnorm`. The RMS reduction
// and normalized-activation staging are byte-identical to the int4 fused kernel
// (and to the standalone `skip_rmsnorm_f16_warp_half4`); only the quantized dot
// differs, reusing the exact block-32 int8 dequant work split from
// `matmul_nbits_gemv_int8_f16` (one byte per weight, per-block uint8 zero point
// defaulting to 128, fp32 accumulation). Specialized on `HasZp` like the int4
// sibling so the symmetric case emits no per-block zero-point load.
template <bool HasZp>
__device__ __forceinline__ void matmul_nbits_gemv_int8_f16_scales_f16_rmsnorm_tpl(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const void* __restrict__ gamma,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int bias_post_round,
const int gamma_is_half,
const float epsilon)
{
extern __shared__ __align__(16) __half staged_activation[];
__shared__ float shared_inv_std;
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
// --- RMS reduction, byte-identical to `skip_rmsnorm_f16_warp_half4`. ---
if (warp == 0) {
const int chunks_per_lane = k / (32 * 4);
const unsigned long long* activation4 =
reinterpret_cast<const unsigned long long*>(activation);
float ss0 = 0.0f;
float ss1 = 0.0f;
float ss2 = 0.0f;
float ss3 = 0.0f;
for (int item = 0; item < chunks_per_lane; ++item) {
const int chunk = lane + item * 32;
MatMulNBitsSkipHalf4 residual;
residual.raw = activation4[chunk];
const float2 rounded0 = __half22float2(residual.pair[0]);
const float2 rounded1 = __half22float2(residual.pair[1]);
ss0 += rounded0.x * rounded0.x;
ss1 += rounded0.y * rounded0.y;
ss2 += rounded1.x * rounded1.x;
ss3 += rounded1.y * rounded1.y;
}
float ss = (ss0 + ss1) + (ss2 + ss3);
for (int off = 16; off > 0; off >>= 1) {
ss += __shfl_down_sync(0xffffffffu, ss, off);
}
if (lane == 0) {
shared_inv_std = 1.0f / sqrtf(ss / (float)k + epsilon);
}
}
__syncthreads();
const float inv_std = shared_inv_std;
// --- Normalized activation, matching the norm kernel's rounded output. ---
for (int j = tid; j < k; j += (int)blockDim.x) {
const float residual = __half2float(activation[j]);
const float scale = load_rmsnorm_gamma(gamma, gamma_is_half, j);
staged_activation[j] = __float2half((residual * inv_std) * scale);
}
__syncthreads();
// --- INT8 dot over the staged, normalized input (mirrors the non-fused
// `matmul_nbits_gemv_int8_f16` work split, fp32 accumulation). ---
const int columns_per_block = (int)blockDim.x >> 5;
const int column = (int)blockIdx.x * columns_per_block + warp;
const __half* __restrict__ scales =
reinterpret_cast<const __half*>(scales_raw);
float value = 0.0f;
if (column < n) {
const int quarter = lane & 3;
for (int block_base = 0; block_base < k_blocks; block_base += 8) {
const int block = block_base + (lane >> 2);
float block_partial = 0.0f;
if (block < k_blocks) {
const int zero_point =
block_zp_int8<HasZp>(zero_points, column, block, k_blocks);
const int depth = block * 32 + quarter * 8;
const long packed_start =
((long)column * k_blocks + block) * 32 + quarter * 8;
if (depth + 8 <= k) {
const uint2 packed_word =
*reinterpret_cast<const uint2*>(packed + packed_start);
const unsigned char* bytes =
reinterpret_cast<const unsigned char*>(&packed_word);
const uint4 act =
*reinterpret_cast<const uint4*>(staged_activation + depth);
const __half* acth = reinterpret_cast<const __half*>(&act);
#pragma unroll
for (int i = 0; i < 8; ++i) {
block_partial += ((float)(int)bytes[i] - (float)zero_point)
* __half2float(acth[i]);
}
} else if (depth < k) {
const int valid = min(8, k - depth);
#pragma unroll
for (int i = 0; i < 8; ++i) {
if (i < valid) {
const int quantized = (int)packed[packed_start + i];
block_partial += ((float)quantized - (float)zero_point)
* __half2float(staged_activation[depth + i]);
}
}
}
}
block_partial += __shfl_down_sync(0xffffffffu, block_partial, 2, 4);
block_partial += __shfl_down_sync(0xffffffffu, block_partial, 1, 4);
if (quarter == 0 && block < k_blocks) {
const float scale =
__half2float(scales[(long)column * k_blocks + block]);
value += block_partial * scale;
}
}
}
value = warp_sum(value);
if (lane == 0 && column < n) {
output[column] = fold_bias_f16(value, bias, column, bias_post_round);
}
}
extern "C" __global__ void matmul_nbits_gemv_int8_f16_scales_f16_rmsnorm(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const void* __restrict__ gamma,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int bias_post_round,
const int gamma_is_half,
const float epsilon)
{
matmul_nbits_gemv_int8_f16_scales_f16_rmsnorm_tpl<false>(activation, packed, scales_raw, zero_points, gamma, bias, output, k, n, k_blocks, bias_post_round, gamma_is_half, epsilon);
}
extern "C" __global__ void matmul_nbits_gemv_int8_f16_scales_f16_rmsnorm_zp(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const void* __restrict__ gamma,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int bias_post_round,
const int gamma_is_half,
const float epsilon)
{
matmul_nbits_gemv_int8_f16_scales_f16_rmsnorm_tpl<true>(activation, packed, scales_raw, zero_points, gamma, bias, output, k, n, k_blocks, bias_post_round, gamma_is_half, epsilon);
}
// Standalone RMS-normalization prologue for the M>1 prefill path of the fused
// GEMV. It reproduces `skip_rmsnorm_f16_warp_half4` (minus the residual add,
// which the preceding GEMV's epilogue already applied) bit-for-bit: identical
// half4 chunking, identical `(ss0+ss1)+(ss2+ss3)` reduction, identical warp
// shuffle, and identical `__floats2half2_rn` output rounding. One warp
// normalizes one token row into `normalized`, which the portable tiled GEMM
// then consumes exactly as it would the standalone norm's fp16 output.
extern "C" __global__ void matmul_nbits_rmsnorm_f16_warp_half4(
const __half* __restrict__ activation,
const void* __restrict__ gamma,
__half* __restrict__ normalized,
const int norm_size,
const int num_groups,
const int gamma_is_half,
const float epsilon)
{
const int g = (int)blockIdx.x;
if (g >= num_groups) return;
const long base = (long)g * norm_size;
const int lane = (int)threadIdx.x;
const int chunks_per_lane = norm_size / (32 * 4);
const unsigned long long* activation4 =
reinterpret_cast<const unsigned long long*>(activation + base);
const unsigned long long* gamma4 =
reinterpret_cast<const unsigned long long*>(gamma);
unsigned long long* normalized4 =
reinterpret_cast<unsigned long long*>(normalized + base);
float ss0 = 0.0f;
float ss1 = 0.0f;
float ss2 = 0.0f;
float ss3 = 0.0f;
for (int item = 0; item < chunks_per_lane; ++item) {
const int chunk = lane + item * 32;
MatMulNBitsSkipHalf4 residual;
residual.raw = activation4[chunk];
const float2 rounded0 = __half22float2(residual.pair[0]);
const float2 rounded1 = __half22float2(residual.pair[1]);
ss0 += rounded0.x * rounded0.x;
ss1 += rounded0.y * rounded0.y;
ss2 += rounded1.x * rounded1.x;
ss3 += rounded1.y * rounded1.y;
}
float ss = (ss0 + ss1) + (ss2 + ss3);
for (int off = 16; off > 0; off >>= 1) {
ss += __shfl_down_sync(0xffffffffu, ss, off);
}
float inv_std = 0.0f;
if (lane == 0) {
inv_std = 1.0f / sqrtf(ss / (float)norm_size + epsilon);
}
inv_std = __shfl_sync(0xffffffffu, inv_std, 0);
for (int item = 0; item < chunks_per_lane; ++item) {
const int chunk = lane + item * 32;
MatMulNBitsSkipHalf4 residual;
MatMulNBitsSkipHalf4 output;
residual.raw = activation4[chunk];
const float2 value0 = __half22float2(residual.pair[0]);
const float2 value1 = __half22float2(residual.pair[1]);
// gamma is only a final multiplicand: an fp16 gamma keeps the wide
// half4 load, an fp32 gamma is read at full precision (matching the
// standalone `skip_rmsnorm_f16_warp_half4`), so fp32-gamma decoders fuse.
float scale0x, scale0y, scale1x, scale1y;
if (gamma_is_half) {
MatMulNBitsSkipHalf4 scale;
scale.raw = gamma4[chunk];
const float2 scale0 = __half22float2(scale.pair[0]);
const float2 scale1 = __half22float2(scale.pair[1]);
scale0x = scale0.x;
scale0y = scale0.y;
scale1x = scale1.x;
scale1y = scale1.y;
} else {
const int j = chunk << 2;
const float* gamma_f = reinterpret_cast<const float*>(gamma);
scale0x = gamma_f[j];
scale0y = gamma_f[j + 1];
scale1x = gamma_f[j + 2];
scale1y = gamma_f[j + 3];
}
output.pair[0] = __floats2half2_rn(
value0.x * inv_std * scale0x,
value0.y * inv_std * scale0y);
output.pair[1] = __floats2half2_rn(
value1.x * inv_std * scale1x,
value1.y * inv_std * scale1y);
normalized4[chunk] = output.raw;
}
}
// SwiGLU activation, byte-identical to the standalone `op_silu` in the
// elementwise kernels: silu(x) = x * sigmoid(x), evaluated in the same
// rounding-stable form so the paired epilogue reproduces the two-op tokens.
__device__ __forceinline__ float gate_up_silu_f32(float x)
{
if (x >= 0.0f) {
const float denominator = __fadd_rn(1.0f, (float)exp((double)-x));
return __fdiv_rn(x, denominator);
}
const float e = (float)exp((double)x);
const float numerator = __fmul_rn(x, e);
return __fdiv_rn(numerator, __fadd_rn(1.0f, e));
}
__device__ __forceinline__ float gate_up_decomposed_silu_f32(float x)
{
float sigmoid;
if (x >= 0.0f) {
sigmoid = 1.0f / (1.0f + (float)exp((double)-x));
} else {
const float e = (float)exp((double)x);
sigmoid = e / (1.0f + e);
}
const float sigmoid_h = __half2float(__float2half_rn(sigmoid));
return __half2float(__float2half_rn(__fmul_rn(x, sigmoid_h)));
}
// Paired gate/up projection + SwiGLU. One warp computes column `column` of BOTH
// the gate and up projections (which share the same activation and the block-32
// fp16 layout of `matmul_nbits_gemv_f16_scales_f16`), then writes
// silu(gate)*up directly. The activation is permuted once per K-tile and reused
// by both accumulators, so the two GEMVs read the activation from registers
// exactly once. The epilogue reproduces the standalone two-op numerics
// (`fp16(gate_acc)`, `fp16(up_acc)`, then `fp16(silu(gate_h)*up_h)`) so greedy
// decoding stays byte-identical. Register-only + warp shuffles: no shared
// memory, so it is portable to sm_53+ and safe on small SMs (no >48KB opt-in).
template <bool HasZp, bool Decomposed, bool FusedSym = false>
__device__ __forceinline__ void matmul_nbits_gemv_f16_gate_up_swiglu_tpl(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed_gate,
const __half* __restrict__ scales_gate,
const unsigned char* __restrict__ packed_up,
const __half* __restrict__ scales_up,
const unsigned char* __restrict__ zero_points_gate,
const unsigned char* __restrict__ zero_points_up,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int blob_size,
const int zp_row_bytes)
{
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int columns_per_block = (int)blockDim.x >> 5;
const int column = (int)blockIdx.x * columns_per_block + warp;
__half2 g0 = __float2half2_rn(0.0f);
__half2 g1 = __float2half2_rn(0.0f);
__half2 g2 = __float2half2_rn(0.0f);
__half2 g3 = __float2half2_rn(0.0f);
__half2 u0 = __float2half2_rn(0.0f);
__half2 u1 = __float2half2_rn(0.0f);
__half2 u2 = __float2half2_rn(0.0f);
__half2 u3 = __float2half2_rn(0.0f);
float gate_tail = 0.0f;
float up_tail = 0.0f;
if (column < n) {
const int lane_depth = lane * 8;
const __half* activation_ptr = activation + lane_depth;
const unsigned char* packed_gate_ptr =
packed_gate + (long)column * k_blocks * blob_size + lane * 4;
const unsigned char* packed_up_ptr =
packed_up + (long)column * k_blocks * blob_size + lane * 4;
const __half* scale_gate_ptr =
scales_gate + (long)column * k_blocks + (lane >> 2);
const __half* scale_up_ptr =
scales_up + (long)column * k_blocks + (lane >> 2);
int depth_base = 0;
for (; depth_base + lane_depth + 8 <= k; depth_base += 256) {
// Permute the shared activation once; both projections reuse it.
const uint4 permuted = permute_activation_f16x8(activation_ptr);
const int block = (depth_base >> 5) + (lane >> 2);
const unsigned int gate_sub2 =
block_sub2<HasZp>(zero_points_gate, column, block, zp_row_bytes);
const unsigned int up_sub2 =
block_sub2<HasZp>(zero_points_up, column, block, zp_row_bytes);
const unsigned int gate_word =
*reinterpret_cast<const unsigned int*>(packed_gate_ptr);
const unsigned int up_word =
*reinterpret_cast<const unsigned int*>(packed_up_ptr);
if constexpr (FusedSym && !HasZp) {
// Symmetric fast dequant: fold `- 8` into the bias constants
// (byte-identical, four fewer f16x2 ops/word). `gate_sub2`/
// `up_sub2` are the constant 8 here and go unused.
accumulate_int4x8_f16_permuted_sym8(
gate_word, permuted, *scale_gate_ptr, g0, g1, g2, g3);
accumulate_int4x8_f16_permuted_sym8(
up_word, permuted, *scale_up_ptr, u0, u1, u2, u3);
} else {
accumulate_int4x8_f16_permuted_zp(
gate_word, permuted, *scale_gate_ptr, gate_sub2, g0, g1, g2, g3);
accumulate_int4x8_f16_permuted_zp(
up_word, permuted, *scale_up_ptr, up_sub2, u0, u1, u2, u3);
}
activation_ptr += 256;
packed_gate_ptr += 128;
packed_up_ptr += 128;
scale_gate_ptr += 8;
scale_up_ptr += 8;
}
const int tail_depth = depth_base + lane_depth;
if (tail_depth < k) {
const unsigned int gate_word =
*reinterpret_cast<const unsigned int*>(packed_gate_ptr);
const unsigned int up_word =
*reinterpret_cast<const unsigned int*>(packed_up_ptr);
const float gate_scale = __half2float(*scale_gate_ptr);
const float up_scale = __half2float(*scale_up_ptr);
const int tail_block = (depth_base >> 5) + (lane >> 2);
const int gate_zp =
block_zp<HasZp>(zero_points_gate, column, tail_block, zp_row_bytes);
const int up_zp =
block_zp<HasZp>(zero_points_up, column, tail_block, zp_row_bytes);
const int valid = min(8, k - tail_depth);
#pragma unroll
for (int i = 0; i < 8; ++i) {
if (i < valid) {
const float a = __half2float(activation_ptr[i]);
const int qg = (int)((gate_word >> (i * 4)) & 15u) - gate_zp;
const int qu = (int)((up_word >> (i * 4)) & 15u) - up_zp;
gate_tail += (float)qg * a * gate_scale;
up_tail += (float)qu * a * up_scale;
}
}
}
}
// Reduce each accumulator in the exact term order of the standalone
// `matmul_nbits_gemv_f16_scales_f16` epilogue so the pre-round sums match.
const float2 g04 = __half22float2(g0);
const float2 g15 = __half22float2(g1);
const float2 g26 = __half22float2(g2);
const float2 g37 = __half22float2(g3);
float gate_value = gate_tail + g04.x;
gate_value += g15.x;
gate_value += g26.x;
gate_value += g37.x;
gate_value += g04.y;
gate_value += g15.y;
gate_value += g26.y;
gate_value += g37.y;
gate_value = warp_sum(gate_value);
const float2 u04 = __half22float2(u0);
const float2 u15 = __half22float2(u1);
const float2 u26 = __half22float2(u2);
const float2 u37 = __half22float2(u3);
float up_value = up_tail + u04.x;
up_value += u15.x;
up_value += u26.x;
up_value += u37.x;
up_value += u04.y;
up_value += u15.y;
up_value += u26.y;
up_value += u37.y;
up_value = warp_sum(up_value);
if (lane == 0 && column < n) {
// Round each projection to fp16 first (matching the separate GEMV
// stores), then compute silu(gate)*up and round once — identical to the
// standalone silu_mul_f16 kernel fed by the two GEMV outputs.
const float gate_h = __half2float(__float2half(gate_value));
const float up_h = __half2float(__float2half(up_value));
const float silu_h = Decomposed
? gate_up_decomposed_silu_f32(gate_h)
: gate_up_silu_f32(gate_h);
output[column] = __float2half_rn(__fmul_rn(silu_h, up_h));
}
}
extern "C" __global__ void matmul_nbits_gemv_f16_gate_up_swiglu(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed_gate,
const __half* __restrict__ scales_gate,
const unsigned char* __restrict__ packed_up,
const __half* __restrict__ scales_up,
const unsigned char* __restrict__ zero_points_gate,
const unsigned char* __restrict__ zero_points_up,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int blob_size,
const int zp_row_bytes)
{
matmul_nbits_gemv_f16_gate_up_swiglu_tpl<false, false>(activation, packed_gate, scales_gate, packed_up, scales_up, zero_points_gate, zero_points_up, output, k, n, k_blocks, blob_size, zp_row_bytes);
}
extern "C" __global__ void matmul_nbits_gemv_f16_gate_up_decomposed_swiglu(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed_gate,
const __half* __restrict__ scales_gate,
const unsigned char* __restrict__ packed_up,
const __half* __restrict__ scales_up,
const unsigned char* __restrict__ zero_points_gate,
const unsigned char* __restrict__ zero_points_up,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int blob_size,
const int zp_row_bytes)
{
matmul_nbits_gemv_f16_gate_up_swiglu_tpl<false, true>(activation, packed_gate, scales_gate, packed_up, scales_up, zero_points_gate, zero_points_up, output, k, n, k_blocks, blob_size, zp_row_bytes);
}
extern "C" __global__ void matmul_nbits_gemv_f16_gate_up_swiglu_zp(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed_gate,
const __half* __restrict__ scales_gate,
const unsigned char* __restrict__ packed_up,
const __half* __restrict__ scales_up,
const unsigned char* __restrict__ zero_points_gate,
const unsigned char* __restrict__ zero_points_up,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int blob_size,
const int zp_row_bytes)
{
matmul_nbits_gemv_f16_gate_up_swiglu_tpl<true, false>(activation, packed_gate, scales_gate, packed_up, scales_up, zero_points_gate, zero_points_up, output, k, n, k_blocks, blob_size, zp_row_bytes);
}
extern "C" __global__ void matmul_nbits_gemv_f16_gate_up_decomposed_swiglu_zp(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed_gate,
const __half* __restrict__ scales_gate,
const unsigned char* __restrict__ packed_up,
const __half* __restrict__ scales_up,
const unsigned char* __restrict__ zero_points_gate,
const unsigned char* __restrict__ zero_points_up,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int blob_size,
const int zp_row_bytes)
{
matmul_nbits_gemv_f16_gate_up_swiglu_tpl<true, true>(activation, packed_gate, scales_gate, packed_up, scales_up, zero_points_gate, zero_points_up, output, k, n, k_blocks, blob_size, zp_row_bytes);
}
// Fused-symmetric (`ONNX_GENAI_GATEUP_VEC`) siblings of the two SYMMETRIC
// gate/up SwiGLU entries above. Byte-identical — only the dequant folds the
// `- 8` symmetric zero point into the bias constants (see
// `int4x8_to_half2x4_sym8`), issuing four fewer f16x2 ops per weight word to
// relieve the issue-bound decode GEMV. No asymmetric `_vec` variant: the `_zp`
// entries carry a per-block zero point that cannot be folded to a constant.
extern "C" __global__ void matmul_nbits_gemv_f16_gate_up_swiglu_vec(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed_gate,
const __half* __restrict__ scales_gate,
const unsigned char* __restrict__ packed_up,
const __half* __restrict__ scales_up,
const unsigned char* __restrict__ zero_points_gate,
const unsigned char* __restrict__ zero_points_up,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int blob_size,
const int zp_row_bytes)
{
matmul_nbits_gemv_f16_gate_up_swiglu_tpl<false, false, true>(activation, packed_gate, scales_gate, packed_up, scales_up, zero_points_gate, zero_points_up, output, k, n, k_blocks, blob_size, zp_row_bytes);
}
extern "C" __global__ void matmul_nbits_gemv_f16_gate_up_decomposed_swiglu_vec(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed_gate,
const __half* __restrict__ scales_gate,
const unsigned char* __restrict__ packed_up,
const __half* __restrict__ scales_up,
const unsigned char* __restrict__ zero_points_gate,
const unsigned char* __restrict__ zero_points_up,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int blob_size,
const int zp_row_bytes)
{
matmul_nbits_gemv_f16_gate_up_swiglu_tpl<false, true, true>(activation, packed_gate, scales_gate, packed_up, scales_up, zero_points_gate, zero_points_up, output, k, n, k_blocks, blob_size, zp_row_bytes);
}
// This is `matmul_nbits_gemv_f16_gate_up_swiglu` preceded by the exact prologue
// of `matmul_nbits_gemv_f16_scales_f16_rmsnorm`: the block reduces the shared
// activation (the residual sum the preceding GEMV epilogue already produced)
// once, stages the normalized activation into shared memory with the same
// rounding, and then both the gate and up GEMVs read that single staged,
// normalized activation. Doing the reduction once — rather than once per
// following GEMV — is the whole point of routing the fan-out-2 post-attention
// `SkipSimplifiedLayerNormalization` through the paired kernel. Every arithmetic
// step mirrors the standalone norm followed by the two-op gate/up SwiGLU, so
// greedy tokens stay bit-for-bit identical.
template <bool HasZp, bool Decomposed, bool FusedSym = false>
__device__ __forceinline__ void matmul_nbits_gemv_f16_gate_up_swiglu_rmsnorm_tpl(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed_gate,
const __half* __restrict__ scales_gate,
const unsigned char* __restrict__ packed_up,
const __half* __restrict__ scales_up,
const unsigned char* __restrict__ zero_points_gate,
const unsigned char* __restrict__ zero_points_up,
const void* __restrict__ gamma,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int gamma_is_half,
const float epsilon)
{
extern __shared__ __align__(16) __half staged_activation[];
__shared__ float shared_inv_std;
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
// --- RMS reduction, byte-identical to `skip_rmsnorm_f16_warp_half4`. ---
if (warp == 0) {
const int chunks_per_lane = k / (32 * 4);
const unsigned long long* activation4 =
reinterpret_cast<const unsigned long long*>(activation);
float ss0 = 0.0f;
float ss1 = 0.0f;
float ss2 = 0.0f;
float ss3 = 0.0f;
for (int item = 0; item < chunks_per_lane; ++item) {
const int chunk = lane + item * 32;
MatMulNBitsSkipHalf4 residual;
residual.raw = activation4[chunk];
const float2 rounded0 = __half22float2(residual.pair[0]);
const float2 rounded1 = __half22float2(residual.pair[1]);
ss0 += rounded0.x * rounded0.x;
ss1 += rounded0.y * rounded0.y;
ss2 += rounded1.x * rounded1.x;
ss3 += rounded1.y * rounded1.y;
}
float ss = (ss0 + ss1) + (ss2 + ss3);
for (int off = 16; off > 0; off >>= 1) {
ss += __shfl_down_sync(0xffffffffu, ss, off);
}
if (lane == 0) {
shared_inv_std = 1.0f / sqrtf(ss / (float)k + epsilon);
}
}
__syncthreads();
const float inv_std = shared_inv_std;
// --- Normalized activation, matching the norm kernel's rounded output. ---
for (int j = tid; j < k; j += (int)blockDim.x) {
const float residual = __half2float(activation[j]);
const float scale = load_rmsnorm_gamma(gamma, gamma_is_half, j);
staged_activation[j] = __float2half((residual * inv_std) * scale);
}
__syncthreads();
// --- Paired gate/up int4 dot over the staged, normalized activation. ---
const int columns_per_block = (int)blockDim.x >> 5;
const int column = (int)blockIdx.x * columns_per_block + warp;
__half2 g0 = __float2half2_rn(0.0f);
__half2 g1 = __float2half2_rn(0.0f);
__half2 g2 = __float2half2_rn(0.0f);
__half2 g3 = __float2half2_rn(0.0f);
__half2 u0 = __float2half2_rn(0.0f);
__half2 u1 = __float2half2_rn(0.0f);
__half2 u2 = __float2half2_rn(0.0f);
__half2 u3 = __float2half2_rn(0.0f);
float gate_tail = 0.0f;
float up_tail = 0.0f;
if (column < n) {
const int lane_depth = lane * 8;
const __half* activation_ptr = staged_activation + lane_depth;
const unsigned char* packed_gate_ptr =
packed_gate + (long)column * k_blocks * blob_size + lane * 4;
const unsigned char* packed_up_ptr =
packed_up + (long)column * k_blocks * blob_size + lane * 4;
const __half* scale_gate_ptr =
scales_gate + (long)column * k_blocks + (lane >> 2);
const __half* scale_up_ptr =
scales_up + (long)column * k_blocks + (lane >> 2);
int depth_base = 0;
if constexpr (HasZp) {
// Asymmetric (zero-point) path — the dominant Phi decode kernel, which
// ncu shows is Long-Scoreboard/global-load-latency bound. Software-
// pipeline the int4 gate/up weight loads: issue the next iteration's
// two 128-byte weight words before consuming the current ones so the
// load latency overlaps this iteration's compute. Pure scheduling
// change (identical accumulation order/ops) → bit-identical to the
// non-prefetched loop. Only the weight words are prefetched; also
// prefetching the small (L1/L2-resident) scales and per-block zero
// points pushed registers 48->56 and the occupancy loss erased the
// latency win. The symmetric (`HasZp == false`) path below keeps its
// exact original instruction stream so Qwen stays byte-identical with
// no register/occupancy change.
unsigned int gate_word_next =
*reinterpret_cast<const unsigned int*>(packed_gate_ptr);
unsigned int up_word_next =
*reinterpret_cast<const unsigned int*>(packed_up_ptr);
for (; depth_base + lane_depth + 8 <= k; depth_base += 256) {
const uint4 permuted = permute_activation_f16x8(activation_ptr);
const int block = (depth_base >> 5) + (lane >> 2);
const unsigned int gate_sub2 =
block_sub2<HasZp>(zero_points_gate, column, block, zp_row_bytes);
const unsigned int up_sub2 =
block_sub2<HasZp>(zero_points_up, column, block, zp_row_bytes);
const unsigned int gate_word = gate_word_next;
const unsigned int up_word = up_word_next;
if (depth_base + 256 + lane_depth + 8 <= k) {
gate_word_next = *reinterpret_cast<const unsigned int*>(
packed_gate_ptr + 128);
up_word_next = *reinterpret_cast<const unsigned int*>(
packed_up_ptr + 128);
}
accumulate_int4x8_f16_permuted_zp(
gate_word, permuted, *scale_gate_ptr, gate_sub2, g0, g1, g2, g3);
accumulate_int4x8_f16_permuted_zp(
up_word, permuted, *scale_up_ptr, up_sub2, u0, u1, u2, u3);
activation_ptr += 256;
packed_gate_ptr += 128;
packed_up_ptr += 128;
scale_gate_ptr += 8;
scale_up_ptr += 8;
}
} else {
for (; depth_base + lane_depth + 8 <= k; depth_base += 256) {
const uint4 permuted = permute_activation_f16x8(activation_ptr);
const int block = (depth_base >> 5) + (lane >> 2);
const unsigned int gate_word =
*reinterpret_cast<const unsigned int*>(packed_gate_ptr);
const unsigned int up_word =
*reinterpret_cast<const unsigned int*>(packed_up_ptr);
if constexpr (FusedSym) {
// Symmetric fast dequant: fold `- 8` into the bias constants
// (byte-identical, four fewer f16x2 ops/word).
accumulate_int4x8_f16_permuted_sym8(
gate_word, permuted, *scale_gate_ptr, g0, g1, g2, g3);
accumulate_int4x8_f16_permuted_sym8(
up_word, permuted, *scale_up_ptr, u0, u1, u2, u3);
} else {
const unsigned int gate_sub2 = block_sub2<HasZp>(
zero_points_gate, column, block, zp_row_bytes);
const unsigned int up_sub2 = block_sub2<HasZp>(
zero_points_up, column, block, zp_row_bytes);
accumulate_int4x8_f16_permuted_zp(
gate_word, permuted, *scale_gate_ptr, gate_sub2, g0, g1, g2, g3);
accumulate_int4x8_f16_permuted_zp(
up_word, permuted, *scale_up_ptr, up_sub2, u0, u1, u2, u3);
}
activation_ptr += 256;
packed_gate_ptr += 128;
packed_up_ptr += 128;
scale_gate_ptr += 8;
scale_up_ptr += 8;
}
}
const int tail_depth = depth_base + lane_depth;
if (tail_depth < k) {
const unsigned int gate_word =
*reinterpret_cast<const unsigned int*>(packed_gate_ptr);
const unsigned int up_word =
*reinterpret_cast<const unsigned int*>(packed_up_ptr);
const float gate_scale = __half2float(*scale_gate_ptr);
const float up_scale = __half2float(*scale_up_ptr);
const int tail_block = (depth_base >> 5) + (lane >> 2);
const int gate_zp =
block_zp<HasZp>(zero_points_gate, column, tail_block, zp_row_bytes);
const int up_zp =
block_zp<HasZp>(zero_points_up, column, tail_block, zp_row_bytes);
const int valid = min(8, k - tail_depth);
#pragma unroll
for (int i = 0; i < 8; ++i) {
if (i < valid) {
const float a = __half2float(activation_ptr[i]);
const int qg = (int)((gate_word >> (i * 4)) & 15u) - gate_zp;
const int qu = (int)((up_word >> (i * 4)) & 15u) - up_zp;
gate_tail += (float)qg * a * gate_scale;
up_tail += (float)qu * a * up_scale;
}
}
}
}
const float2 g04 = __half22float2(g0);
const float2 g15 = __half22float2(g1);
const float2 g26 = __half22float2(g2);
const float2 g37 = __half22float2(g3);
float gate_value = gate_tail + g04.x;
gate_value += g15.x;
gate_value += g26.x;
gate_value += g37.x;
gate_value += g04.y;
gate_value += g15.y;
gate_value += g26.y;
gate_value += g37.y;
gate_value = warp_sum(gate_value);
const float2 u04 = __half22float2(u0);
const float2 u15 = __half22float2(u1);
const float2 u26 = __half22float2(u2);
const float2 u37 = __half22float2(u3);
float up_value = up_tail + u04.x;
up_value += u15.x;
up_value += u26.x;
up_value += u37.x;
up_value += u04.y;
up_value += u15.y;
up_value += u26.y;
up_value += u37.y;
up_value = warp_sum(up_value);
if (lane == 0 && column < n) {
const float gate_h = __half2float(__float2half(gate_value));
const float up_h = __half2float(__float2half(up_value));
const float silu_h = Decomposed
? gate_up_decomposed_silu_f32(gate_h)
: gate_up_silu_f32(gate_h);
output[column] = __float2half_rn(__fmul_rn(silu_h, up_h));
}
}
extern "C" __global__ void matmul_nbits_gemv_f16_gate_up_swiglu_rmsnorm(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed_gate,
const __half* __restrict__ scales_gate,
const unsigned char* __restrict__ packed_up,
const __half* __restrict__ scales_up,
const unsigned char* __restrict__ zero_points_gate,
const unsigned char* __restrict__ zero_points_up,
const void* __restrict__ gamma,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int gamma_is_half,
const float epsilon)
{
matmul_nbits_gemv_f16_gate_up_swiglu_rmsnorm_tpl<false, false>(activation, packed_gate, scales_gate, packed_up, scales_up, zero_points_gate, zero_points_up, gamma, output, k, n, k_blocks, blob_size, zp_row_bytes, gamma_is_half, epsilon);
}
extern "C" __global__ void matmul_nbits_gemv_f16_gate_up_decomposed_swiglu_rmsnorm(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed_gate,
const __half* __restrict__ scales_gate,
const unsigned char* __restrict__ packed_up,
const __half* __restrict__ scales_up,
const unsigned char* __restrict__ zero_points_gate,
const unsigned char* __restrict__ zero_points_up,
const void* __restrict__ gamma,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int gamma_is_half,
const float epsilon)
{
matmul_nbits_gemv_f16_gate_up_swiglu_rmsnorm_tpl<false, true>(activation, packed_gate, scales_gate, packed_up, scales_up, zero_points_gate, zero_points_up, gamma, output, k, n, k_blocks, blob_size, zp_row_bytes, gamma_is_half, epsilon);
}
// Fused-symmetric (`ONNX_GENAI_GATEUP_VEC`) siblings of the two SYMMETRIC
// RMS-norm-fused gate/up SwiGLU entries — the dominant qwen2.5-14b decode
// kernel is `..decomposed_swiglu_rmsnorm`. Byte-identical (folds the `- 8`
// symmetric zero point into the dequant bias constants; see
// `int4x8_to_half2x4_sym8`), only the issued instruction count drops.
extern "C" __global__ void matmul_nbits_gemv_f16_gate_up_swiglu_rmsnorm_vec(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed_gate,
const __half* __restrict__ scales_gate,
const unsigned char* __restrict__ packed_up,
const __half* __restrict__ scales_up,
const unsigned char* __restrict__ zero_points_gate,
const unsigned char* __restrict__ zero_points_up,
const void* __restrict__ gamma,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int gamma_is_half,
const float epsilon)
{
matmul_nbits_gemv_f16_gate_up_swiglu_rmsnorm_tpl<false, false, true>(activation, packed_gate, scales_gate, packed_up, scales_up, zero_points_gate, zero_points_up, gamma, output, k, n, k_blocks, blob_size, zp_row_bytes, gamma_is_half, epsilon);
}
extern "C" __global__ void matmul_nbits_gemv_f16_gate_up_decomposed_swiglu_rmsnorm_vec(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed_gate,
const __half* __restrict__ scales_gate,
const unsigned char* __restrict__ packed_up,
const __half* __restrict__ scales_up,
const unsigned char* __restrict__ zero_points_gate,
const unsigned char* __restrict__ zero_points_up,
const void* __restrict__ gamma,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int gamma_is_half,
const float epsilon)
{
matmul_nbits_gemv_f16_gate_up_swiglu_rmsnorm_tpl<false, true, true>(activation, packed_gate, scales_gate, packed_up, scales_up, zero_points_gate, zero_points_up, gamma, output, k, n, k_blocks, blob_size, zp_row_bytes, gamma_is_half, epsilon);
}
// Occupancy-raised (`ONNX_GENAI_GATEUP_OCC`) siblings of the two SYMMETRIC
// RMS-norm-fused `_vec` entries. IDENTICAL kernel body and math — the only
// change is `__launch_bounds__(256, 8)`, which caps the register allocation at
// 32 regs/thread so the SM can co-resident 8 blocks (100% theoretical vs 75%
// register-limited). The dominant qwen2.5-14b decode kernel
// (`..decomposed_swiglu_rmsnorm`) is Short-Scoreboard/shared-load-latency bound
// (~51% of stall cycles waiting on the staged-activation LDS); the extra
// resident warps hide that latency. `__launch_bounds__` only constrains
// register allocation — same instruction stream, same fp16 accumulate order,
// same RMS reduction — so it is BYTE-IDENTICAL to the `_vec` entries above.
extern "C" __global__ void __launch_bounds__(256, 8)
matmul_nbits_gemv_f16_gate_up_swiglu_rmsnorm_vec_occ(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed_gate,
const __half* __restrict__ scales_gate,
const unsigned char* __restrict__ packed_up,
const __half* __restrict__ scales_up,
const unsigned char* __restrict__ zero_points_gate,
const unsigned char* __restrict__ zero_points_up,
const void* __restrict__ gamma,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int gamma_is_half,
const float epsilon)
{
matmul_nbits_gemv_f16_gate_up_swiglu_rmsnorm_tpl<false, false, true>(activation, packed_gate, scales_gate, packed_up, scales_up, zero_points_gate, zero_points_up, gamma, output, k, n, k_blocks, blob_size, zp_row_bytes, gamma_is_half, epsilon);
}
extern "C" __global__ void __launch_bounds__(256, 8)
matmul_nbits_gemv_f16_gate_up_decomposed_swiglu_rmsnorm_vec_occ(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed_gate,
const __half* __restrict__ scales_gate,
const unsigned char* __restrict__ packed_up,
const __half* __restrict__ scales_up,
const unsigned char* __restrict__ zero_points_gate,
const unsigned char* __restrict__ zero_points_up,
const void* __restrict__ gamma,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int gamma_is_half,
const float epsilon)
{
matmul_nbits_gemv_f16_gate_up_swiglu_rmsnorm_tpl<false, true, true>(activation, packed_gate, scales_gate, packed_up, scales_up, zero_points_gate, zero_points_up, gamma, output, k, n, k_blocks, blob_size, zp_row_bytes, gamma_is_half, epsilon);
}
extern "C" __global__ void matmul_nbits_gemv_f16_gate_up_swiglu_rmsnorm_zp(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed_gate,
const __half* __restrict__ scales_gate,
const unsigned char* __restrict__ packed_up,
const __half* __restrict__ scales_up,
const unsigned char* __restrict__ zero_points_gate,
const unsigned char* __restrict__ zero_points_up,
const void* __restrict__ gamma,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int gamma_is_half,
const float epsilon)
{
matmul_nbits_gemv_f16_gate_up_swiglu_rmsnorm_tpl<true, false>(activation, packed_gate, scales_gate, packed_up, scales_up, zero_points_gate, zero_points_up, gamma, output, k, n, k_blocks, blob_size, zp_row_bytes, gamma_is_half, epsilon);
}
extern "C" __global__ void matmul_nbits_gemv_f16_gate_up_decomposed_swiglu_rmsnorm_zp(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed_gate,
const __half* __restrict__ scales_gate,
const unsigned char* __restrict__ packed_up,
const __half* __restrict__ scales_up,
const unsigned char* __restrict__ zero_points_gate,
const unsigned char* __restrict__ zero_points_up,
const void* __restrict__ gamma,
__half* __restrict__ output,
const int k,
const int n,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int gamma_is_half,
const float epsilon)
{
matmul_nbits_gemv_f16_gate_up_swiglu_rmsnorm_tpl<true, true>(activation, packed_gate, scales_gate, packed_up, scales_up, zero_points_gate, zero_points_up, gamma, output, k, n, k_blocks, blob_size, zp_row_bytes, gamma_is_half, epsilon);
}
// Down projection specialization: a 256-thread CTA (8 warps) computes `COLS`
// columns and parallelizes over block-32 K tiles. Each thread loads its assigned
// activation block directly into registers and reuses it across all `COLS`
// columns, then the 8 warps combine their per-column partials through shared
// memory.
//
// `COLS` is a pure grid-fill knob: every output column is still reduced
// *entirely within one CTA* by all 256 threads striding the same K tiles in the
// same order, so the fp32 accumulation is bit-identical regardless of `COLS` —
// only the CTA count (grid = ceil(N / COLS)) changes. Tall-skinny down/output
// projections have a small N, so on many-SM devices the default 8-column launch
// underfills the machine (e.g. Qwen2.5-7B down: N=3584 -> 448 CTAs, ~0.57
// waves/SM on an H200). Halving `COLS` doubles the grid to fill the idle SMs on
// this latency-bound M=1 GEMV without changing the numerics; the host picks
// `COLS` from the device multiprocessor count (see `select_down_columns`).
template <int COLS>
__device__ __forceinline__ void matmul_nbits_gemv_f16_scales_f16_down_tpl(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int n,
const int k_blocks,
const int blob_size,
const int bias_post_round)
{
__shared__ float warp_sums[8][COLS];
const __half* __restrict__ scales =
reinterpret_cast<const __half*>(scales_raw);
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int column_base = (int)blockIdx.x * COLS;
float values[COLS];
#pragma unroll
for (int i = 0; i < COLS; ++i) {
values[i] = 0.0f;
}
for (int block = tid; block < k_blocks; block += (int)blockDim.x) {
const __half* activation_block = activation + block * 32;
const uint4 activation0 = permute_activation_f16x8(activation_block);
const uint4 activation1 = permute_activation_f16x8(activation_block + 8);
const uint4 activation2 = permute_activation_f16x8(activation_block + 16);
const uint4 activation3 = permute_activation_f16x8(activation_block + 24);
#pragma unroll
for (int tile_column = 0; tile_column < COLS; ++tile_column) {
const int column = column_base + tile_column;
if (column < n) {
const long packed_start =
((long)column * k_blocks + block) * blob_size;
const uint4 packed_weights =
*reinterpret_cast<const uint4*>(packed + packed_start);
const __half scale = scales[(long)column * k_blocks + block];
values[tile_column] += dot_int4x32_f16_permuted_scaled(
packed_weights,
activation0,
activation1,
activation2,
activation3,
scale);
}
}
}
#pragma unroll
for (int tile_column = 0; tile_column < COLS; ++tile_column) {
const float value = warp_sum(values[tile_column]);
if (lane == 0) {
warp_sums[warp][tile_column] = value;
}
}
__syncthreads();
if (warp == 0 && lane < COLS) {
const int column = column_base + lane;
float value = warp_sums[0][lane];
value += warp_sums[1][lane];
value += warp_sums[2][lane];
value += warp_sums[3][lane];
value += warp_sums[4][lane];
value += warp_sums[5][lane];
value += warp_sums[6][lane];
value += warp_sums[7][lane];
if (column < n) {
output[column] = fold_bias_f16(value, bias, column, bias_post_round);
}
}
}
// Default 8-column down projection (grid = ceil(N/8)).
extern "C" __global__ void matmul_nbits_gemv_f16_scales_f16_down(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round)
{
(void)k;
(void)block_size;
(void)zero_points;
(void)zp_row_bytes;
(void)scales_fp16;
matmul_nbits_gemv_f16_scales_f16_down_tpl<8>(
activation, packed, scales_raw, bias, output, n, k_blocks, blob_size,
bias_post_round);
}
// Grid-fill down projection variants: fewer columns per CTA -> proportionally
// larger grid, bit-identical output. Selected on grid-starved (small-N) down
// shapes to fill the multiprocessors on latency-bound decode.
extern "C" __global__ void matmul_nbits_gemv_f16_scales_f16_down_c4(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round)
{
(void)k;
(void)block_size;
(void)zero_points;
(void)zp_row_bytes;
(void)scales_fp16;
matmul_nbits_gemv_f16_scales_f16_down_tpl<4>(
activation, packed, scales_raw, bias, output, n, k_blocks, blob_size,
bias_post_round);
}
extern "C" __global__ void matmul_nbits_gemv_f16_scales_f16_down_c2(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round)
{
(void)k;
(void)block_size;
(void)zero_points;
(void)zp_row_bytes;
(void)scales_fp16;
matmul_nbits_gemv_f16_scales_f16_down_tpl<2>(
activation, packed, scales_raw, bias, output, n, k_blocks, blob_size,
bias_post_round);
}
// Model-agnostic fp16 int4/int8 decode GEMV supporting any power-of-two
// block_size. One warp per output column. Each lane owns contiguous 8-element K
// chunks and strides by 256 (= 32 lanes * 8) across the reduction. Unlike the
// tuned block-32 kernels, the scale / zero-point block index is derived from the
// real block_size (block = depth / block_size), so a lane's 8-element chunk
// always resolves to the block it belongs to for any block width that is a
// multiple of 8 (all supported power-of-two block sizes >= 16). The `bits`
// scalar selects the packed layout: int4 unpacks two nibbles per byte with an
// optional int4 zero point (default 8, packed two block-nibbles per byte); int8
// reads one unsigned byte per weight with an optional uint8 zero point (default
// 128, one per block) — byte-for-byte the tuned block-32 int8 layout generalized
// to any block width. fp32 accumulation is preserved; the kernel is
// register-only (capture-safe).
extern "C" __global__ void matmul_nbits_gemv_f16_general_bs(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round,
const int bits)
{
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int columns_per_block = (int)blockDim.x >> 5;
const int column = (int)blockIdx.x * columns_per_block + warp;
float value = 0.0f;
if (column < n) {
for (int depth = lane * 8; depth < k; depth += 256) {
const int block = depth / block_size;
const int within = depth - block * block_size;
const long blob_base = ((long)column * k_blocks + block) * blob_size;
float scale;
if (scales_fp16) {
scale = __half2float(
reinterpret_cast<const __half*>(scales)[(long)column * k_blocks + block]);
} else {
scale =
reinterpret_cast<const float*>(scales)[(long)column * k_blocks + block];
}
const int valid = min(8, k - depth);
float partial = 0.0f;
if (bits == 8) {
const int zero_point =
zero_points ? (int)zero_points[(long)column * k_blocks + block] : 128;
const unsigned char* block_bytes = packed + blob_base + within;
#pragma unroll
for (int i = 0; i < 8; ++i) {
if (i < valid) {
const int q = (int)block_bytes[i] - zero_point;
partial += (float)q * __half2float(activation[depth + i]);
}
}
} else {
int zero_point = 8;
if (zero_points) {
const unsigned char zp =
zero_points[(long)column * zp_row_bytes + (block >> 1)];
zero_point = (block & 1) ? (zp >> 4) : (zp & 15);
}
const unsigned int packed_word =
*reinterpret_cast<const unsigned int*>(packed + blob_base + (within >> 1));
if (valid == 8) {
// Fast path: LOP3 int4->fp16 dequant + 128-bit activation
// load. Byte-identical to the scalar loop below (same
// ascending-order fp32 products) but replaces the per-nibble
// shift/and/convert with 4 lop3 + f16x2 debias, cutting the
// dequant-ALU pressure that dominates the block!=32 GEMV.
const unsigned int sub2 = int4_zero_point_sub2(zero_point);
partial = dot_int4x8_f16_sub(packed_word, activation + depth, sub2);
} else {
#pragma unroll
for (int i = 0; i < 8; ++i) {
if (i < valid) {
const int q = (int)((packed_word >> (i * 4)) & 15u) - zero_point;
partial += (float)q * __half2float(activation[depth + i]);
}
}
}
}
value += partial * scale;
}
}
value = warp_sum(value);
if (lane == 0 && column < n) {
output[column] = fold_bias_f16(value, bias, column, bias_post_round);
}
}
// Split-K counterpart of `matmul_nbits_gemv_f16_general_bs`: K_SPLIT warps
// cooperate on one output column, each walking a strided subset of the 256-wide
// K steps, then summing their fp32 partials through shared memory. The launch
// grid is K_SPLIT x larger than the single-warp kernel, which fills the SMs on
// the grid-starved, latency-bound block!=32 decode GEMV (the medium
// projections run at ~0.5 waves/SM single-warp). Each warp keeps the same fp32
// LOP3 dequant loop as the single-warp kernel, so the only numeric difference
// is the new K-slice partial-sum association (near-equal, not byte-identical) —
// the same trade the block-32 split-K entries already ship by default.
extern "C" __global__ void matmul_nbits_gemv_f16_general_bs_splitk(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round,
const int bits)
{
constexpr int K_SPLIT = 4; // must match Rust GENERAL_BS_SPLITK
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int warps_per_block = (int)blockDim.x >> 5;
const int cols_per_block = warps_per_block / K_SPLIT;
const int col_local = warp / K_SPLIT;
const int ks = warp % K_SPLIT;
const int column = (int)blockIdx.x * cols_per_block + col_local;
__shared__ float partials[8][K_SPLIT];
float value = 0.0f;
if (column < n) {
for (int depth = ks * 256 + lane * 8; depth < k; depth += K_SPLIT * 256) {
const int block = depth / block_size;
const int within = depth - block * block_size;
const long blob_base = ((long)column * k_blocks + block) * blob_size;
float scale;
if (scales_fp16) {
scale = __half2float(
reinterpret_cast<const __half*>(scales)[(long)column * k_blocks + block]);
} else {
scale =
reinterpret_cast<const float*>(scales)[(long)column * k_blocks + block];
}
const int valid = min(8, k - depth);
float partial = 0.0f;
if (bits == 8) {
const int zero_point =
zero_points ? (int)zero_points[(long)column * k_blocks + block] : 128;
const unsigned char* block_bytes = packed + blob_base + within;
#pragma unroll
for (int i = 0; i < 8; ++i) {
if (i < valid) {
const int q = (int)block_bytes[i] - zero_point;
partial += (float)q * __half2float(activation[depth + i]);
}
}
} else {
int zero_point = 8;
if (zero_points) {
const unsigned char zp =
zero_points[(long)column * zp_row_bytes + (block >> 1)];
zero_point = (block & 1) ? (zp >> 4) : (zp & 15);
}
const unsigned int packed_word =
*reinterpret_cast<const unsigned int*>(packed + blob_base + (within >> 1));
if (valid == 8) {
const unsigned int sub2 = int4_zero_point_sub2(zero_point);
partial = dot_int4x8_f16_sub(packed_word, activation + depth, sub2);
} else {
#pragma unroll
for (int i = 0; i < 8; ++i) {
if (i < valid) {
const int q = (int)((packed_word >> (i * 4)) & 15u) - zero_point;
partial += (float)q * __half2float(activation[depth + i]);
}
}
}
}
value += partial * scale;
}
}
value = warp_sum(value);
if (lane == 0) {
partials[col_local][ks] = (column < n) ? value : 0.0f;
}
__syncthreads();
if (ks == 0 && lane == 0 && column < n) {
float acc = 0.0f;
#pragma unroll
for (int s = 0; s < K_SPLIT; ++s) {
acc += partials[col_local][s];
}
output[column] = fold_bias_f16(acc, bias, column, bias_post_round);
}
}
// Wide-load counterpart of `matmul_nbits_gemv_f16_general_bs` (see
// `gemv_int4_wide_lane_dot`). Same launch geometry (one warp per output column,
// 8 columns per 256-thread CTA) and same fp32/warp-sum reduction, but each lane
// streams 32 nibbles/step via a pipelined 128-bit weight load to lift DRAM
// throughput toward ORT's on the wide (already occupancy-filled) gate_up shape.
extern "C" __global__ void matmul_nbits_gemv_f16_general_bs_wide(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round,
const int bits)
{
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int columns_per_block = (int)blockDim.x >> 5;
const int column = (int)blockIdx.x * columns_per_block + warp;
float value = 0.0f;
if (column < n) {
value = gemv_int4_wide_lane_dot(
activation, packed, scales, zero_points, k, block_size, k_blocks,
blob_size, zp_row_bytes, scales_fp16, column, lane * 32, 32 * 32);
}
value = warp_sum(value);
if (lane == 0 && column < n) {
output[column] = fold_bias_f16(value, bias, column, bias_post_round);
}
}
// Wide-load counterpart of `matmul_nbits_gemv_f16_general_bs_splitk`: K_SPLIT
// warps cooperate on one output column, each walking a `32 * K_SPLIT`-strided
// set of 32-nibble chunks with the pipelined 128-bit loads, then summing their
// fp32 partials through shared memory (same grid-fill as the narrow split-K).
extern "C" __global__ void matmul_nbits_gemv_f16_general_bs_splitk_wide(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round,
const int bits)
{
constexpr int K_SPLIT = 4; // must match Rust GENERAL_BS_SPLITK
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int warps_per_block = (int)blockDim.x >> 5;
const int cols_per_block = warps_per_block / K_SPLIT;
const int col_local = warp / K_SPLIT;
const int ks = warp % K_SPLIT;
const int column = (int)blockIdx.x * cols_per_block + col_local;
__shared__ float partials[8][K_SPLIT];
float value = 0.0f;
if (column < n) {
value = gemv_int4_wide_lane_dot(
activation, packed, scales, zero_points, k, block_size, k_blocks,
blob_size, zp_row_bytes, scales_fp16, column, ks * 32 * 32 + lane * 32,
K_SPLIT * 32 * 32);
}
value = warp_sum(value);
if (lane == 0) {
partials[col_local][ks] = (column < n) ? value : 0.0f;
}
__syncthreads();
if (ks == 0 && lane == 0 && column < n) {
float acc = 0.0f;
#pragma unroll
for (int s = 0; s < K_SPLIT; ++s) {
acc += partials[col_local][s];
}
output[column] = fold_bias_f16(acc, bias, column, bias_post_round);
}
}
// Interleaved + biased (symmetric-only, OPT-IN) sibling of
// `matmul_nbits_gemv_f16_general_bs_splitk_wide`. Identical split-K geometry
// (K_SPLIT warps per column, shared-memory fp32 partial reduction), but consumes
// offline-interleaved weights and folds the symmetric -8 bias inside the LOP3
// converter. Because each lane's fp32 partial is bit-identical to the
// non-interleaved split-K wide kernel (see `gemv_int4_wide_lane_dot_interleaved`)
// and the K_SPLIT reduction order is unchanged, the output is byte-identical to
// `matmul_nbits_gemv_f16_general_bs_splitk_wide` on symmetric weights.
// `zero_points`/`zp_row_bytes`/`bits` are accepted for launch-signature parity
// but unused (the dispatch only routes symmetric int4 nodes here).
extern "C" __global__ void matmul_nbits_gemv_f16_general_bs_splitk_wide_interleaved(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round,
const int bits)
{
(void)zero_points;
(void)zp_row_bytes;
(void)bits;
constexpr int K_SPLIT = 4; // must match Rust GENERAL_BS_SPLITK
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int warps_per_block = (int)blockDim.x >> 5;
const int cols_per_block = warps_per_block / K_SPLIT;
const int col_local = warp / K_SPLIT;
const int ks = warp % K_SPLIT;
const int column = (int)blockIdx.x * cols_per_block + col_local;
__shared__ float partials[8][K_SPLIT];
float value = 0.0f;
if (column < n) {
value = gemv_int4_wide_lane_dot_interleaved(
activation, packed, scales, k, block_size, k_blocks,
blob_size, scales_fp16, column, ks * 32 * 32 + lane * 32,
K_SPLIT * 32 * 32);
}
value = warp_sum(value);
if (lane == 0) {
partials[col_local][ks] = (column < n) ? value : 0.0f;
}
__syncthreads();
if (ks == 0 && lane == 0 && column < n) {
float acc = 0.0f;
#pragma unroll
for (int s = 0; s < K_SPLIT; ++s) {
acc += partials[col_local][s];
}
output[column] = fold_bias_f16(acc, bias, column, bias_post_round);
}
}
// Multicol x split-K hybrid of `matmul_nbits_gemv_f16_general_bs_splitk_wide`.
// Combines the split-K grid-fill (K_SPLIT warps cooperate on one column group,
// summing their fp32 partials through shared memory) with the register-blocked
// wide-load multicol dot (each warp accumulates WIDE_NC output columns, issuing
// WIDE_NC independent 128-bit weight loads per chunk). The split-K path was
// picked for the grid-starved medium-N projections (down_proj N~4096, qkv/attn-
// out) *because* single-warp multicol under-fills the SMs there (~<1 wave); this
// hybrid restores the multicol memory-level parallelism (the lever that lifts the
// gate_up `wide_multicol` kernel to ~37% DRAM peak) WHILE keeping K_SPLIT so the
// grid still fills the device. Each lane's fp32 partial for (column, ks) uses the
// SAME depth0 (`ks*32*32 + lane*32`) and stride (`K_SPLIT*32*32`) as the
// single-column split-K wide kernel, the per-column accumulation order inside
// `gemv_int4_wide_lane_dot_multicol` is unchanged, and the K_SPLIT shared-memory
// reduction order is unchanged, so the output is BYTE-IDENTICAL to
// `matmul_nbits_gemv_f16_general_bs_splitk_wide`.
extern "C" __global__ void matmul_nbits_gemv_f16_general_bs_splitk_wide_multicol(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round,
const int bits)
{
constexpr int K_SPLIT = 4; // must match Rust GENERAL_BS_SPLITK_MULTICOL
// 256-thread CTA (8 warps) => MAX_COL_GROUPS column groups per block.
constexpr int MAX_COL_GROUPS = 8 / K_SPLIT;
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int warps_per_block = (int)blockDim.x >> 5;
const int col_groups_per_block = warps_per_block / K_SPLIT;
const int group_local = warp / K_SPLIT;
const int ks = warp % K_SPLIT;
const long col_base =
((long)blockIdx.x * col_groups_per_block + group_local) * (long)WIDE_NC;
__shared__ float partials[MAX_COL_GROUPS][WIDE_NC][K_SPLIT];
float values[WIDE_NC];
gemv_int4_wide_lane_dot_multicol(
activation, packed, scales, zero_points, k, block_size, k_blocks,
blob_size, zp_row_bytes, scales_fp16, col_base, n,
ks * 32 * 32 + lane * 32, K_SPLIT * 32 * 32, values);
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
const float reduced = warp_sum(values[c]);
if (lane == 0) {
partials[group_local][c][ks] = reduced;
}
}
__syncthreads();
if (ks == 0 && lane == 0) {
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
const long column = col_base + c;
if (column < n) {
float acc = 0.0f;
#pragma unroll
for (int s = 0; s < K_SPLIT; ++s) {
acc += partials[group_local][c][s];
}
output[column] = fold_bias_f16(acc, bias, column, bias_post_round);
}
}
}
}
// Interleaved + biased (symmetric-only, OPT-IN) sibling of
// `matmul_nbits_gemv_f16_general_bs_splitk_wide_multicol`. Same multicol x
// split-K geometry, but consumes offline nibble-interleaved weights and folds
// the symmetric -8 bias inside the LOP3 converter (dropping the per-block
// zero-point subtract and the `prmt.b32` activation reorder). Because each lane's
// fp32 partial is bit-identical to the non-interleaved multicol dot on symmetric
// weights and the K_SPLIT reduction order is unchanged, the output is
// byte-identical to `matmul_nbits_gemv_f16_general_bs_splitk_wide_multicol`.
// `zero_points`/`zp_row_bytes`/`bits` are accepted for launch-signature parity
// but unused (the dispatch only routes symmetric int4 nodes here).
extern "C" __global__ void matmul_nbits_gemv_f16_general_bs_splitk_wide_multicol_interleaved(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round,
const int bits)
{
(void)zero_points;
(void)zp_row_bytes;
(void)bits;
constexpr int K_SPLIT = 4; // must match Rust GENERAL_BS_SPLITK_MULTICOL
constexpr int MAX_COL_GROUPS = 8 / K_SPLIT;
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int warps_per_block = (int)blockDim.x >> 5;
const int col_groups_per_block = warps_per_block / K_SPLIT;
const int group_local = warp / K_SPLIT;
const int ks = warp % K_SPLIT;
const long col_base =
((long)blockIdx.x * col_groups_per_block + group_local) * (long)WIDE_NC;
__shared__ float partials[MAX_COL_GROUPS][WIDE_NC][K_SPLIT];
float values[WIDE_NC];
gemv_int4_wide_lane_dot_multicol_interleaved(
activation, packed, scales, k, block_size, k_blocks,
blob_size, scales_fp16, col_base, n,
ks * 32 * 32 + lane * 32, K_SPLIT * 32 * 32, values);
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
const float reduced = warp_sum(values[c]);
if (lane == 0) {
partials[group_local][c][ks] = reduced;
}
}
__syncthreads();
if (ks == 0 && lane == 0) {
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
const long column = col_base + c;
if (column < n) {
float acc = 0.0f;
#pragma unroll
for (int s = 0; s < K_SPLIT; ++s) {
acc += partials[group_local][c][s];
}
output[column] = fold_bias_f16(acc, bias, column, bias_post_round);
}
}
}
}
// Column register-blocked wide-load GEMV (see `gemv_int4_wide_lane_dot_multicol`).
// Same launch geometry as `matmul_nbits_gemv_f16_general_bs_wide` (256-thread
// CTA, one warp per group), but every warp emits WIDE_NC output columns, so the
// grid covers `columns_per_block = 8 * WIDE_NC` columns per block. Decoding each
// activation sub-word once and reusing it across WIDE_NC columns relieves the
// L1/TEX-throughput limiter of the single-column wide kernel while staying
// byte-identical (per-column fp32 accumulation order unchanged).
extern "C" __global__ void matmul_nbits_gemv_f16_general_bs_wide_multicol(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round,
const int bits)
{
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int warps_per_block = (int)blockDim.x >> 5;
const long col_base =
((long)blockIdx.x * warps_per_block + warp) * (long)WIDE_NC;
float values[WIDE_NC];
gemv_int4_wide_lane_dot_multicol(
activation, packed, scales, zero_points, k, block_size, k_blocks,
blob_size, zp_row_bytes, scales_fp16, col_base, n, lane * 32, 32 * 32,
values);
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
const float reduced = warp_sum(values[c]);
const long column = col_base + c;
if (lane == 0 && column < n) {
output[column] = fold_bias_f16(reduced, bias, column, bias_post_round);
}
}
}
// Interleaved + biased (symmetric-only, OPT-IN) sibling of
// `matmul_nbits_gemv_f16_general_bs_wide_multicol`. Identical launch geometry
// and grid; consumes offline-interleaved weights and folds the symmetric -8
// bias inside the LOP3 converter, dropping the per-block zero-point subtract and
// the `prmt.b32` activation reorder. Byte-identical output to the fp32 multicol
// kernel on symmetric weights (see `gemv_int4_wide_lane_dot_multicol_interleaved`).
// `zero_points`/`zp_row_bytes` are accepted for launch-signature parity but
// unused (the dispatch only routes symmetric nodes here).
extern "C" __global__ void matmul_nbits_gemv_f16_general_bs_wide_multicol_interleaved(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round,
const int bits)
{
(void)zero_points;
(void)zp_row_bytes;
(void)bits;
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int warps_per_block = (int)blockDim.x >> 5;
const long col_base =
((long)blockIdx.x * warps_per_block + warp) * (long)WIDE_NC;
float values[WIDE_NC];
gemv_int4_wide_lane_dot_multicol_interleaved(
activation, packed, scales, k, block_size, k_blocks,
blob_size, scales_fp16, col_base, n, lane * 32, 32 * 32,
values);
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
const float reduced = warp_sum(values[c]);
const long column = col_base + c;
if (lane == 0 && column < n) {
output[column] = fold_bias_f16(reduced, bias, column, bias_post_round);
}
}
}
// Offline (once-per-weight) int4 nibble-interleave pass for the
// `ONNX_GENAI_INTERLEAVE_DEQUANT` lever. Reads the packed weight buffer as
// 32-bit words and rewrites each word from natural nibble order
// [e7 e6 e5 e4 | e3 e2 e1 e0] to TRT-LLM order [e7 e5 e3 e1 | e6 e4 e2 e0]
// (even elements to the low four nibble slots, odd to the high four). This is a
// pure per-word permutation, independent of block layout, so it applies to any
// int4 MatMulNBits weight whose byte count is a multiple of 4. Run once into a
// cached device buffer; the decode GEMV then reads the interleaved buffer.
extern "C" __global__ void matmul_nbits_interleave_int4(
const unsigned int* __restrict__ src,
unsigned int* __restrict__ dst,
const unsigned long words)
{
const unsigned long idx =
(unsigned long)blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= words) {
return;
}
const unsigned int w = src[idx];
const unsigned int e0 = (w >> 0) & 0xfu;
const unsigned int e1 = (w >> 4) & 0xfu;
const unsigned int e2 = (w >> 8) & 0xfu;
const unsigned int e3 = (w >> 12) & 0xfu;
const unsigned int e4 = (w >> 16) & 0xfu;
const unsigned int e5 = (w >> 20) & 0xfu;
const unsigned int e6 = (w >> 24) & 0xfu;
const unsigned int e7 = (w >> 28) & 0xfu;
// Physical slots [n0..n7] = [e0,e2,e4,e6,e1,e3,e5,e7].
dst[idx] = (e0 << 0) | (e2 << 4) | (e4 << 8) | (e6 << 12)
| (e1 << 16) | (e3 << 20) | (e5 << 24) | (e7 << 28);
}
// fp16 mixed-precision sibling of `matmul_nbits_gemv_f16_general_bs_wide_multicol`
// (see `gemv_int4_fp16_lane_dot_multicol`). Same launch geometry and grid; the
// only difference is the per-chunk fp16 __hfma2 MAC. Opt-in (gated by
// `use_gemv_fp16`) and accuracy-gated (NOT byte-identical to the fp32 path).
extern "C" __global__ void matmul_nbits_gemv_f16_general_bs_wide_multicol_fp16(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round,
const int bits)
{
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int warps_per_block = (int)blockDim.x >> 5;
const long col_base =
((long)blockIdx.x * warps_per_block + warp) * (long)WIDE_NC;
float values[WIDE_NC];
gemv_int4_fp16_lane_dot_multicol(
activation, packed, scales, zero_points, k, block_size, k_blocks,
blob_size, zp_row_bytes, scales_fp16, col_base, n, lane * 32, 32 * 32,
values);
#pragma unroll
for (int c = 0; c < WIDE_NC; ++c) {
const float reduced = warp_sum(values[c]);
const long column = col_base + c;
if (lane == 0 && column < n) {
output[column] = fold_bias_f16(reduced, bias, column, bias_post_round);
}
}
}
// Model-agnostic fp16 int4/int8 prefill GEMM supporting any power-of-two
// block_size. Identical 16x16 tiling and fp32 accumulation as the tuned
// block-32 GEMM, but the reduction walks K in fixed 32-wide tiles and derives
// the block index from the real block_size (block = depth / block_size), so the
// K-tile width is decoupled from the block width. For block_size == 32 this is
// numerically identical to matmul_nbits_gemm_f16 (block == tile).
extern "C" __global__ void matmul_nbits_gemm_f16_general_bs(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int m,
const int k,
const int n,
const int k_blocks,
const int bits,
const int scales_fp16,
const int bias_post_round,
const int bias_row_stride,
const int block_size,
const int blob_size)
{
__shared__ float activation_tile[16][32];
__shared__ float weight_tile[32][16];
const int tid = (int)threadIdx.y * 16 + (int)threadIdx.x;
const int row = (int)blockIdx.y * 16 + (int)threadIdx.y;
const int column = (int)blockIdx.x * 16 + (int)threadIdx.x;
float value = 0.0f;
const int k_tiles = (k + 31) / 32;
for (int tile = 0; tile < k_tiles; ++tile) {
#pragma unroll
for (int load = tid; load < 16 * 32; load += 16 * 16) {
const int tile_row = load >> 5;
const int within = load & 31;
const int depth = tile * 32 + within;
const int global_row = (int)blockIdx.y * 16 + tile_row;
activation_tile[tile_row][within] =
global_row < m && depth < k
? __half2float(activation[(long)global_row * k + depth])
: 0.0f;
}
#pragma unroll
for (int load = tid; load < 32 * 16; load += 16 * 16) {
const int tile_column = load >> 5;
const int within = load & 31;
const int depth = tile * 32 + within;
const int global_column = (int)blockIdx.x * 16 + tile_column;
float weight = 0.0f;
if (global_column < n && depth < k) {
const int block = depth / block_size;
const int within_block = depth - block * block_size;
const long scale_index = (long)global_column * k_blocks + block;
const float scale = scales_fp16
? __half2float(
reinterpret_cast<const __half*>(scales_raw)[scale_index])
: reinterpret_cast<const float*>(scales_raw)[scale_index];
const long blob_base =
((long)global_column * k_blocks + block) * blob_size;
int quantized;
int zero_point;
if (bits == 8) {
quantized = (int)packed[blob_base + within_block];
zero_point = zero_points ? (int)zero_points[scale_index] : 128;
} else {
const unsigned char byte =
packed[blob_base + (within_block >> 1)];
quantized = (within_block & 1) ? (byte >> 4) : (byte & 15);
zero_point = 8;
if (zero_points) {
const int zp_row_bytes = (k_blocks + 1) >> 1;
const unsigned char zp =
zero_points[(long)global_column * zp_row_bytes + (block >> 1)];
zero_point = (block & 1) ? (zp >> 4) : (zp & 15);
}
}
weight = ((float)quantized - (float)zero_point) * scale;
}
weight_tile[within][tile_column] = weight;
}
__syncthreads();
if (row < m && column < n) {
#pragma unroll
for (int within = 0; within < 32; ++within) {
value += activation_tile[threadIdx.y][within]
* weight_tile[within][threadIdx.x];
}
}
__syncthreads();
}
if (row < m && column < n) {
const __half* row_bias = bias ? bias + (long)row * bias_row_stride : bias;
output[(long)row * n + column] =
fold_bias_f16(value, row_bias, column, bias_post_round);
}
}
"#;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum F16GemvVariant {
General,
DownProjection,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct F16GemvSelection {
variant: F16GemvVariant,
reason: &'static str,
}
fn select_f16_gemv_variant(
k: usize,
n: usize,
block_size: usize,
scales_fp16: bool,
has_zero_points: bool,
) -> F16GemvSelection {
let down_eligible = !has_zero_points
&& scales_fp16
&& block_size == GEMV_F16_DOWN_BLOCK_SIZE
&& k.is_multiple_of(GEMV_F16_DOWN_BLOCK_SIZE)
&& k > n;
if down_eligible {
F16GemvSelection {
variant: F16GemvVariant::DownProjection,
reason: "variant=down_projection;class=tall_skinny(K>N);block_size=32;\
scales=fp16;K%32==0",
}
} else {
F16GemvSelection {
variant: F16GemvVariant::General,
reason: if has_zero_points {
"variant=general;zero_points=explicit;down_projection requires symmetric zp=8"
} else {
"variant=general;class=not(tall_skinny K>N & block_size=32 & \
scales=fp16 & K%32==0)"
},
}
}
}
const DOWN_FILL_CTAS_PER_SM: usize = 12;
fn select_down_columns(n: usize, multiprocessor_count: u32) -> (usize, &'static str) {
let target = (multiprocessor_count.max(1) as usize).saturating_mul(DOWN_FILL_CTAS_PER_SM);
for (cols, entry) in [
(GEMV_F16_DOWN_COLUMNS_PER_BLOCK, GEMV_F16_DOWN_ENTRY),
(4usize, GEMV_F16_DOWN_C4_ENTRY),
] {
if n.div_ceil(cols) >= target {
return (cols, entry);
}
}
(2usize, GEMV_F16_DOWN_C2_ENTRY)
}
const ACCURACY4_GEMV_FILL_CTAS_PER_SM: usize = 12;
const F16_SYMMETRIC_SPLITK_TARGET_WARPS_PER_SM: usize = 16;
const GEMV_F16_PIPE_WELL_OCCUPIED_CTAS_PER_SM: usize = 32;
const GENERAL_BS_SPLITK_TARGET_CTAS_PER_SM: usize = 16;
fn use_general_bs_splitk(k: usize, n: usize, bits: usize, multiprocessor_count: u32) -> bool {
if let Some(forced) = general_bs_splitk_override() {
return forced;
}
let single_warp_ctas = n.div_ceil((GEMV_F16_LARGE_THREADS / 32) as usize);
bits == 4
&& k >= 512
&& !(n <= GEMV_F16_SMALL_N_MAX && k <= GEMV_F16_SMALL_N_MAX)
&& single_warp_ctas
< (multiprocessor_count.max(1) as usize)
.saturating_mul(GENERAL_BS_SPLITK_TARGET_CTAS_PER_SM)
}
fn general_bs_splitk_override() -> Option<bool> {
match std::env::var("ONNX_GENAI_GENERAL_SPLITK").ok().as_deref() {
Some("1") | Some("true") | Some("on") => Some(true),
Some("0") | Some("false") | Some("off") => Some(false),
_ => None,
}
}
fn use_gemv_wideload(bits: usize, block_size: usize, k: usize) -> bool {
if bits != 4 || !block_size.is_multiple_of(32) || !k.is_multiple_of(32) {
return false;
}
!matches!(
std::env::var("ONNX_GENAI_GEMV_WIDELOAD").ok().as_deref(),
Some("0") | Some("false") | Some("off")
)
}
fn use_gemv_wide_multicol() -> bool {
!matches!(
std::env::var("ONNX_GENAI_GEMV_WIDE_MULTICOL")
.ok()
.as_deref(),
Some("0") | Some("false") | Some("off")
)
}
fn use_gemv_splitk_multicol() -> bool {
!matches!(
std::env::var("ONNX_GENAI_GEMV_SPLITK_MULTICOL")
.ok()
.as_deref(),
Some("0") | Some("false") | Some("off")
)
}
fn splitk_smalln_prefers_single_column(n: usize, multiprocessor_count: u32) -> bool {
match std::env::var("ONNX_GENAI_GEMV_SPLITK_SMALLN_SINGLECOL")
.ok()
.as_deref()
{
Some("1") | Some("true") | Some("on") => return true,
Some("0") | Some("false") | Some("off") => return false,
_ => {}
}
let multicol_cols_per_cta = (8 / GENERAL_BS_SPLITK_MULTICOL) * GEMV_F16_WIDE_MULTICOL_NC;
let multicol_grid = n.div_ceil(multicol_cols_per_cta.max(1));
multicol_grid < multiprocessor_count.max(1) as usize
}
fn use_gemv_fp16() -> bool {
matches!(
std::env::var("ONNX_GENAI_GEMV_FP16").ok().as_deref(),
Some("1") | Some("true") | Some("on")
)
}
fn dequant_f16_gemm_enabled() -> bool {
!matches!(
std::env::var("ONNX_GENAI_DEQUANT_F16_GEMM").ok().as_deref(),
Some("0") | Some("false") | Some("off")
)
}
fn dequant_f16_gemm_min_m() -> usize {
std::env::var("ONNX_GENAI_DEQUANT_F16_GEMM_MIN_M")
.ok()
.and_then(|value| value.parse::<usize>().ok())
.filter(|min_m| *min_m > 1)
.unwrap_or(8)
}
fn dequant_f16_gemm_max_scratch_bytes() -> usize {
std::env::var("ONNX_GENAI_DEQUANT_F16_GEMM_MAX_SCRATCH")
.ok()
.and_then(|value| value.parse::<usize>().ok())
.unwrap_or(1 << 30)
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct Bf16DirectOut {
staging: CUdeviceptr,
dst: CUdeviceptr,
elements: usize,
taken: bool,
}
thread_local! {
static BF16_DIRECT_OUT: std::cell::Cell<Option<Bf16DirectOut>> =
const { std::cell::Cell::new(None) };
}
thread_local! {
static BF16_DIRECT_OUT_STORES: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
}
fn take_bf16_direct_out(staging: CUdeviceptr, elements: usize) -> Option<CUdeviceptr> {
BF16_DIRECT_OUT.with(|cell| match cell.get() {
Some(mut offer)
if !offer.taken && offer.staging == staging && offer.elements == elements =>
{
offer.taken = true;
cell.set(Some(offer));
BF16_DIRECT_OUT_STORES.with(|count| count.set(count.get() + 1));
Some(offer.dst)
}
_ => None,
})
}
fn interleave_dequant_enabled() -> bool {
matches!(
std::env::var("ONNX_GENAI_INTERLEAVE_DEQUANT")
.ok()
.as_deref(),
Some("1") | Some("true") | Some("on")
)
}
fn use_scales_f16_zp_splitk_pf() -> bool {
!matches!(
std::env::var("ONNX_GENAI_ZP_SPLITK_PREFETCH")
.ok()
.as_deref(),
Some("0") | Some("false") | Some("off")
)
}
impl crate::interleave_cache::InterleaveDevice for CudaRuntime {
fn interleave_device_id(&self) -> u64 {
self.runtime_id()
}
fn interleave_alloc(&self, bytes: usize) -> Result<CUdeviceptr> {
self.alloc_raw(bytes)
}
unsafe fn interleave_free(&self, ptr: CUdeviceptr) {
let _ = unsafe { self.free_raw(ptr) };
}
fn interleave_build(&self, src: CUdeviceptr, dst: CUdeviceptr, bytes: usize) -> Result<()> {
launch_interleave_int4(self, src, dst, bytes)
}
fn interleave_is_capturing(&self) -> Result<bool> {
self.is_capturing()
}
fn interleave_frees_are_observed(&self) -> bool {
!self.weights_may_be_paged()
}
fn interleave_drain_before_free(&self) -> Result<()> {
self.drain_for_unmap()
}
}
fn launch_interleave_int4(
runtime: &CudaRuntime,
src: CUdeviceptr,
dst: CUdeviceptr,
bytes: usize,
) -> Result<()> {
let words = (bytes / 4) as u64;
let function = runtime.nvrtc_function(GEMV_F16_MODULE, GEMV_F16_SRC, INTERLEAVE_INT4_ENTRY)?;
let src_ptr = cuptr(src as usize as *const c_void);
let dst_ptr = cuptr(dst as usize as *const c_void);
const THREADS: u32 = 256;
let grid = u32::try_from(words.div_ceil(THREADS as u64))
.unwrap_or(u32::MAX)
.max(1);
let mut builder = runtime.stream().launch_builder(&function);
builder.arg(&src_ptr).arg(&dst_ptr).arg(&words);
unsafe {
builder.launch(LaunchConfig {
grid_dim: (grid, 1, 1),
block_dim: (THREADS, 1, 1),
shared_mem_bytes: 0,
})
}
.map(|_| ())
.map_err(|err| driver_err("launch MatMulNBits int4 interleave", err))
}
fn ensure_interleaved(
runtime: &Arc<CudaRuntime>,
packed: CUdeviceptr,
bytes: usize,
) -> Result<(CUdeviceptr, bool)> {
runtime.ensure_interleaved_int4(packed, bytes)
}
fn use_f16_symmetric_splitk(
k: usize,
n: usize,
multiprocessor_count: u32,
max_threads_per_block: u32,
) -> bool {
let eligible =
k >= 512 && k.is_multiple_of(32) && max_threads_per_block >= GEMV_F16_LARGE_THREADS;
if !eligible {
return false;
}
n < (multiprocessor_count.max(1) as usize)
.saturating_mul(F16_SYMMETRIC_SPLITK_TARGET_WARPS_PER_SM)
}
fn use_scales_f16_zp_splitk_gate(k: usize, n: usize, multiprocessor_count: u32) -> bool {
match std::env::var("ONNX_GENAI_SCALES_F16_ZP_SPLITK")
.ok()
.as_deref()
{
Some("1") | Some("true") | Some("on") => return true,
Some("0") | Some("false") | Some("off") => return false,
_ => {}
}
let grid_starved = n
< (multiprocessor_count.max(1) as usize)
.saturating_mul(F16_SYMMETRIC_SPLITK_TARGET_WARPS_PER_SM);
grid_starved || k >= ZP_SPLITK_BANDWIDTH_MIN_K
}
fn scales_f16_zp_split_factor(n: usize, multiprocessor_count: u32) -> usize {
let warps_per_cta = (GEMV_F16_LARGE_THREADS / 32) as usize;
let blocks_at_default = n.div_ceil(warps_per_cta / GEMV_F16_SCALES_F16_ZP_SPLITK);
if blocks_at_default < multiprocessor_count.max(1) as usize {
GEMV_F16_SCALES_F16_ZP_SPLITK8
} else {
GEMV_F16_SCALES_F16_ZP_SPLITK
}
}
const ZP_SPLITK_BANDWIDTH_MIN_K: usize = 6656;
fn select_accuracy4_gemv_warps(n: usize, multiprocessor_count: u32) -> u32 {
if let Some(forced) = std::env::var("ONNX_GENAI_ACC4_WARPS")
.ok()
.and_then(|value| value.parse::<u32>().ok())
.filter(|warps| matches!(warps, 1 | 2 | 4 | 8))
{
return forced;
}
let target =
(multiprocessor_count.max(1) as usize).saturating_mul(ACCURACY4_GEMV_FILL_CTAS_PER_SM);
for warps in [8usize, 4, 2] {
if n.div_ceil(warps) >= target {
return warps as u32;
}
}
1
}
fn use_accuracy4_stage64(
n: usize,
multiprocessor_count: u32,
compute_capability: (u32, u32),
max_shared_memory_per_block: u32,
) -> bool {
if max_shared_memory_per_block < GEMV_ACCURACY4_STAGE64_SHARED_BYTES {
return false;
}
let resident_warps = crate::arch::decode_resident_warps_per_sm(compute_capability) as usize;
let resident_ctas = resident_warps / (GEMV_ACCURACY4_THREADS as usize / 32);
let one_wave = (multiprocessor_count.max(1) as usize).saturating_mul(resident_ctas);
n.div_ceil(GEMV_ACCURACY4_COLUMNS_PER_BLOCK) < one_wave
}
fn down_columns_override() -> Option<(usize, &'static str)> {
match std::env::var("ONNX_GENAI_DOWN_COLS").ok()?.as_str() {
"8" => Some((GEMV_F16_DOWN_COLUMNS_PER_BLOCK, GEMV_F16_DOWN_ENTRY)),
"4" => Some((4, GEMV_F16_DOWN_C4_ENTRY)),
"2" => Some((2, GEMV_F16_DOWN_C2_ENTRY)),
_ => None,
}
}
fn scales_f16_pipeline_enabled() -> bool {
!matches!(
std::env::var("ONNX_GENAI_GEMV_PIPELINE").ok().as_deref(),
Some("0") | Some("false") | Some("off")
)
}
fn int8_symmetric_splitk_enabled() -> bool {
!std::env::var_os("ONNX_GENAI_CUDA_DISABLE_INT8_SYMMETRIC_SPLITK")
.is_some_and(|value| value != "0" && !value.is_empty())
}
fn scales_f16_pipe_well_occupied(
n: usize,
columns_per_block: usize,
multiprocessor_count: u32,
) -> bool {
match std::env::var("ONNX_GENAI_GEMV_PIPE_WELLOCC")
.ok()
.as_deref()
{
Some("0") | Some("false") | Some("off") => return false,
Some("1") | Some("true") | Some("on") => return true,
_ => {}
}
let ctas = n.div_ceil(columns_per_block.max(1));
let threshold = (multiprocessor_count.max(1) as usize)
.saturating_mul(GEMV_F16_PIPE_WELL_OCCUPIED_CTAS_PER_SM);
ctas >= threshold
}
fn gate_up_vec_enabled() -> bool {
!matches!(
std::env::var("ONNX_GENAI_GATEUP_VEC").ok().as_deref(),
Some("0") | Some("false") | Some("off")
)
}
fn gate_up_occ_enabled() -> bool {
!matches!(
std::env::var("ONNX_GENAI_GATEUP_OCC").ok().as_deref(),
Some("0") | Some("false") | Some("off")
)
}
pub struct MatMulNBitsFactory {
pub runtime: Arc<CudaRuntime>,
}
impl KernelFactory for MatMulNBitsFactory {
fn create(&self, node: &Node, _input_shapes: &[Vec<usize>]) -> Result<Box<dyn Kernel>> {
let k = required_positive_attr(node, "K")?;
let n = required_positive_attr(node, "N")?;
let bits = optional_int_attr(node, "bits")?.unwrap_or(4);
if !matches!(bits, 4 | 8) {
return Err(error(format!(
"MatMulNBits CUDA supports bits in {{4, 8}}, got bits={bits}. Why: the native \
kernels implement packed int4 and int8 layouts. How to fix: export bits=4 or \
bits=8, or select another execution provider"
)));
}
let weight_prepacked = optional_int_attr(node, "weight_prepacked")?.unwrap_or(0);
if weight_prepacked != 0 {
return Err(error(format!(
"weight_prepacked={weight_prepacked} is unsupported: CUDA only supports the standard (non-prepacked) layout"
)));
}
let block_size = required_positive_attr(node, "block_size")?;
if block_size < 16 || !block_size.is_power_of_two() {
return Err(error(format!(
"block_size must be a power of two and at least 16, got {block_size}"
)));
}
let accuracy_level = node
.attr("accuracy_level")
.and_then(|value| value.as_int())
.unwrap_or(0);
let accuracy4_workspace = if bits == 4 && accuracy_level == 4 {
Some(Mutex::new(Accuracy4Workspace::new(
self.runtime.clone(),
k,
block_size,
)?))
} else {
None
};
Ok(Box::new(MatMulNBitsKernel {
runtime: self.runtime.clone(),
k,
n,
bits: bits as usize,
block_size,
accuracy_level,
accuracy4_workspace,
marlin_repack_cache: marlin_gemm::RepackCache::new(self.runtime.clone()),
constant_inputs: [false; 8],
fold_bias_post_round: node
.attr(crate::optimizer::MATMUL_NBITS_FOLDED_BIAS_ATTR)
.and_then(onnx_runtime_ir::Attribute::as_int)
== Some(1),
gate_up_swiglu: node
.attr(crate::optimizer::GATE_UP_SWIGLU_FUSION_ATTR)
.and_then(onnx_runtime_ir::Attribute::as_int)
== Some(1),
decomposed_silu: node
.attr(crate::optimizer::DECOMPOSED_SILU_ATTR)
.and_then(onnx_runtime_ir::Attribute::as_int)
== Some(1),
rmsnorm_prologue: node
.attr(crate::optimizer::MATMUL_NBITS_RMSNORM_PROLOGUE_ATTR)
.and_then(onnx_runtime_ir::Attribute::as_int)
== Some(1),
rmsnorm_epsilon: node
.attr(crate::optimizer::MATMUL_NBITS_RMSNORM_EPSILON_ATTR)
.and_then(onnx_runtime_ir::Attribute::as_float)
.unwrap_or(1e-5),
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(self.runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(self.runtime.clone())),
}))
}
}
#[derive(Debug)]
struct Accuracy4Workspace {
allocation: Arc<GraphDeviceAllocation>,
quantized_activation: CUdeviceptr,
activation_scale: CUdeviceptr,
padded_k: usize,
}
impl Accuracy4Workspace {
fn new(runtime: Arc<CudaRuntime>, k: usize, block_size: usize) -> Result<Self> {
let padded_k = k.div_ceil(block_size) * block_size;
let k_blocks = padded_k / block_size;
let scale_bytes = k_blocks * std::mem::size_of::<f32>();
let allocation = GraphDeviceAllocation::allocate(&runtime, padded_k + scale_bytes)?;
let quantized_activation = allocation.ptr();
Ok(Self {
allocation,
quantized_activation,
activation_scale: quantized_activation + padded_k as CUdeviceptr,
padded_k,
})
}
}
impl Accuracy4Workspace {
fn device_graph_resource(&self) -> DeviceGraphResource {
GraphDeviceAllocation::device_graph_resource(&self.allocation)
}
}
#[derive(Debug)]
struct Bf16Scratch {
runtime: Arc<CudaRuntime>,
allocation: Option<Arc<GraphDeviceAllocation>>,
cap: usize,
used: bool,
}
impl Bf16Scratch {
fn new(runtime: Arc<CudaRuntime>) -> Self {
Self {
runtime,
allocation: None,
cap: 0,
used: false,
}
}
fn begin_call(&mut self) {
self.used = false;
}
fn ensure(&mut self, bytes: usize) -> Result<CUdeviceptr> {
self.used = true;
if bytes > self.cap {
if self.allocation.is_some() {
self.runtime.drain_for_unmap()?;
}
self.allocation = Some(GraphDeviceAllocation::allocate(
&self.runtime,
bytes.max(1),
)?);
self.cap = bytes;
}
Ok(self
.allocation
.as_ref()
.map_or(0, |allocation| allocation.ptr()))
}
fn device_graph_resource(&self) -> Option<DeviceGraphResource> {
if !self.used {
return None;
}
self.allocation
.as_ref()
.map(GraphDeviceAllocation::device_graph_resource)
}
}
#[derive(Debug)]
struct Bf16ConstCache {
runtime: Arc<CudaRuntime>,
allocation: Option<Arc<GraphDeviceAllocation>>,
cap: usize,
used: bool,
slots: Vec<(CUdeviceptr, usize, usize)>,
}
impl Bf16ConstCache {
fn new(runtime: Arc<CudaRuntime>) -> Self {
Self {
runtime,
allocation: None,
cap: 0,
used: false,
slots: Vec::new(),
}
}
fn begin_call(&mut self) {
self.used = false;
}
fn staged(&mut self, consts: &[(CUdeviceptr, usize)]) -> Result<Vec<CUdeviceptr>> {
self.used = true;
let matches =
self.slots.len() == consts.len()
&& self.slots.iter().zip(consts).all(
|((src, numel, _), (want_src, want_numel))| {
*src == *want_src && *numel == *want_numel
},
);
if !matches {
self.rebuild(consts)?;
}
Ok(self
.slots
.iter()
.map(|(_, _, offset)| {
self.allocation
.as_ref()
.map_or(0, |allocation| allocation.ptr())
+ *offset as CUdeviceptr
})
.collect())
}
fn rebuild(&mut self, consts: &[(CUdeviceptr, usize)]) -> Result<()> {
const ALIGN: usize = 256;
let f16 = std::mem::size_of::<half::f16>();
let round = |bytes: usize| bytes.div_ceil(ALIGN) * ALIGN;
let mut slots = Vec::with_capacity(consts.len());
let mut total = 0usize;
for (src, numel) in consts {
slots.push((*src, *numel, total));
total += round(numel * f16);
}
if total > self.cap {
if self.allocation.is_some() {
self.runtime.drain_for_unmap()?;
}
self.allocation = Some(GraphDeviceAllocation::allocate(
&self.runtime,
total.max(1),
)?);
self.cap = total;
}
let base = self
.allocation
.as_ref()
.map_or(0, |allocation| allocation.ptr());
for (src, numel, offset) in &slots {
super::cast::launch_cast_raw(
&self.runtime,
cuptr(*src as *const c_void),
DataType::BFloat16,
base + *offset as CUdeviceptr,
DataType::Float16,
*numel,
)?;
}
self.slots = slots;
Ok(())
}
fn device_graph_resource(&self) -> Option<DeviceGraphResource> {
if !self.used {
return None;
}
self.allocation
.as_ref()
.map(GraphDeviceAllocation::device_graph_resource)
}
}
#[derive(Debug)]
pub struct MatMulNBitsKernel {
runtime: Arc<CudaRuntime>,
k: usize,
n: usize,
bits: usize,
block_size: usize,
accuracy_level: i64,
accuracy4_workspace: Option<Mutex<Accuracy4Workspace>>,
marlin_repack_cache: marlin_gemm::RepackCache,
constant_inputs: [bool; 8],
fold_bias_post_round: bool,
gate_up_swiglu: bool,
decomposed_silu: bool,
rmsnorm_prologue: bool,
rmsnorm_epsilon: f32,
last_call_capture_safe: AtomicBool,
bf16_scratch: Mutex<Bf16Scratch>,
bf16_const_cache: Mutex<Bf16ConstCache>,
}
impl MatMulNBitsKernel {
fn marlin_weight_inputs_constant(&self) -> bool {
marlin_weight_inputs_are_constant(&self.constant_inputs, self.gate_up_swiglu)
}
fn release_marlin_repack_for_decode(&self) {
if self.marlin_repack_cache.retained_bytes() > 0
&& !self.runtime.is_capturing().unwrap_or(true)
{
self.marlin_repack_cache.release_all();
}
}
fn run(
&self,
inputs: &[TensorView],
outputs: &mut [TensorMut],
workspace: Option<WorkspaceView>,
) -> Result<()> {
self.last_call_capture_safe.store(false, Ordering::Relaxed);
self.marlin_repack_cache.begin_call();
self.bf16_scratch
.lock()
.map_err(|_| error("BFloat16 scratch lock poisoned"))?
.begin_call();
self.bf16_const_cache
.lock()
.map_err(|_| error("BFloat16 constant cache lock poisoned"))?
.begin_call();
let max_inputs = if self.gate_up_swiglu { 8 } else { 7 };
if !(3..=max_inputs).contains(&inputs.len()) || outputs.len() != 1 {
return Err(error(format!(
"expected 3 to {max_inputs} inputs and 1 output, got {} inputs and {} outputs",
inputs.len(),
outputs.len()
)));
}
if inputs[0].dtype == DataType::Float16 {
if self.gate_up_swiglu {
return self.run_f16_gate_up_swiglu(inputs, outputs, workspace);
}
return self.run_f16(inputs, outputs, workspace);
}
if inputs[0].dtype == DataType::BFloat16 {
return self.run_bf16(inputs, outputs, workspace);
}
require_dtype("A", inputs[0].dtype, DataType::Float32)?;
require_dtype("B", inputs[1].dtype, DataType::Uint8)?;
require_dtype("scales", inputs[2].dtype, DataType::Float32)?;
require_dtype("Y", outputs[0].dtype, DataType::Float32)?;
let a_shape = inputs[0].shape;
if a_shape.is_empty() || a_shape[a_shape.len() - 1] != self.k {
return Err(error(format!(
"A must have rank >= 1 and last dimension K={}, got {:?}",
self.k, a_shape
)));
}
let expected_output_shape = [&a_shape[..a_shape.len() - 1], &[self.n]].concat();
if outputs[0].shape != expected_output_shape {
return Err(error(format!(
"Y must have shape {expected_output_shape:?}, got {:?}",
outputs[0].shape
)));
}
let k_blocks = self.k.div_ceil(self.block_size);
let blob_size = self.block_size * self.bits / 8;
require_shape("B", inputs[1].shape, &[self.n, k_blocks, blob_size])?;
require_flat_or_matrix_shape("scales", inputs[2].shape, self.n, k_blocks)?;
let zero_points = optional_input(inputs, 3);
let zp_row_bytes = (k_blocks * self.bits).div_ceil(8);
if let Some(zp) = zero_points {
require_dtype("zero_points", zp.dtype, DataType::Uint8)?;
require_flat_or_matrix_shape("zero_points", zp.shape, self.n, zp_row_bytes)?;
}
let group_indices = optional_input(inputs, 4);
if let Some(g_idx) = group_indices {
require_dtype("g_idx", g_idx.dtype, DataType::Int32)?;
if !g_idx.is_contiguous() {
return Err(error(
"g_idx must be contiguous on the CUDA execution provider",
));
}
let padded_k = k_blocks * self.block_size;
if g_idx.shape != [self.k] && g_idx.shape != [padded_k] {
return Err(error(format!(
"g_idx must have shape [{}] or [{padded_k}], got {:?}",
self.k, g_idx.shape
)));
}
let mut bytes = vec![0u8; g_idx.numel() * 4];
unsafe {
self.runtime
.dtoh(&mut bytes, cuptr(g_idx.data_ptr::<u8>() as *const c_void))?
};
for (index, value) in bytes.chunks_exact(4).enumerate() {
let group = i32::from_ne_bytes([value[0], value[1], value[2], value[3]]);
if group < 0 || group as usize >= k_blocks {
return Err(error(format!(
"g_idx[{index}]={group} is outside 0..{k_blocks}"
)));
}
}
}
let bias = optional_input(inputs, 5);
if let Some(bias) = bias {
require_dtype("bias", bias.dtype, DataType::Float32)?;
require_shape("bias", bias.shape, &[self.n])?;
}
for (name, contiguous) in [
("A", inputs[0].is_contiguous()),
("B", inputs[1].is_contiguous()),
("scales", inputs[2].is_contiguous()),
(
"zero_points",
zero_points.is_none_or(TensorView::is_contiguous),
),
("g_idx", group_indices.is_none_or(TensorView::is_contiguous)),
("bias", bias.is_none_or(TensorView::is_contiguous)),
("Y", outputs[0].is_contiguous()),
] {
if !contiguous {
return Err(error(format!(
"{name} must be contiguous on the CUDA execution provider"
)));
}
}
let m = a_shape[..a_shape.len() - 1].iter().product::<usize>();
crate::trace::record_kernel_metrics(inputs, outputs, || {
let mut flops = (m as u64)
.saturating_mul(self.n as u64)
.saturating_mul(self.k as u64)
.saturating_mul(2);
if bias.is_some() {
flops = flops.saturating_add((m as u64).saturating_mul(self.n as u64));
}
flops
});
self.last_call_capture_safe
.store(m == 1 && group_indices.is_none(), Ordering::Relaxed);
if m == 1 && group_indices.is_none() {
self.release_marlin_repack_for_decode();
if self.bits == 4
&& self.block_size == 128
&& let Some(zero_points) = zero_points
{
onnx_runtime_ep_api::record_kernel_variant!(
"gemv_int4_f32_block128",
"M==1 decode: bits=4, block_size=128, asymmetric → specialized \
shift-indexed int4 f32 GEMV (bit-identical to the generic path)"
);
return self.launch_int4_f32_gemv_block128(
&inputs[0],
&inputs[1],
&inputs[2],
zero_points,
bias,
&mut outputs[0],
k_blocks,
);
}
if self.bits == 8 && self.block_size == 128 && zero_points.is_some() {
onnx_runtime_ep_api::record_kernel_variant!(
"gemv_int8_f32_block128",
"M==1 decode: bits=8, block_size=128, asymmetric → specialized \
shift-indexed int8 f32 GEMV (bit-identical to the generic path)"
);
return self.launch_int8_f32_gemv_block128(
&inputs[0],
&inputs[1],
&inputs[2],
zero_points,
bias,
&mut outputs[0],
k_blocks,
);
}
if self.bits == 8 && self.block_size == 32 {
onnx_runtime_ep_api::record_kernel_variant!(
"gemv_int8_f32",
"M==1 decode: bits=8, block_size=32 → direct capture-safe f32 GEMV"
);
return self.launch_int8_f32_gemv(
&inputs[0],
&inputs[1],
&inputs[2],
zero_points,
bias,
&mut outputs[0],
k_blocks,
);
}
if self.bits == 8 {
onnx_runtime_ep_api::record_kernel_variant!(
"gemv_f32_general_bs_int8",
"M==1 decode: bits=8, block_size={} → model-agnostic f32 GEMV \
(fp32 accumulation, any power-of-two block_size)",
self.block_size
);
return self.launch_f32_gemv(
&inputs[0],
&inputs[1],
&inputs[2],
zero_points,
bias,
&mut outputs[0],
k_blocks,
blob_size,
zp_row_bytes,
);
}
if self.accuracy_level == 4 && self.block_size == 32 && zero_points.is_none() {
onnx_runtime_ep_api::record_kernel_variant!(
"gemv_accuracy4_int8",
"M==1 decode: accuracy_level==4, block_size==32, symmetric (no zero_points) \
→ int8-quantized-activation capture-safe GEMV"
);
return self.launch_accuracy4_gemv(
&inputs[0],
&inputs[1],
&inputs[2],
bias,
&mut outputs[0],
k_blocks,
);
}
if self.bits == 4 && self.accuracy_level == 4 {
if self.block_size == 32 {
onnx_runtime_ep_api::record_kernel_variant!(
"gemv_accuracy4_blockwise_int8",
"M==1 decode: bits=4, accuracy_level==4, block_size=32 → int8-quantized \
activation quantized ONCE then parallelized blockwise GEMV (grid filled \
from the device SM count; bit-identical to the tiled accuracy4 GEMM)"
);
return self.launch_accuracy4_gemv_blockwise(
&inputs[0],
&inputs[1],
&inputs[2],
zero_points,
bias,
&mut outputs[0],
k_blocks,
blob_size,
zp_row_bytes,
);
}
onnx_runtime_ep_api::record_kernel_variant!(
"gemv_f32_general_bs_int4",
"M==1 decode: bits=4, accuracy_level==4, block_size={} (≠32) → \
model-agnostic f32 GEMV (fp32 activations, higher precision than the \
int8-activation reference)",
self.block_size
);
return self.launch_f32_gemv(
&inputs[0],
&inputs[1],
&inputs[2],
zero_points,
bias,
&mut outputs[0],
k_blocks,
blob_size,
zp_row_bytes,
);
}
if self.accuracy_level != 4 {
onnx_runtime_ep_api::record_kernel_variant!(
"gemv_f32",
"M==1 decode: accuracy_level={} (non-accuracy4) → direct f32 GEMV",
self.accuracy_level
);
return self.launch_f32_gemv(
&inputs[0],
&inputs[1],
&inputs[2],
zero_points,
bias,
&mut outputs[0],
k_blocks,
blob_size,
zp_row_bytes,
);
}
}
if self.bits == 4
&& self.accuracy_level == 4
&& self.block_size == 32
&& group_indices.is_none()
{
onnx_runtime_ep_api::record_kernel_variant!(
"gemm_tiled_accuracy4",
"M={} (GEMV requires M==1), accuracy_level==4, block_size=32, no g_idx → \
tiled accuracy4 GEMM (int8-quantized activations)",
m
);
return self.launch_accuracy4(
&inputs[0],
&inputs[1],
&inputs[2],
zero_points,
bias,
&mut outputs[0],
m,
k_blocks,
blob_size,
zp_row_bytes,
);
}
onnx_runtime_ep_api::record_kernel_variant!(
"dequant_cublas_gemm",
"M={}, accuracy_level={}, g_idx={} → dequantize weights to f32 then cuBLAS GEMM \
(general prefill / grouped path)",
m,
self.accuracy_level,
group_indices.is_some()
);
let weight = self.runtime.alloc_raw(self.k * self.n * 4)?;
let result = self
.launch_dequant(
&inputs[1],
&inputs[2],
zero_points,
group_indices,
weight,
k_blocks,
blob_size,
zp_row_bytes,
)
.and_then(|()| {
let params = GemmParams {
dtype: GemmDtype::F32,
a: cuptr(inputs[0].data_ptr::<u8>() as *const c_void),
b: weight,
c: cuptr(outputs[0].data_ptr_mut::<u8>() as *const c_void),
m,
k: self.k,
n: self.n,
batch: 1,
a_batch_stride: m * self.k,
b_batch_stride: 0,
epilogue: bias.map(|bias| GemmEpilogue {
kind: GemmEpilogueKind::Bias,
bias: cuptr(bias.data_ptr::<u8>() as *const c_void),
}),
};
unsafe {
blas::governed_gemm(
self.runtime.blas(),
self.runtime.stream_ptr(),
¶ms,
workspace,
"MatMulNBits",
)
}
})
.and_then(|()| self.runtime.synchronize());
let free_weight = unsafe { self.runtime.free_raw(weight) };
result.and(free_weight)
}
fn uses_dequant_cublas_workspace(
&self,
dtype: DataType,
a_shape: &[usize],
group_indices_present: bool,
) -> bool {
if dtype != DataType::Float32 || a_shape.is_empty() || a_shape[a_shape.len() - 1] != self.k
{
return false;
}
let m = a_shape[..a_shape.len() - 1].iter().product::<usize>();
if m == 1 && !group_indices_present {
return false;
}
!(self.bits == 4
&& self.accuracy_level == 4
&& self.block_size == 32
&& !group_indices_present)
}
fn uses_dequant_f16_cublas_workspace(
&self,
dtype: DataType,
a_shape: &[usize],
group_indices_present: bool,
) -> bool {
if !matches!(dtype, DataType::Float16 | DataType::BFloat16)
|| a_shape.is_empty()
|| a_shape[a_shape.len() - 1] != self.k
|| group_indices_present
{
return false;
}
let m = a_shape[..a_shape.len() - 1].iter().product::<usize>();
dequant_f16_gemm_enabled()
&& m >= dequant_f16_gemm_min_m()
&& self.k.saturating_mul(self.n).saturating_mul(2)
<= dequant_f16_gemm_max_scratch_bytes()
}
fn workspace_requirement_for(
&self,
inputs: &[TensorMetadata<'_>],
) -> Result<WorkspaceRequirement> {
let Some(a) = inputs.first() else {
return Ok(WorkspaceRequirement::NONE);
};
let group_indices_present = inputs.get(4).is_some_and(|input| input.present);
let fused = self.gate_up_swiglu;
let dtype = if self.uses_dequant_cublas_workspace(a.dtype, a.shape, group_indices_present) {
GemmDtype::F32
} else if self.uses_dequant_f16_cublas_workspace(
a.dtype,
a.shape,
!fused && group_indices_present,
) {
GemmDtype::F16
} else {
return Ok(WorkspaceRequirement::NONE);
};
let m = a.shape[..a.shape.len() - 1].iter().product::<usize>();
if dtype == GemmDtype::F16 {
let mut bytes = 0;
let biases: &[Option<u64>] = if fused { &[None] } else { &[None, Some(0)] };
for bias in biases {
let params = self.dequant_f16_gemm_ex(1, 1, 1, m, *bias);
bytes = bytes.max(blas::gemm_ex_workspace_bytes(self.runtime.blas(), ¶ms)?);
}
return Ok(blas::governed_workspace_requirement(bytes));
}
let params = GemmParams {
dtype,
a: 1,
b: 1,
c: 1,
m,
k: self.k,
n: self.n,
batch: 1,
a_batch_stride: m * self.k,
b_batch_stride: 0,
epilogue: inputs
.get(5)
.filter(|bias| bias.present)
.map(|_| GemmEpilogue {
kind: GemmEpilogueKind::Bias,
bias: 0,
}),
};
let bytes = blas::gemm_workspace_bytes(self.runtime.blas(), ¶ms)?;
Ok(blas::governed_workspace_requirement(bytes))
}
fn run_bf16(
&self,
inputs: &[TensorView],
outputs: &mut [TensorMut],
workspace: Option<WorkspaceView>,
) -> Result<()> {
require_dtype("A", inputs[0].dtype, DataType::BFloat16)?;
require_dtype("Y", outputs[0].dtype, DataType::BFloat16)?;
let cache_slots: &[usize] = if self.gate_up_swiglu { &[2, 4] } else { &[2] };
let is_cached = |index: usize, dtype: DataType| {
dtype == DataType::BFloat16 && cache_slots.contains(&index)
};
const ALIGN: usize = 256;
let f16 = std::mem::size_of::<half::f16>();
let round = |bytes: usize| bytes.div_ceil(ALIGN) * ALIGN;
let mut offsets: Vec<Option<usize>> = Vec::with_capacity(inputs.len());
let mut total = 0usize;
for (index, input) in inputs.iter().enumerate() {
if input.dtype == DataType::BFloat16 && !is_cached(index, input.dtype) {
offsets.push(Some(total));
total += round(input.numel() * f16);
} else {
offsets.push(None);
}
}
let out_off = total;
let out_n = outputs[0].numel();
total += round(out_n * f16);
let cached: Vec<(usize, CUdeviceptr)> = {
let consts: Vec<(CUdeviceptr, usize)> = cache_slots
.iter()
.filter(|index| **index < inputs.len() && is_cached(**index, inputs[**index].dtype))
.map(|index| {
(
cuptr(inputs[*index].data_ptr::<u8>() as *const c_void) as CUdeviceptr,
inputs[*index].numel(),
)
})
.collect();
if consts.is_empty() {
Vec::new()
} else {
let mut cache = self
.bf16_const_cache
.lock()
.map_err(|_| error("MatMulNBits bf16 const cache mutex poisoned"))?;
let ptrs = cache.staged(&consts)?;
cache_slots
.iter()
.filter(|index| {
**index < inputs.len() && is_cached(**index, inputs[**index].dtype)
})
.zip(ptrs)
.map(|(index, ptr)| (*index, ptr))
.collect()
}
};
let mut arena = self
.bf16_scratch
.lock()
.map_err(|_| error("MatMulNBits bf16 scratch mutex poisoned"))?;
let base = arena.ensure(total)?;
let mut f16_inputs: Vec<TensorView> = Vec::with_capacity(inputs.len());
for (index, (input, offset)) in inputs.iter().zip(offsets.iter()).enumerate() {
if let Some((_, cached_ptr)) = cached.iter().find(|(slot, _)| *slot == index) {
f16_inputs.push(TensorView::new(
DevicePtr(raw_ptr(*cached_ptr) as *const c_void),
DataType::Float16,
input.shape,
input.strides,
input.device,
));
continue;
}
match offset {
Some(off) => {
let ptr = base + *off as CUdeviceptr;
super::cast::launch_cast_raw(
&self.runtime,
cuptr(input.data_ptr::<u8>() as *const c_void),
DataType::BFloat16,
ptr,
DataType::Float16,
input.numel(),
)?;
f16_inputs.push(TensorView::new(
DevicePtr(raw_ptr(ptr) as *const c_void),
DataType::Float16,
input.shape,
input.strides,
input.device,
));
}
None => f16_inputs.push(*input),
}
}
let out_ptr = base + out_off as CUdeviceptr;
let out_shape = outputs[0].shape;
let out_strides = outputs[0].strides;
let out_device = outputs[0].device;
let mut y_f16 = TensorMut::new(
DevicePtrMut(raw_ptr(out_ptr)),
DataType::Float16,
out_shape,
out_strides,
out_device,
);
let dst_ptr = cuptr(outputs[0].data_ptr_mut::<u8>() as *const c_void);
let offer = Some(Bf16DirectOut {
staging: out_ptr,
dst: dst_ptr,
elements: out_n,
taken: false,
});
let previous = BF16_DIRECT_OUT.with(|cell| cell.replace(offer));
let run_result = self.run(&f16_inputs, std::slice::from_mut(&mut y_f16), workspace);
let settled = BF16_DIRECT_OUT.with(|cell| cell.replace(previous));
run_result?;
if !settled.is_some_and(|offer| offer.taken) {
super::cast::launch_cast_raw(
&self.runtime,
out_ptr,
DataType::Float16,
dst_ptr,
DataType::BFloat16,
out_n,
)?;
}
drop(arena);
Ok(())
}
fn run_f16(
&self,
inputs: &[TensorView],
outputs: &mut [TensorMut],
workspace: Option<WorkspaceView>,
) -> Result<()> {
require_dtype("A", inputs[0].dtype, DataType::Float16)?;
require_dtype("B", inputs[1].dtype, DataType::Uint8)?;
let scales_fp16 = match inputs[2].dtype {
DataType::Float16 => true,
DataType::Float32 => false,
other => {
return Err(error(format!(
"scales must have dtype Float16 or Float32 for fp16 activations, got {other:?}"
)));
}
};
require_dtype("Y", outputs[0].dtype, DataType::Float16)?;
let a_shape = inputs[0].shape;
if a_shape.is_empty() || a_shape[a_shape.len() - 1] != self.k {
return Err(error(format!(
"A must have rank >= 1 and last dimension K={}, got {:?}",
self.k, a_shape
)));
}
let expected_output_shape = [&a_shape[..a_shape.len() - 1], &[self.n]].concat();
if outputs[0].shape != expected_output_shape {
return Err(error(format!(
"Y must have shape {expected_output_shape:?}, got {:?}",
outputs[0].shape
)));
}
let k_blocks = self.k.div_ceil(self.block_size);
let blob_size = self.block_size * self.bits / 8;
require_shape("B", inputs[1].shape, &[self.n, k_blocks, blob_size])?;
require_flat_or_matrix_shape("scales", inputs[2].shape, self.n, k_blocks)?;
let zero_points = optional_input(inputs, 3);
let zp_row_bytes = (k_blocks * self.bits).div_ceil(8);
if let Some(zero_points) = zero_points {
require_dtype("zero_points", zero_points.dtype, DataType::Uint8)?;
require_flat_or_matrix_shape("zero_points", zero_points.shape, self.n, zp_row_bytes)?;
}
let group_indices = optional_input(inputs, 4);
let bias = optional_input(inputs, 5);
let rows = a_shape[..a_shape.len() - 1].iter().product::<usize>();
if let Some(bias) = bias {
require_dtype("bias", bias.dtype, DataType::Float16)?;
if bias.numel() != self.n && bias.numel() != rows * self.n {
return Err(error(format!(
"bias must have {} elements (broadcast [N]) or {} elements (per-token \
[1, S, N] residual), got {:?}",
self.n,
rows * self.n,
bias.shape
)));
}
}
let gamma = optional_input(inputs, 6);
if self.rmsnorm_prologue {
let gamma = gamma.ok_or_else(|| {
error("rmsnorm_prologue fusion requires the normalization weight at input 6")
})?;
require_gamma_dtype(gamma.dtype)?;
require_shape("gamma", gamma.shape, &[self.k])?;
}
for (name, contiguous) in [
("A", inputs[0].is_contiguous()),
("B", inputs[1].is_contiguous()),
("scales", inputs[2].is_contiguous()),
(
"zero_points",
zero_points.is_none_or(TensorView::is_contiguous),
),
("bias", bias.is_none_or(TensorView::is_contiguous)),
("gamma", gamma.is_none_or(TensorView::is_contiguous)),
("Y", outputs[0].is_contiguous()),
] {
if !contiguous {
return Err(error(format!(
"{name} must be contiguous on the CUDA execution provider"
)));
}
}
let m = a_shape[..a_shape.len() - 1].iter().product::<usize>();
crate::trace::record_kernel_metrics(inputs, outputs, || {
let mut flops = (m as u64)
.saturating_mul(self.n as u64)
.saturating_mul(self.k as u64)
.saturating_mul(2);
if bias.is_some() {
flops = flops.saturating_add((m as u64).saturating_mul(self.n as u64));
}
if self.rmsnorm_prologue {
let elements = (m as u64).saturating_mul(self.k as u64);
flops = flops
.saturating_add(elements.saturating_mul(4))
.saturating_add((m as u64).saturating_mul(4));
}
flops
});
if group_indices.is_some() {
return Err(error(
"MatMulNBits CUDA fp16 activations do not support g_idx. Why: the block-32 fp16 \
kernels map each K block directly to its scale and zero point and do not implement \
group remapping. How to fix: omit g_idx, provide f32 activations, or select another \
execution provider",
));
}
if m > 1 {
if m <= decode_gemv_loop_max_m()
&& decode_gemv_loop_rows_aligned(self.k)
&& !self.gate_up_swiglu
&& !self.decomposed_silu
{
onnx_runtime_ep_api::record_kernel_variant!(
"gemv_f16_batched_loop",
"M={} small-batch decode: {} single-row decode GEMV launches (one per row), \
each byte-identical to M==1 decode; skips the tiled prefill GEMM's \
M-independent full-weight-grid pass",
m,
m
);
self.last_call_capture_safe.store(true, Ordering::Relaxed);
let per_token_bias =
bias.filter(|b| b.numel() == m * self.n && m * self.n != self.n);
let a_base = inputs[0].data_ptr::<u8>();
let y_base = outputs[0].data_ptr_mut::<u8>();
let bias_base = per_token_bias.map(|b| b.data_ptr::<u8>());
let a_row_shape = [1usize, self.k];
let a_row_strides = [self.k as i64, 1];
let y_row_shape = [1usize, self.n];
let y_row_strides = [self.n as i64, 1];
let a_row_bytes = self.k * 2; let y_row_bytes = self.n * 2; for row in 0..m {
let a_row = TensorView::new(
DevicePtr(a_base.wrapping_add(row * a_row_bytes) as *const c_void),
DataType::Float16,
&a_row_shape,
&a_row_strides,
inputs[0].device,
);
let mut y_row = TensorMut::new(
DevicePtrMut(y_base.wrapping_add(row * y_row_bytes) as *mut c_void),
DataType::Float16,
&y_row_shape,
&y_row_strides,
outputs[0].device,
);
let bias_row = bias_base.map(|base| {
TensorView::new(
DevicePtr(base.wrapping_add(row * y_row_bytes) as *const c_void),
DataType::Float16,
&y_row_shape,
&y_row_strides,
inputs[0].device,
)
});
let bias_ref = match bias_row {
Some(ref b) => Some(b),
None => bias,
};
self.dispatch_f16_decode_gemv_row(
&a_row,
&inputs[1],
&inputs[2],
scales_fp16,
zero_points,
bias_ref,
gamma,
&mut y_row,
k_blocks,
blob_size,
zp_row_bytes,
)?;
}
return Ok(());
}
self.last_call_capture_safe.store(false, Ordering::Relaxed);
let bias_row_stride = match bias {
Some(bias) if bias.numel() == m * self.n && m * self.n != self.n => self.n,
_ => 0,
};
if dequant_f16_gemm_enabled()
&& m >= dequant_f16_gemm_min_m()
&& !self.gate_up_swiglu
&& !self.decomposed_silu
&& !self.rmsnorm_prologue
&& group_indices.is_none()
{
match self.try_dequant_f16_cublas_gemm(
&inputs[0],
&inputs[1],
&inputs[2],
scales_fp16,
zero_points,
bias,
&mut outputs[0],
m,
bias_row_stride,
k_blocks,
blob_size,
zp_row_bytes,
workspace,
) {
Ok(true) => {
onnx_runtime_ep_api::record_kernel_variant!(
"gemm_dequant_f16_cublas",
"M={} prefill: dequantize int{} weights to [K, N] fp16, then \
cuBLASLt fp16 tensor-core GEMM with f32 accumulation",
m,
self.bits
);
return Ok(());
}
Ok(false) => {
}
Err(_err) => {
}
}
}
if marlin_gemm::marlin_m_gt_1_enabled()
&& self.bits == 4
&& !self.gate_up_swiglu
&& !self.decomposed_silu
&& !self.rmsnorm_prologue
&& self.marlin_weight_inputs_constant()
&& marlin_gemm::device_supports_marlin(
self.runtime.capabilities().compute_capability(),
)
{
match self.try_launch_marlin_gemm(
&inputs[0],
&inputs[1],
&inputs[2],
scales_fp16,
zero_points,
bias,
&mut outputs[0],
m,
bias_row_stride,
) {
Ok(Some(warm)) => {
self.last_call_capture_safe.store(warm, Ordering::Relaxed);
onnx_runtime_ep_api::record_kernel_variant!(
"gemm_marlin_int4",
"M={} prefill/verify: fp16 activation, int4, block_size={}, \
zero_points={} → Marlin SM80 mma.sync int4 tensor-core GEMM \
(static grid, capture-safe when weights are pre-repacked)",
m,
self.block_size,
zero_points.is_some()
);
return Ok(());
}
Ok(None) => {
}
Err(_err) => {
}
}
}
if self.rmsnorm_prologue {
let gamma = gamma.ok_or_else(|| {
error("rmsnorm_prologue fusion requires the normalization weight at input 6")
})?;
if marlin_gemm::marlin_m_gt_1_enabled()
&& self.bits == 4
&& !self.gate_up_swiglu
&& !self.decomposed_silu
&& self.marlin_weight_inputs_constant()
&& marlin_gemm::device_supports_marlin(
self.runtime.capabilities().compute_capability(),
)
{
match self.try_launch_marlin_gemm_rmsnorm(
&inputs[0],
&inputs[1],
&inputs[2],
scales_fp16,
zero_points,
gamma,
bias,
&mut outputs[0],
m,
bias_row_stride,
) {
Ok(Some(warm)) => {
self.last_call_capture_safe.store(warm, Ordering::Relaxed);
onnx_runtime_ep_api::record_kernel_variant!(
"gemm_marlin_int4_rmsnorm",
"M={} prefill/verify: RMS-normalization prologue \
(SkipSimplifiedLayerNormalization folded) into pooled scratch, \
then Marlin SM80 mma.sync int4 tensor-core GEMM (capture-safe \
when weights + scratch are pre-warmed)",
m
);
return Ok(());
}
Ok(None) => {
}
Err(_err) => {
}
}
}
onnx_runtime_ep_api::record_kernel_variant!(
"gemm_f16_tiled_rmsnorm",
"M={} prefill: RMS-normalization prologue (SkipSimplifiedLayerNormalization \
folded) into a per-token scratch, then portable 16x16 tiled GEMM with fp32 \
accumulation; not advertised as CUDA-graph capture-safe",
m
);
return self.launch_f16_gemm_rmsnorm_prefill(
&inputs[0],
&inputs[1],
&inputs[2],
zero_points,
gamma,
bias,
&mut outputs[0],
m,
k_blocks,
bias_row_stride,
);
}
onnx_runtime_ep_api::record_kernel_variant!(
"gemm_f16_tiled",
"M={} prefill: fp16 activation, bits={}, block_size={}, zero_points={}, \
scales={} → portable 16x16 CUDA-core tiled GEMM with fp32 accumulation; \
not advertised as CUDA-graph capture-safe",
m,
self.bits,
self.block_size,
zero_points.is_some(),
if scales_fp16 { "fp16" } else { "fp32" }
);
return self.launch_f16_gemm(
&inputs[0],
&inputs[1],
&inputs[2],
scales_fp16,
zero_points,
bias,
&mut outputs[0],
m,
k_blocks,
blob_size,
bias_row_stride,
);
}
self.last_call_capture_safe.store(true, Ordering::Relaxed);
self.dispatch_f16_decode_gemv_row(
&inputs[0],
&inputs[1],
&inputs[2],
scales_fp16,
zero_points,
bias,
gamma,
&mut outputs[0],
k_blocks,
blob_size,
zp_row_bytes,
)
}
#[allow(clippy::too_many_arguments)]
fn dispatch_f16_decode_gemv_row(
&self,
a_row: &TensorView,
packed: &TensorView,
scales: &TensorView,
scales_fp16: bool,
zero_points: Option<&TensorView>,
bias: Option<&TensorView>,
gamma: Option<&TensorView>,
y_row: &mut TensorMut,
k_blocks: usize,
blob_size: usize,
zp_row_bytes: usize,
) -> Result<()> {
if self.bits == 8 && self.block_size == 32 {
if self.rmsnorm_prologue {
let gamma = gamma.ok_or_else(|| {
error("rmsnorm_prologue fusion requires the normalization weight at input 6")
})?;
if !scales_fp16 {
return Err(error(
"rmsnorm_prologue fusion requires fp16 scales (the fused kernel replicates \
the fp16 general scales path)",
));
}
onnx_runtime_ep_api::record_kernel_variant!(
"gemv_int8_f16_scales_f16_rmsnorm",
"M==1 decode: fp16 activation, bits=8, block_size=32, fp16 scales, \
zero_points={} → int8 GEMV with fused RMS-normalization prologue \
(SkipSimplifiedLayerNormalization folded)",
zero_points.is_some()
);
return self.launch_int8_f16_gemv_rmsnorm(
a_row,
packed,
scales,
zero_points,
gamma,
bias,
y_row,
k_blocks,
);
}
onnx_runtime_ep_api::record_kernel_variant!(
"gemv_int8_f16",
"M==1 decode: fp16 activation, bits=8, block_size=32, zero_points={} → direct \
capture-safe GEMV",
zero_points.is_some()
);
return self.launch_int8_f16_gemv(
a_row,
packed,
scales,
scales_fp16,
zero_points,
bias,
y_row,
k_blocks,
);
}
if self.rmsnorm_prologue {
let gamma = gamma.ok_or_else(|| {
error("rmsnorm_prologue fusion requires the normalization weight at input 6")
})?;
if !scales_fp16 {
return Err(error(
"rmsnorm_prologue fusion requires fp16 scales (the fused kernel replicates \
the fp16 general scales path)",
));
}
onnx_runtime_ep_api::record_kernel_variant!(
"gemv_f16_scales_f16_rmsnorm",
"M==1 decode: fp16 activation, bits=4, block_size=32, fp16 scales, \
zero_points={} → general GEMV with fused RMS-normalization prologue \
(SkipSimplifiedLayerNormalization folded)",
zero_points.is_some()
);
return self.launch_f16_gemv_rmsnorm(
a_row,
packed,
scales,
zero_points,
gamma,
bias,
y_row,
k_blocks,
blob_size,
zp_row_bytes,
);
}
self.launch_f16_gemv(
a_row,
packed,
scales,
scales_fp16,
zero_points,
bias,
y_row,
k_blocks,
blob_size,
zp_row_bytes,
)
}
#[allow(clippy::too_many_arguments)]
#[allow(clippy::too_many_arguments)]
fn try_launch_marlin_gemm(
&self,
activation: &TensorView,
packed: &TensorView,
scales: &TensorView,
scales_fp16: bool,
zero_points: Option<&TensorView>,
bias: Option<&TensorView>,
output: &mut TensorMut,
m: usize,
bias_row_stride: usize,
) -> Result<Option<bool>> {
if self.block_size == 0
|| !self.k.is_multiple_of(16)
|| !self.k.is_multiple_of(self.block_size)
{
return Ok(None);
}
self.runtime
.require_nvrtc_half_headers("MatMulNBits Marlin int4 tensor-core GEMM")?;
let packed_ptr = cuptr(packed.data_ptr::<u8>() as *const c_void);
let (weights_ptr, warm) = self.marlin_repack_cache.ensure_repacked(
packed_ptr,
self.n,
self.k,
self.block_size,
)?;
let args = marlin_gemm::MarlinGemmArgs {
activation: cuptr(activation.data_ptr::<u8>() as *const c_void),
weights: weights_ptr,
scales: cuptr(scales.data_ptr::<u8>() as *const c_void),
zero_points: zero_points.map(|t| cuptr(t.data_ptr::<u8>() as *const c_void)),
bias: bias.map(|t| cuptr(t.data_ptr::<u8>() as *const c_void)),
output: cuptr(output.data_ptr_mut::<u8>() as *const c_void),
m,
k: self.k,
n: self.n,
group_size: self.block_size,
scales_fp16,
bias_post_round: self.fold_bias_post_round && bias.is_some(),
bias_row_stride,
};
let split_warm = self.maybe_launch_marlin_splitk(&args)?;
Ok(Some(warm && split_warm))
}
fn maybe_launch_marlin_splitk(&self, args: &marlin_gemm::MarlinGemmArgs) -> Result<bool> {
if marlin_gemm::marlin_splitk_enabled() {
let k_blocks = self.k / self.block_size.max(1);
let sm = self.runtime.capabilities().multiprocessor_count();
let split_k = marlin_gemm::choose_split_k(args.m, args.n, k_blocks, sm);
if split_k > 1 {
let bytes = marlin_gemm::splitk_partials_len(split_k, args.m, args.n)
* std::mem::size_of::<f32>();
let (partials, warm) = self.marlin_repack_cache.ensure_scratch(4, bytes)?;
marlin_gemm::launch_marlin_gemm_splitk(&self.runtime, args, split_k, partials)?;
return Ok(warm);
}
}
marlin_gemm::launch_marlin_gemm(&self.runtime, args)?;
Ok(true)
}
#[allow(clippy::too_many_arguments)]
fn try_launch_marlin_gemm_rmsnorm(
&self,
activation: &TensorView,
packed: &TensorView,
scales: &TensorView,
scales_fp16: bool,
zero_points: Option<&TensorView>,
gamma: &TensorView,
bias: Option<&TensorView>,
output: &mut TensorMut,
m: usize,
bias_row_stride: usize,
) -> Result<Option<bool>> {
if self.block_size == 0
|| !self.k.is_multiple_of(16)
|| !self.k.is_multiple_of(self.block_size)
{
return Ok(None);
}
self.runtime
.require_nvrtc_half_headers("MatMulNBits Marlin int4 tensor-core GEMM (rmsnorm)")?;
let packed_ptr = cuptr(packed.data_ptr::<u8>() as *const c_void);
let (weights_ptr, weights_warm) = self.marlin_repack_cache.ensure_repacked(
packed_ptr,
self.n,
self.k,
self.block_size,
)?;
let scratch_bytes = m * self.k * std::mem::size_of::<half::f16>();
let (scratch, scratch_warm) = self.marlin_repack_cache.ensure_scratch(0, scratch_bytes)?;
self.launch_rmsnorm_prefill(activation, gamma, scratch, m)?;
let args = marlin_gemm::MarlinGemmArgs {
activation: scratch,
weights: weights_ptr,
scales: cuptr(scales.data_ptr::<u8>() as *const c_void),
zero_points: zero_points.map(|t| cuptr(t.data_ptr::<u8>() as *const c_void)),
bias: bias.map(|t| cuptr(t.data_ptr::<u8>() as *const c_void)),
output: cuptr(output.data_ptr_mut::<u8>() as *const c_void),
m,
k: self.k,
n: self.n,
group_size: self.block_size,
scales_fp16,
bias_post_round: self.fold_bias_post_round && bias.is_some(),
bias_row_stride,
};
let split_warm = self.maybe_launch_marlin_splitk(&args)?;
Ok(Some(weights_warm && scratch_warm && split_warm))
}
#[allow(clippy::too_many_arguments)]
fn try_dequant_f16_gate_up_prefill(
&self,
activation: &TensorView,
packed_gate: &TensorView,
scales_gate: &TensorView,
packed_up: &TensorView,
scales_up: &TensorView,
scales_fp16: bool,
zp_gate: Option<&TensorView>,
zp_up: Option<&TensorView>,
gamma: Option<&TensorView>,
output: &mut TensorMut,
m: usize,
k_blocks: usize,
blob_size: usize,
zp_row_bytes: usize,
workspace: Option<WorkspaceView>,
) -> Result<bool> {
let Some(block_shift) = self.dequant_f16_block_shift() else {
return Ok(false);
};
let weight_bytes = self
.k
.checked_mul(self.n)
.and_then(|elems| elems.checked_mul(2))
.ok_or_else(|| error("dequantized f16 weight size overflowed"))?;
if weight_bytes > dequant_f16_gemm_max_scratch_bytes() {
return Ok(false);
}
self.runtime
.require_nvrtc_half_headers("MatMulNBits dequant f16 gate/up SwiGLU GEMM")?;
let act_ptr = if let Some(gamma) = gamma {
let norm_bytes = m * self.k * std::mem::size_of::<half::f16>();
let (norm, _warm) = self.marlin_repack_cache.ensure_scratch(0, norm_bytes)?;
self.launch_rmsnorm_prefill(activation, gamma, norm, m)?;
norm
} else {
cuptr(activation.data_ptr::<u8>() as *const c_void)
};
let out_bytes = output.byte_size();
let (gate_buf, _warm) = self.marlin_repack_cache.ensure_scratch(1, out_bytes)?;
let (weight_gate, _warm) = self.marlin_repack_cache.ensure_scratch(5, weight_bytes)?;
let (weight_up, _warm) = self.marlin_repack_cache.ensure_scratch(6, weight_bytes)?;
let output_ptr = cuptr(output.data_ptr_mut::<u8>() as *const c_void);
for (packed, scales, zero_points, weight, out) in [
(packed_gate, scales_gate, zp_gate, weight_gate, gate_buf),
(packed_up, scales_up, zp_up, weight_up, output_ptr),
] {
self.launch_dequant_f16(
packed,
scales,
scales_fp16,
zero_points,
weight,
k_blocks,
blob_size,
zp_row_bytes,
block_shift,
)?;
let params = self.dequant_f16_gemm_ex(act_ptr, weight, out, m, None);
unsafe {
blas::governed_gemm_ex(
self.runtime.blas(),
self.runtime.stream_ptr(),
¶ms,
workspace,
"MatMulNBits",
)
}?;
}
crate::kernels::elementwise::launch_silu_mul_f16_raw(
&self.runtime,
gate_buf,
output_ptr,
output_ptr,
output.numel(),
self.decomposed_silu,
)?;
Ok(true)
}
#[allow(clippy::too_many_arguments)]
fn try_launch_marlin_gate_up_prefill(
&self,
activation: &TensorView,
packed_gate: &TensorView,
scales_gate: &TensorView,
packed_up: &TensorView,
scales_up: &TensorView,
zp_gate: Option<&TensorView>,
zp_up: Option<&TensorView>,
gamma: Option<&TensorView>,
output: &mut TensorMut,
m: usize,
) -> Result<Option<bool>> {
if self.block_size == 0
|| !self.k.is_multiple_of(16)
|| !self.k.is_multiple_of(self.block_size)
{
return Ok(None);
}
self.runtime
.require_nvrtc_half_headers("MatMulNBits Marlin int4 gate/up SwiGLU GEMM")?;
let packed_gate_ptr = cuptr(packed_gate.data_ptr::<u8>() as *const c_void);
let packed_up_ptr = cuptr(packed_up.data_ptr::<u8>() as *const c_void);
let (weights_gate, gate_w_warm) = self.marlin_repack_cache.ensure_repacked(
packed_gate_ptr,
self.n,
self.k,
self.block_size,
)?;
let (weights_up, up_w_warm) = self.marlin_repack_cache.ensure_repacked(
packed_up_ptr,
self.n,
self.k,
self.block_size,
)?;
let (act_ptr, norm_warm) = if let Some(gamma) = gamma {
let norm_bytes = m * self.k * std::mem::size_of::<half::f16>();
let (norm, norm_warm) = self.marlin_repack_cache.ensure_scratch(0, norm_bytes)?;
self.launch_rmsnorm_prefill(activation, gamma, norm, m)?;
(norm, norm_warm)
} else {
(cuptr(activation.data_ptr::<u8>() as *const c_void), true)
};
let out_bytes = output.byte_size();
let (gate_buf, gate_s_warm) = self.marlin_repack_cache.ensure_scratch(1, out_bytes)?;
let output_ptr = cuptr(output.data_ptr_mut::<u8>() as *const c_void);
let gate_args = marlin_gemm::MarlinGemmArgs {
activation: act_ptr,
weights: weights_gate,
scales: cuptr(scales_gate.data_ptr::<u8>() as *const c_void),
zero_points: zp_gate.map(|t| cuptr(t.data_ptr::<u8>() as *const c_void)),
bias: None,
output: gate_buf,
m,
k: self.k,
n: self.n,
group_size: self.block_size,
scales_fp16: true,
bias_post_round: false,
bias_row_stride: 0,
};
let gate_split_warm = self.maybe_launch_marlin_splitk(&gate_args)?;
let up_args = marlin_gemm::MarlinGemmArgs {
activation: act_ptr,
weights: weights_up,
scales: cuptr(scales_up.data_ptr::<u8>() as *const c_void),
zero_points: zp_up.map(|t| cuptr(t.data_ptr::<u8>() as *const c_void)),
bias: None,
output: output_ptr,
m,
k: self.k,
n: self.n,
group_size: self.block_size,
scales_fp16: true,
bias_post_round: false,
bias_row_stride: 0,
};
let up_split_warm = self.maybe_launch_marlin_splitk(&up_args)?;
crate::kernels::elementwise::launch_silu_mul_f16_raw(
&self.runtime,
gate_buf,
output_ptr,
output_ptr,
output.numel(),
self.decomposed_silu,
)?;
Ok(Some(
gate_w_warm
&& up_w_warm
&& norm_warm
&& gate_s_warm
&& gate_split_warm
&& up_split_warm,
))
}
#[allow(clippy::too_many_arguments)]
fn launch_f16_gemm(
&self,
activation: &TensorView,
packed: &TensorView,
scales: &TensorView,
scales_fp16: bool,
zero_points: Option<&TensorView>,
bias: Option<&TensorView>,
output: &mut TensorMut,
m: usize,
k_blocks: usize,
blob_size: usize,
bias_row_stride: usize,
) -> Result<()> {
self.runtime
.require_nvrtc_half_headers("MatMulNBits fp16 prefill GEMM")?;
let general_block_size = self.block_size != 32;
let entry = if general_block_size {
GEMM_F16_GENERAL_BS_ENTRY
} else {
GEMM_F16_ENTRY
};
let function = self
.runtime
.nvrtc_function(GEMV_F16_MODULE, GEMV_F16_SRC, entry)?;
let activation_ptr = cuptr(activation.data_ptr::<u8>() as *const c_void);
let packed_ptr = cuptr(packed.data_ptr::<u8>() as *const c_void);
let scales_ptr = cuptr(scales.data_ptr::<u8>() as *const c_void);
let zero_points_ptr = zero_points
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let bias_ptr = bias
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let output_ptr = cuptr(output.data_ptr_mut::<u8>() as *const c_void);
let m_i32 = as_i32("M", m)?;
let k = as_i32("K", self.k)?;
let n = as_i32("N", self.n)?;
let k_blocks = as_i32("K block count", k_blocks)?;
let bits = as_i32("bits", self.bits)?;
let scales_fp16_flag = scales_fp16 as i32;
let bias_post_round_flag: i32 = (self.fold_bias_post_round && bias.is_some()) as i32;
let bias_row_stride_i32 = as_i32("bias row stride", bias_row_stride)?;
let block_size_i32 = as_i32("block_size", self.block_size)?;
let blob_size_i32 = as_i32("block blob size", blob_size)?;
let grid_x = u32::try_from(self.n.div_ceil(GEMM_F16_TILE))
.map_err(|_| error(format!("N={} exceeds CUDA prefill grid limits", self.n)))?;
let grid_y = u32::try_from(m.div_ceil(GEMM_F16_TILE))
.map_err(|_| error(format!("M={m} exceeds CUDA prefill grid limits")))?;
let mut builder = self.runtime.stream().launch_builder(&function);
builder
.arg(&activation_ptr)
.arg(&packed_ptr)
.arg(&scales_ptr)
.arg(&zero_points_ptr)
.arg(&bias_ptr)
.arg(&output_ptr)
.arg(&m_i32)
.arg(&k)
.arg(&n)
.arg(&k_blocks)
.arg(&bits)
.arg(&scales_fp16_flag)
.arg(&bias_post_round_flag)
.arg(&bias_row_stride_i32);
if general_block_size {
builder.arg(&block_size_i32).arg(&blob_size_i32);
}
unsafe {
builder.launch(LaunchConfig {
grid_dim: (grid_x, grid_y, 1),
block_dim: (GEMM_F16_TILE as u32, GEMM_F16_TILE as u32, 1),
shared_mem_bytes: 0,
})
}
.map(|_| ())
.map_err(|err| driver_err("launch MatMulNBits fp16 prefill GEMM", err))
}
#[allow(clippy::too_many_arguments)]
fn launch_f16_gemm_rmsnorm_prefill(
&self,
activation: &TensorView,
packed: &TensorView,
scales: &TensorView,
zero_points: Option<&TensorView>,
gamma: &TensorView,
bias: Option<&TensorView>,
output: &mut TensorMut,
m: usize,
k_blocks: usize,
bias_row_stride: usize,
) -> Result<()> {
let scratch = self
.runtime
.alloc_raw(m * self.k * std::mem::size_of::<half::f16>())?;
let scratch_shape = [m, self.k];
let scratch_strides = [self.k as i64, 1];
let normalized = TensorView::new(
DevicePtr(raw_ptr(scratch) as *const c_void),
DataType::Float16,
&scratch_shape,
&scratch_strides,
activation.device,
);
let result = self
.launch_rmsnorm_prefill(activation, gamma, scratch, m)
.and_then(|()| {
self.launch_f16_gemm(
&normalized,
packed,
scales,
true,
zero_points,
bias,
output,
m,
k_blocks,
self.block_size * self.bits / 8,
bias_row_stride,
)
});
let free_scratch = unsafe { self.runtime.free_raw(scratch) };
result.and(free_scratch)
}
fn launch_rmsnorm_prefill(
&self,
activation: &TensorView,
gamma: &TensorView,
normalized: CUdeviceptr,
m: usize,
) -> Result<()> {
self.runtime
.require_nvrtc_half_headers("MatMulNBits fp16 RMS-norm prefill prologue")?;
let function =
self.runtime
.nvrtc_function(GEMV_F16_MODULE, GEMV_F16_SRC, RMSNORM_PREFILL_ENTRY)?;
let activation_ptr = cuptr(activation.data_ptr::<u8>() as *const c_void);
let gamma_ptr = cuptr(gamma.data_ptr::<u8>() as *const c_void);
let normalized_ptr = normalized;
let norm_size = as_i32("K", self.k)?;
let num_groups = as_i32("M", m)?;
let gamma_is_half: i32 = (gamma.dtype == DataType::Float16) as i32;
let epsilon = self.rmsnorm_epsilon;
let mut builder = self.runtime.stream().launch_builder(&function);
builder
.arg(&activation_ptr)
.arg(&gamma_ptr)
.arg(&normalized_ptr)
.arg(&norm_size)
.arg(&num_groups)
.arg(&gamma_is_half)
.arg(&epsilon);
unsafe {
builder.launch(LaunchConfig {
grid_dim: (m as u32, 1, 1),
block_dim: (RMSNORM_PREFILL_THREADS, 1, 1),
shared_mem_bytes: 0,
})
}
.map(|_| ())
.map_err(|err| driver_err("launch MatMulNBits fp16 RMS-norm prefill prologue", err))
}
#[allow(clippy::too_many_arguments)]
fn launch_int8_f16_gemv(
&self,
activation: &TensorView,
packed: &TensorView,
scales: &TensorView,
scales_fp16: bool,
zero_points: Option<&TensorView>,
bias: Option<&TensorView>,
output: &mut TensorMut,
k_blocks: usize,
) -> Result<()> {
self.runtime
.require_nvrtc_half_headers("MatMulNBits int8 fp16 GEMV")?;
let large_path_eligible = self.k.is_multiple_of(256)
&& !(self.n <= GEMV_F16_SMALL_N_MAX && self.k <= GEMV_F16_SMALL_N_MAX);
let capabilities = self.runtime.capabilities();
let use_splitk = large_path_eligible
&& if zero_points.is_some() {
true
} else {
int8_symmetric_splitk_enabled()
&& use_f16_symmetric_splitk(
self.k,
self.n,
capabilities.multiprocessor_count(),
capabilities.max_threads_per_block(),
)
};
let entry = if use_splitk {
GEMV_INT8_F16_SPLITK_ENTRY
} else {
GEMV_INT8_F16_ENTRY
};
let function = self
.runtime
.nvrtc_function(GEMV_F16_MODULE, GEMV_F16_SRC, entry)?;
let activation_ptr = cuptr(activation.data_ptr::<u8>() as *const c_void);
let packed_ptr = cuptr(packed.data_ptr::<u8>() as *const c_void);
let scales_ptr = cuptr(scales.data_ptr::<u8>() as *const c_void);
let zero_points_ptr = zero_points
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let bias_ptr = bias
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let output_ptr = cuptr(output.data_ptr_mut::<u8>() as *const c_void);
let k = as_i32("K", self.k)?;
let n = as_i32("N", self.n)?;
let k_blocks = as_i32("K block count", k_blocks)?;
let scales_fp16_flag = scales_fp16 as i32;
let bias_post_round_flag: i32 = (self.fold_bias_post_round && bias.is_some()) as i32;
let threads = if self.n <= GEMV_F16_SMALL_N_MAX && self.k <= GEMV_F16_SMALL_N_MAX {
GEMV_F16_SMALL_THREADS
} else {
GEMV_F16_LARGE_THREADS
};
let columns_per_block = if use_splitk {
(threads / 32) as usize / GEMV_INT8_F16_SPLITK
} else {
(threads / 32) as usize
};
let mut builder = self.runtime.stream().launch_builder(&function);
builder
.arg(&activation_ptr)
.arg(&packed_ptr)
.arg(&scales_ptr)
.arg(&zero_points_ptr)
.arg(&bias_ptr)
.arg(&output_ptr)
.arg(&k)
.arg(&n)
.arg(&k_blocks)
.arg(&scales_fp16_flag)
.arg(&bias_post_round_flag);
unsafe {
builder.launch(LaunchConfig {
grid_dim: (self.n.div_ceil(columns_per_block) as u32, 1, 1),
block_dim: (threads, 1, 1),
shared_mem_bytes: 0,
})
}
.map(|_| ())
.map_err(|err| driver_err("launch MatMulNBits int8 fp16 GEMV", err))
}
#[allow(clippy::too_many_arguments)]
fn launch_int8_f16_gemv_rmsnorm(
&self,
activation: &TensorView,
packed: &TensorView,
scales: &TensorView,
zero_points: Option<&TensorView>,
gamma: &TensorView,
bias: Option<&TensorView>,
output: &mut TensorMut,
k_blocks: usize,
) -> Result<()> {
self.runtime
.require_nvrtc_half_headers("MatMulNBits int8 fp16 RMS-norm-prologue GEMV")?;
if bias.is_some() {
if self.fold_bias_post_round {
onnx_runtime_ep_api::record_kernel_variant_stage!(
"bias",
"qkv_bias_fused",
"folded standalone Add(MatMulNBits, bias) into GEMV epilogue with \
fp16-after-round semantics fp16(fp16(acc)+bias) (token-identity preserved)"
);
} else {
onnx_runtime_ep_api::record_kernel_variant_stage!(
"bias",
"bias_native",
"native MatMulNBits bias: single-round epilogue fp16(acc+bias)"
);
}
}
let entry = if zero_points.is_some() {
GEMV_INT8_F16_SCALES_F16_RMSNORM_ZP_ENTRY
} else {
GEMV_INT8_F16_SCALES_F16_RMSNORM_ENTRY
};
let function = self
.runtime
.nvrtc_function(GEMV_F16_MODULE, GEMV_F16_SRC, entry)?;
let activation_ptr = cuptr(activation.data_ptr::<u8>() as *const c_void);
let packed_ptr = cuptr(packed.data_ptr::<u8>() as *const c_void);
let scales_ptr = cuptr(scales.data_ptr::<u8>() as *const c_void);
let zero_points_ptr = zero_points
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let gamma_ptr = cuptr(gamma.data_ptr::<u8>() as *const c_void);
let bias_ptr = bias
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let output_ptr = cuptr(output.data_ptr_mut::<u8>() as *const c_void);
let k = as_i32("K", self.k)?;
let n = as_i32("N", self.n)?;
let k_blocks = as_i32("K block count", k_blocks)?;
let bias_post_round_flag: i32 = (self.fold_bias_post_round && bias.is_some()) as i32;
let gamma_is_half: i32 = (gamma.dtype == DataType::Float16) as i32;
let epsilon = self.rmsnorm_epsilon;
let threads = if self.n <= GEMV_F16_SMALL_N_MAX && self.k <= GEMV_F16_SMALL_N_MAX {
GEMV_F16_SMALL_THREADS
} else {
GEMV_F16_LARGE_THREADS
};
let columns_per_block = (threads / 32) as usize;
let shared_mem_bytes = (self.k * std::mem::size_of::<half::f16>()) as u32;
let mut builder = self.runtime.stream().launch_builder(&function);
builder
.arg(&activation_ptr)
.arg(&packed_ptr)
.arg(&scales_ptr)
.arg(&zero_points_ptr)
.arg(&gamma_ptr)
.arg(&bias_ptr)
.arg(&output_ptr)
.arg(&k)
.arg(&n)
.arg(&k_blocks)
.arg(&bias_post_round_flag)
.arg(&gamma_is_half)
.arg(&epsilon);
self.runtime
.configure_dynamic_shared_memory(&function, shared_mem_bytes)?;
unsafe {
builder.launch(LaunchConfig {
grid_dim: (self.n.div_ceil(columns_per_block) as u32, 1, 1),
block_dim: (threads, 1, 1),
shared_mem_bytes,
})
}
.map(|_| ())
.map_err(|err| driver_err("launch MatMulNBits int8 fp16 RMS-norm-prologue GEMV", err))
}
fn run_f16_gate_up_swiglu(
&self,
inputs: &[TensorView],
outputs: &mut [TensorMut],
workspace: Option<WorkspaceView>,
) -> Result<()> {
if !(5..=8).contains(&inputs.len()) || outputs.len() != 1 {
return Err(error(format!(
"gate/up SwiGLU fusion expects 5 to 8 inputs [x, W_gate, scales_gate, W_up, \
scales_up, (gamma), (zp_gate, zp_up)] and 1 output, got {} inputs and {} outputs",
inputs.len(),
outputs.len()
)));
}
require_dtype("A", inputs[0].dtype, DataType::Float16)?;
require_dtype("W_gate", inputs[1].dtype, DataType::Uint8)?;
require_dtype("scales_gate", inputs[2].dtype, DataType::Float16)?;
require_dtype("W_up", inputs[3].dtype, DataType::Uint8)?;
require_dtype("scales_up", inputs[4].dtype, DataType::Float16)?;
require_dtype("Y", outputs[0].dtype, DataType::Float16)?;
let gamma = if self.rmsnorm_prologue {
let gamma = optional_input(inputs, 5).ok_or_else(|| {
error("rmsnorm_prologue fusion requires the normalization weight at input 5")
})?;
require_gamma_dtype(gamma.dtype)?;
require_shape("gamma", gamma.shape, &[self.k])?;
if !gamma.is_contiguous() {
return Err(error(
"gamma must be contiguous on the CUDA execution provider".to_string(),
));
}
Some(gamma)
} else {
None
};
let a_shape = inputs[0].shape;
if a_shape.is_empty() || a_shape[a_shape.len() - 1] != self.k {
return Err(error(format!(
"A must have rank >= 1 and last dimension K={}, got {:?}",
self.k, a_shape
)));
}
let m = a_shape[..a_shape.len() - 1].iter().product::<usize>();
let expected_output_shape = [&a_shape[..a_shape.len() - 1], &[self.n]].concat();
if outputs[0].shape != expected_output_shape {
return Err(error(format!(
"Y must have shape {expected_output_shape:?}, got {:?}",
outputs[0].shape
)));
}
if self.block_size != 32 || self.bits != 4 {
return Err(error(format!(
"gate/up SwiGLU fusion received bits={} and block_size={}. Why: the fused fp16 \
path implements the block-32 packed int4 layout. How to fix: export 4-bit \
MatMulNBits weights with block_size=32 or disable this fusion",
self.bits, self.block_size
)));
}
let k_blocks = self.k.div_ceil(self.block_size);
let blob_size = self.block_size / 2;
require_shape("W_gate", inputs[1].shape, &[self.n, k_blocks, blob_size])?;
require_shape("W_up", inputs[3].shape, &[self.n, k_blocks, blob_size])?;
require_flat_or_matrix_shape("scales_gate", inputs[2].shape, self.n, k_blocks)?;
require_flat_or_matrix_shape("scales_up", inputs[4].shape, self.n, k_blocks)?;
let zp_gate = optional_input(inputs, 6);
let zp_up = optional_input(inputs, 7);
if zp_gate.is_some() != zp_up.is_some() {
return Err(error(
"gate/up SwiGLU fusion requires zero points for both projections or neither"
.to_string(),
));
}
let zp_row_bytes = (k_blocks * self.bits).div_ceil(8);
for (name, zp) in [("zp_gate", zp_gate), ("zp_up", zp_up)] {
if let Some(zp) = zp {
require_dtype(name, zp.dtype, DataType::Uint8)?;
require_flat_or_matrix_shape(name, zp.shape, self.n, zp_row_bytes)?;
if !zp.is_contiguous() {
return Err(error(format!(
"{name} must be contiguous on the CUDA execution provider"
)));
}
}
}
for (name, contiguous) in [
("A", inputs[0].is_contiguous()),
("W_gate", inputs[1].is_contiguous()),
("scales_gate", inputs[2].is_contiguous()),
("W_up", inputs[3].is_contiguous()),
("scales_up", inputs[4].is_contiguous()),
("Y", outputs[0].is_contiguous()),
] {
if !contiguous {
return Err(error(format!(
"{name} must be contiguous on the CUDA execution provider"
)));
}
}
crate::trace::record_kernel_metrics(inputs, outputs, || {
let rows = m as u64;
let mut flops = rows
.saturating_mul(self.n as u64)
.saturating_mul(self.k as u64)
.saturating_mul(4)
.saturating_add(rows.saturating_mul(self.n as u64).saturating_mul(5));
if self.rmsnorm_prologue {
let elements = rows.saturating_mul(self.k as u64);
flops = flops
.saturating_add(elements.saturating_mul(4))
.saturating_add(rows.saturating_mul(4));
}
flops
});
if m == 0 {
self.last_call_capture_safe.store(false, Ordering::Relaxed);
onnx_runtime_ep_api::record_kernel_variant!(
"gate_up_swiglu_empty",
"M=0 gate/up SwiGLU has an empty output and requires no CUDA launch"
);
return Ok(());
}
if m > 1 {
if m <= decode_gemv_loop_max_m()
&& decode_gemv_loop_rows_aligned(self.k)
&& let Some(gamma) = gamma
{
onnx_runtime_ep_api::record_kernel_variant!(
"gate_up_swiglu_rmsnorm_batched_loop",
"M={} small-batch rmsnorm decode: {} single-row fused gate/up SwiGLU \
GEMV launches (one per row), each byte-identical to M==1 decode; keeps \
the batch decode graph a single captured subgraph instead of a \
per-MLP-layer eager seam",
m,
m
);
self.last_call_capture_safe.store(true, Ordering::Relaxed);
let a_base = inputs[0].data_ptr::<u8>();
let y_base = outputs[0].data_ptr_mut::<u8>();
let a_row_shape = [1usize, self.k];
let a_row_strides = [self.k as i64, 1];
let y_row_shape = [1usize, self.n];
let y_row_strides = [self.n as i64, 1];
let a_row_bytes = self.k * 2; let y_row_bytes = self.n * 2; for row in 0..m {
let a_row = TensorView::new(
DevicePtr(a_base.wrapping_add(row * a_row_bytes) as *const c_void),
DataType::Float16,
&a_row_shape,
&a_row_strides,
inputs[0].device,
);
let mut y_row = TensorMut::new(
DevicePtrMut(y_base.wrapping_add(row * y_row_bytes) as *mut c_void),
DataType::Float16,
&y_row_shape,
&y_row_strides,
outputs[0].device,
);
self.launch_gate_up_swiglu_rmsnorm(
&a_row,
&inputs[1],
&inputs[2],
&inputs[3],
&inputs[4],
zp_gate,
zp_up,
gamma,
&mut y_row,
k_blocks,
blob_size,
zp_row_bytes,
)?;
}
return Ok(());
}
self.last_call_capture_safe.store(false, Ordering::Relaxed);
if dequant_f16_gemm_enabled() && m >= dequant_f16_gemm_min_m() {
match self.try_dequant_f16_gate_up_prefill(
&inputs[0],
&inputs[1],
&inputs[2],
&inputs[3],
&inputs[4],
inputs[2].dtype == DataType::Float16,
zp_gate,
zp_up,
gamma,
&mut outputs[0],
m,
k_blocks,
blob_size,
zp_row_bytes,
workspace,
) {
Ok(true) => {
onnx_runtime_ep_api::record_kernel_variant!(
"gate_up_swiglu_dequant_f16_cublas",
"M={} prefill: dequantize both int{} projections to [N, K] fp16, \
two cuBLASLt fp16 tensor-core GEMMs, then the fp16 SiluMul \
epilogue",
m,
self.bits
);
return Ok(());
}
Ok(false) => {
if std::env::var_os("ONNX_GENAI_DEQUANT_F16_DEBUG").is_some() {
eprintln!("dequant_f16 gate/up: ineligible m={m}");
}
}
Err(_err) => {
if std::env::var_os("ONNX_GENAI_DEQUANT_F16_DEBUG").is_some() {
eprintln!("dequant_f16 gate/up: {_err}");
}
}
}
}
if marlin_gemm::marlin_m_gt_1_enabled()
&& self.bits == 4
&& self.marlin_weight_inputs_constant()
&& marlin_gemm::device_supports_marlin(
self.runtime.capabilities().compute_capability(),
)
{
match self.try_launch_marlin_gate_up_prefill(
&inputs[0],
&inputs[1],
&inputs[2],
&inputs[3],
&inputs[4],
zp_gate,
zp_up,
gamma,
&mut outputs[0],
m,
) {
Ok(Some(warm)) => {
self.last_call_capture_safe.store(warm, Ordering::Relaxed);
onnx_runtime_ep_api::record_kernel_variant!(
"gate_up_swiglu_marlin_prefill",
"M={} prefill/verify: {}paired gate/up Marlin SM80 mma.sync int4 \
tensor-core GEMMs followed by fp16 SiluMul (capture-safe when \
weights + scratch are pre-warmed)",
m,
if gamma.is_some() {
"RMS-normalization prologue then "
} else {
""
}
);
return Ok(());
}
Ok(None) => {
}
Err(_err) => {
}
}
}
if let Some(gamma) = gamma {
onnx_runtime_ep_api::record_kernel_variant!(
"gate_up_swiglu_rmsnorm_prefill",
"M={} prefill: RMS-normalization prologue into scratch, then two portable \
block-32 int4 fp16 tiled GEMMs with fp32 accumulation, followed by fp16 \
SiluMul; not advertised as CUDA-graph capture-safe",
m
);
return self.launch_gate_up_swiglu_rmsnorm_prefill(
&inputs[0],
&inputs[1],
&inputs[2],
&inputs[3],
&inputs[4],
zp_gate,
zp_up,
gamma,
&mut outputs[0],
m,
k_blocks,
zp_row_bytes,
);
}
onnx_runtime_ep_api::record_kernel_variant!(
"gate_up_swiglu_prefill",
"M={} prefill: two portable block-32 int4 fp16 tiled GEMMs with fp32 \
accumulation, followed by fp16 SiluMul; not advertised as CUDA-graph \
capture-safe",
m
);
return self.launch_gate_up_swiglu_prefill(
&inputs[0],
&inputs[1],
&inputs[2],
&inputs[3],
&inputs[4],
zp_gate,
zp_up,
&mut outputs[0],
m,
k_blocks,
);
}
if let Some(gamma) = gamma {
onnx_runtime_ep_api::record_kernel_variant!(
"gate_up_swiglu_rmsnorm_fused",
"fp16 block-32 M==1 decode: fused RMS-normalization prologue + paired gate/up \
int4 GEMV + SwiGLU (silu(gate)*up) in one capture-safe kernel; the RMS \
reduction runs once for both projections and reproduces the standalone norm \
plus two-op fp16 rounding for byte-identical greedy tokens"
);
self.last_call_capture_safe.store(true, Ordering::Relaxed);
return self.launch_gate_up_swiglu_rmsnorm(
&inputs[0],
&inputs[1],
&inputs[2],
&inputs[3],
&inputs[4],
zp_gate,
zp_up,
gamma,
&mut outputs[0],
k_blocks,
blob_size,
zp_row_bytes,
);
}
onnx_runtime_ep_api::record_kernel_variant!(
"gate_up_swiglu_fused",
"fp16 block-32 M==1 decode: fused paired gate/up int4 GEMV + SwiGLU \
(silu(gate)*up) in one capture-safe kernel; reproduces the two-op fp16 \
rounding for byte-identical greedy tokens"
);
self.last_call_capture_safe.store(true, Ordering::Relaxed);
self.launch_gate_up_swiglu(
&inputs[0],
&inputs[1],
&inputs[2],
&inputs[3],
&inputs[4],
zp_gate,
zp_up,
&mut outputs[0],
k_blocks,
blob_size,
zp_row_bytes,
)
}
#[allow(clippy::too_many_arguments)]
fn launch_gate_up_swiglu_prefill(
&self,
activation: &TensorView,
packed_gate: &TensorView,
scales_gate: &TensorView,
packed_up: &TensorView,
scales_up: &TensorView,
zp_gate: Option<&TensorView>,
zp_up: Option<&TensorView>,
output: &mut TensorMut,
m: usize,
k_blocks: usize,
) -> Result<()> {
let scratch = self.runtime.alloc_raw(output.byte_size())?;
let scratch_shape = output.shape.to_vec();
let scratch_strides = output.strides.to_vec();
let mut gate_output = TensorMut::new(
DevicePtrMut(raw_ptr(scratch)),
DataType::Float16,
&scratch_shape,
&scratch_strides,
output.device,
);
let result = (|| {
self.launch_f16_gemm(
activation,
packed_gate,
scales_gate,
true,
zp_gate,
None,
&mut gate_output,
m,
k_blocks,
self.block_size * self.bits / 8,
0,
)?;
self.launch_f16_gemm(
activation,
packed_up,
scales_up,
true,
zp_up,
None,
output,
m,
k_blocks,
self.block_size * self.bits / 8,
0,
)?;
let output_ptr = cuptr(output.data_ptr_mut::<u8>() as *const c_void);
crate::kernels::elementwise::launch_silu_mul_f16_raw(
&self.runtime,
scratch,
output_ptr,
output_ptr,
output.numel(),
self.decomposed_silu,
)
})();
let free_scratch = unsafe { self.runtime.free_raw(scratch) };
result.and(free_scratch)
}
#[allow(clippy::too_many_arguments)]
fn launch_gate_up_swiglu(
&self,
activation: &TensorView,
packed_gate: &TensorView,
scales_gate: &TensorView,
packed_up: &TensorView,
scales_up: &TensorView,
zp_gate: Option<&TensorView>,
zp_up: Option<&TensorView>,
output: &mut TensorMut,
k_blocks: usize,
blob_size: usize,
zp_row_bytes: usize,
) -> Result<()> {
self.runtime
.require_nvrtc_half_headers("MatMulNBits fp16 gate/up SwiGLU GEMV")?;
let has_zp = zp_gate.is_some() || zp_up.is_some();
let vec = !has_zp && gate_up_vec_enabled();
let entry = match (self.decomposed_silu, has_zp, vec) {
(true, true, _) => GATE_UP_DECOMPOSED_SWIGLU_ZP_ENTRY,
(true, false, true) => GATE_UP_DECOMPOSED_SWIGLU_VEC_ENTRY,
(true, false, false) => GATE_UP_DECOMPOSED_SWIGLU_ENTRY,
(false, true, _) => GATE_UP_SWIGLU_ZP_ENTRY,
(false, false, true) => GATE_UP_SWIGLU_VEC_ENTRY,
(false, false, false) => GATE_UP_SWIGLU_ENTRY,
};
let function = self
.runtime
.nvrtc_function(GEMV_F16_MODULE, GEMV_F16_SRC, entry)?;
let activation_ptr = cuptr(activation.data_ptr::<u8>() as *const c_void);
let packed_gate_ptr = cuptr(packed_gate.data_ptr::<u8>() as *const c_void);
let scales_gate_ptr = cuptr(scales_gate.data_ptr::<u8>() as *const c_void);
let packed_up_ptr = cuptr(packed_up.data_ptr::<u8>() as *const c_void);
let scales_up_ptr = cuptr(scales_up.data_ptr::<u8>() as *const c_void);
let zp_gate_ptr = zp_gate
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let zp_up_ptr = zp_up
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let output_ptr = cuptr(output.data_ptr_mut::<u8>() as *const c_void);
let k = as_i32("K", self.k)?;
let n = as_i32("N", self.n)?;
let k_blocks = as_i32("K block count", k_blocks)?;
let blob_size = as_i32("block blob size", blob_size)?;
let zp_row_bytes = as_i32("zero-point row byte count", zp_row_bytes)?;
let threads = GATE_UP_SWIGLU_THREADS;
let columns_per_block = (threads / 32) as usize;
let mut builder = self.runtime.stream().launch_builder(&function);
builder
.arg(&activation_ptr)
.arg(&packed_gate_ptr)
.arg(&scales_gate_ptr)
.arg(&packed_up_ptr)
.arg(&scales_up_ptr)
.arg(&zp_gate_ptr)
.arg(&zp_up_ptr)
.arg(&output_ptr)
.arg(&k)
.arg(&n)
.arg(&k_blocks)
.arg(&blob_size)
.arg(&zp_row_bytes);
unsafe {
builder.launch(LaunchConfig {
grid_dim: (self.n.div_ceil(columns_per_block) as u32, 1, 1),
block_dim: (threads, 1, 1),
shared_mem_bytes: 0,
})
}
.map(|_| ())
.map_err(|err| driver_err("launch MatMulNBits fp16 gate/up SwiGLU GEMV", err))
}
#[allow(clippy::too_many_arguments)]
fn launch_gate_up_swiglu_rmsnorm(
&self,
activation: &TensorView,
packed_gate: &TensorView,
scales_gate: &TensorView,
packed_up: &TensorView,
scales_up: &TensorView,
zp_gate: Option<&TensorView>,
zp_up: Option<&TensorView>,
gamma: &TensorView,
output: &mut TensorMut,
k_blocks: usize,
blob_size: usize,
zp_row_bytes: usize,
) -> Result<()> {
self.runtime
.require_nvrtc_half_headers("MatMulNBits fp16 gate/up SwiGLU RMS-norm GEMV")?;
let has_zp = zp_gate.is_some() || zp_up.is_some();
let vec = !has_zp && gate_up_vec_enabled();
let occ = !has_zp && gate_up_occ_enabled();
let entry = match (self.decomposed_silu, has_zp, occ, vec) {
(true, true, _, _) => GATE_UP_DECOMPOSED_SWIGLU_RMSNORM_ZP_ENTRY,
(true, false, true, _) => GATE_UP_DECOMPOSED_SWIGLU_RMSNORM_VEC_OCC_ENTRY,
(true, false, false, true) => GATE_UP_DECOMPOSED_SWIGLU_RMSNORM_VEC_ENTRY,
(true, false, false, false) => GATE_UP_DECOMPOSED_SWIGLU_RMSNORM_ENTRY,
(false, true, _, _) => GATE_UP_SWIGLU_RMSNORM_ZP_ENTRY,
(false, false, true, _) => GATE_UP_SWIGLU_RMSNORM_VEC_OCC_ENTRY,
(false, false, false, true) => GATE_UP_SWIGLU_RMSNORM_VEC_ENTRY,
(false, false, false, false) => GATE_UP_SWIGLU_RMSNORM_ENTRY,
};
let function = self
.runtime
.nvrtc_function(GEMV_F16_MODULE, GEMV_F16_SRC, entry)?;
let activation_ptr = cuptr(activation.data_ptr::<u8>() as *const c_void);
let packed_gate_ptr = cuptr(packed_gate.data_ptr::<u8>() as *const c_void);
let scales_gate_ptr = cuptr(scales_gate.data_ptr::<u8>() as *const c_void);
let packed_up_ptr = cuptr(packed_up.data_ptr::<u8>() as *const c_void);
let scales_up_ptr = cuptr(scales_up.data_ptr::<u8>() as *const c_void);
let zp_gate_ptr = zp_gate
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let zp_up_ptr = zp_up
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let gamma_ptr = cuptr(gamma.data_ptr::<u8>() as *const c_void);
let output_ptr = cuptr(output.data_ptr_mut::<u8>() as *const c_void);
let k = as_i32("K", self.k)?;
let n = as_i32("N", self.n)?;
let k_blocks = as_i32("K block count", k_blocks)?;
let blob_size = as_i32("block blob size", blob_size)?;
let zp_row_bytes = as_i32("zero-point row byte count", zp_row_bytes)?;
let gamma_is_half: i32 = (gamma.dtype == DataType::Float16) as i32;
let epsilon = self.rmsnorm_epsilon;
let threads = GATE_UP_SWIGLU_THREADS;
let columns_per_block = (threads / 32) as usize;
let shared_mem_bytes = (self.k * std::mem::size_of::<half::f16>()) as u32;
let mut builder = self.runtime.stream().launch_builder(&function);
builder
.arg(&activation_ptr)
.arg(&packed_gate_ptr)
.arg(&scales_gate_ptr)
.arg(&packed_up_ptr)
.arg(&scales_up_ptr)
.arg(&zp_gate_ptr)
.arg(&zp_up_ptr)
.arg(&gamma_ptr)
.arg(&output_ptr)
.arg(&k)
.arg(&n)
.arg(&k_blocks)
.arg(&blob_size)
.arg(&zp_row_bytes)
.arg(&gamma_is_half)
.arg(&epsilon);
self.runtime
.configure_dynamic_shared_memory(&function, shared_mem_bytes)?;
unsafe {
builder.launch(LaunchConfig {
grid_dim: (self.n.div_ceil(columns_per_block) as u32, 1, 1),
block_dim: (threads, 1, 1),
shared_mem_bytes,
})
}
.map(|_| ())
.map_err(|err| driver_err("launch MatMulNBits fp16 gate/up SwiGLU RMS-norm GEMV", err))
}
#[allow(clippy::too_many_arguments)]
fn launch_gate_up_swiglu_rmsnorm_prefill(
&self,
activation: &TensorView,
packed_gate: &TensorView,
scales_gate: &TensorView,
packed_up: &TensorView,
scales_up: &TensorView,
zp_gate: Option<&TensorView>,
zp_up: Option<&TensorView>,
gamma: &TensorView,
output: &mut TensorMut,
m: usize,
k_blocks: usize,
_zp_row_bytes: usize,
) -> Result<()> {
let scratch = self
.runtime
.alloc_raw(m * self.k * std::mem::size_of::<half::f16>())?;
let scratch_shape = [m, self.k];
let scratch_strides = [self.k as i64, 1];
let normalized = TensorView::new(
DevicePtr(raw_ptr(scratch) as *const c_void),
DataType::Float16,
&scratch_shape,
&scratch_strides,
activation.device,
);
let result = self
.launch_rmsnorm_prefill(activation, gamma, scratch, m)
.and_then(|()| {
self.launch_gate_up_swiglu_prefill(
&normalized,
packed_gate,
scales_gate,
packed_up,
scales_up,
zp_gate,
zp_up,
output,
m,
k_blocks,
)
});
let free_scratch = unsafe { self.runtime.free_raw(scratch) };
result.and(free_scratch)
}
#[allow(clippy::too_many_arguments)]
fn launch_f16_gemv(
&self,
activation: &TensorView,
packed: &TensorView,
scales: &TensorView,
scales_fp16: bool,
zero_points: Option<&TensorView>,
bias: Option<&TensorView>,
output: &mut TensorMut,
k_blocks: usize,
blob_size: usize,
zp_row_bytes: usize,
) -> Result<()> {
let selection = select_f16_gemv_variant(
self.k,
self.n,
self.block_size,
scales_fp16,
zero_points.is_some(),
);
let variant_name = if self.block_size != 32 {
"gemv_f16_general_bs"
} else {
match selection.variant {
F16GemvVariant::DownProjection => "gemv_f16_down_projection",
F16GemvVariant::General => "gemv_f16_general",
}
};
onnx_runtime_ep_api::record_kernel_variant!(
variant_name,
"fp16-activation x int{} M==1 decode GEMV: block_size={}; zero_points={}; {}",
self.bits,
self.block_size,
zero_points.is_some(),
selection.reason
);
if bias.is_some() {
if self.fold_bias_post_round {
onnx_runtime_ep_api::record_kernel_variant_stage!(
"bias",
"qkv_bias_fused",
"folded standalone Add(MatMulNBits, bias) into GEMV epilogue with \
fp16-after-round semantics fp16(fp16(acc)+bias) (token-identity preserved)"
);
} else {
onnx_runtime_ep_api::record_kernel_variant_stage!(
"bias",
"bias_native",
"native MatMulNBits bias: single-round epilogue fp16(acc+bias)"
);
}
}
self.launch_f16_gemv_variant(
activation,
packed,
scales,
scales_fp16,
zero_points,
bias,
output,
k_blocks,
blob_size,
zp_row_bytes,
selection,
)
}
#[allow(clippy::too_many_arguments)]
fn launch_f16_gemv_variant(
&self,
activation: &TensorView,
packed: &TensorView,
scales: &TensorView,
scales_fp16: bool,
zero_points: Option<&TensorView>,
bias: Option<&TensorView>,
output: &mut TensorMut,
k_blocks: usize,
blob_size: usize,
zp_row_bytes: usize,
selection: F16GemvSelection,
) -> Result<()> {
self.runtime
.require_nvrtc_half_headers("MatMulNBits fp16 GEMV")?;
let capabilities = self.runtime.capabilities();
let use_scales_f16_zp_splitk = self.block_size == 32
&& scales_fp16
&& zero_points.is_some()
&& matches!(selection.variant, F16GemvVariant::General)
&& self.k.is_multiple_of(256)
&& !(self.n <= GEMV_F16_SMALL_N_MAX && self.k <= GEMV_F16_SMALL_N_MAX)
&& use_scales_f16_zp_splitk_gate(
self.k,
self.n,
capabilities.multiprocessor_count(),
);
let scales_f16_zp_factor =
scales_f16_zp_split_factor(self.n, capabilities.multiprocessor_count());
let use_scales_f16_symmetric_splitk = self.block_size == 32
&& scales_fp16
&& zero_points.is_none()
&& matches!(selection.variant, F16GemvVariant::General)
&& use_f16_symmetric_splitk(
self.k,
self.n,
capabilities.multiprocessor_count(),
capabilities.max_threads_per_block(),
);
let use_scales_f16_splitk = use_scales_f16_zp_splitk || use_scales_f16_symmetric_splitk;
let use_scales_f16_pipeline = self.block_size == 32
&& scales_fp16
&& matches!(selection.variant, F16GemvVariant::General)
&& !use_scales_f16_splitk
&& scales_f16_pipeline_enabled();
let pipe_columns_per_block =
if self.n <= GEMV_F16_SMALL_N_MAX && self.k <= GEMV_F16_SMALL_N_MAX {
(GEMV_F16_SMALL_THREADS / 32) as usize
} else {
(GEMV_F16_LARGE_THREADS / 32) as usize
};
let use_scales_f16_pipe = use_scales_f16_pipeline
&& !scales_f16_pipe_well_occupied(
self.n,
pipe_columns_per_block,
self.runtime.capabilities().multiprocessor_count(),
);
let down_choice = down_columns_override().unwrap_or_else(|| {
select_down_columns(self.n, self.runtime.capabilities().multiprocessor_count())
});
let use_general_splitk = self.block_size != 32
&& use_general_bs_splitk(
self.k,
self.n,
self.bits,
self.runtime.capabilities().multiprocessor_count(),
);
let entry = if self.block_size != 32 {
if use_general_splitk {
if use_gemv_wideload(self.bits, self.block_size, self.k) {
if use_gemv_splitk_multicol()
&& !splitk_smalln_prefers_single_column(
self.n,
self.runtime.capabilities().multiprocessor_count(),
)
{
GEMV_F16_GENERAL_BS_SPLITK_WIDE_MULTICOL_ENTRY
} else {
GEMV_F16_GENERAL_BS_SPLITK_WIDE_ENTRY
}
} else {
GEMV_F16_GENERAL_BS_SPLITK_ENTRY
}
} else if use_gemv_wideload(self.bits, self.block_size, self.k) {
if use_gemv_wide_multicol() {
if use_gemv_fp16() {
GEMV_F16_GENERAL_BS_WIDE_MULTICOL_FP16_ENTRY
} else {
GEMV_F16_GENERAL_BS_WIDE_MULTICOL_ENTRY
}
} else {
GEMV_F16_GENERAL_BS_WIDE_ENTRY
}
} else {
GEMV_F16_GENERAL_BS_ENTRY
}
} else {
match selection.variant {
F16GemvVariant::DownProjection => down_choice.1,
F16GemvVariant::General if scales_fp16 => {
if zero_points.is_some() {
if use_scales_f16_zp_splitk {
if scales_f16_zp_factor == GEMV_F16_SCALES_F16_ZP_SPLITK8 {
GEMV_F16_SCALES_F16_ZP_SPLITK8_ENTRY
} else {
GEMV_F16_SCALES_F16_ZP_SPLITK_ENTRY
}
} else if use_scales_f16_pipe {
GEMV_F16_SCALES_F16_ZP_PIPE_ENTRY
} else {
GEMV_F16_SCALES_F16_ZP_ENTRY
}
} else {
if use_scales_f16_symmetric_splitk {
GEMV_F16_SCALES_F16_SPLITK_ENTRY
} else if use_scales_f16_pipe {
GEMV_F16_SCALES_F16_PIPE_ENTRY
} else {
GEMV_F16_SCALES_F16_ENTRY
}
}
}
F16GemvVariant::General => GEMV_F16_ENTRY,
}
};
let interleaved_entry = if entry == GEMV_F16_GENERAL_BS_WIDE_MULTICOL_ENTRY {
Some(GEMV_F16_GENERAL_BS_WIDE_MULTICOL_INTERLEAVED_ENTRY)
} else if entry == GEMV_F16_GENERAL_BS_SPLITK_WIDE_MULTICOL_ENTRY {
Some(GEMV_F16_GENERAL_BS_SPLITK_WIDE_MULTICOL_INTERLEAVED_ENTRY)
} else if entry == GEMV_F16_GENERAL_BS_SPLITK_WIDE_ENTRY {
Some(GEMV_F16_GENERAL_BS_SPLITK_WIDE_INTERLEAVED_ENTRY)
} else {
None
};
let orig_packed_ptr = cuptr(packed.data_ptr::<u8>() as *const c_void);
let interleave_on = interleave_dequant_enabled() && self.bits == 4 && zero_points.is_none();
let (entry, packed_ptr) = if interleave_on && let Some(target) = interleaved_entry {
let bytes = self.n.saturating_mul(k_blocks).saturating_mul(blob_size);
match ensure_interleaved(&self.runtime, orig_packed_ptr, bytes) {
Ok((iptr, warm)) => {
if !warm {
self.last_call_capture_safe.store(false, Ordering::Relaxed);
}
onnx_runtime_ep_api::record_kernel_variant_stage!(
"dequant",
"interleaved_biased",
"TRT-LLM interleaved+biased int4 dequant: offline nibble-interleave + \
folded symmetric -8 bias drops the per-block sub.f16x2 and the prmt.b32 \
activation reorder (byte-identical to the fp32 wide/multicol/split-K path)"
);
(target, iptr)
}
Err(_) => (entry, orig_packed_ptr),
}
} else {
(entry, orig_packed_ptr)
};
let entry = if use_scales_f16_zp_splitk_pf() {
if entry == GEMV_F16_SCALES_F16_ZP_SPLITK_ENTRY {
GEMV_F16_SCALES_F16_ZP_SPLITK_PF_ENTRY
} else if entry == GEMV_F16_SCALES_F16_SPLITK_ENTRY {
GEMV_F16_SCALES_F16_SPLITK_PF_ENTRY
} else if entry == GEMV_F16_SCALES_F16_ZP_SPLITK8_ENTRY {
GEMV_F16_SCALES_F16_ZP_SPLITK8_PF_ENTRY
} else {
entry
}
} else {
entry
};
let function = self
.runtime
.nvrtc_function(GEMV_F16_MODULE, GEMV_F16_SRC, entry)?;
let activation_ptr = cuptr(activation.data_ptr::<u8>() as *const c_void);
let scales_ptr = cuptr(scales.data_ptr::<u8>() as *const c_void);
let zero_points_ptr = zero_points
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let bias_ptr = bias
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let bf16_direct_capable = entry == GEMV_F16_SCALES_F16_SPLITK_ENTRY
|| entry == GEMV_F16_SCALES_F16_ZP_SPLITK_ENTRY
|| entry == GEMV_F16_SCALES_F16_SPLITK_PF_ENTRY
|| entry == GEMV_F16_SCALES_F16_ZP_SPLITK_PF_ENTRY
|| entry == GEMV_F16_SCALES_F16_ZP_SPLITK8_ENTRY
|| entry == GEMV_F16_SCALES_F16_ZP_SPLITK8_PF_ENTRY;
let staging_ptr = cuptr(output.data_ptr_mut::<u8>() as *const c_void);
let direct_bf16_ptr = bf16_direct_capable
.then(|| take_bf16_direct_out(staging_ptr, self.n))
.flatten();
let output_ptr = direct_bf16_ptr.unwrap_or(staging_ptr);
let out_bf16_flag: i32 = direct_bf16_ptr.is_some() as i32;
let k = as_i32("K", self.k)?;
let n = as_i32("N", self.n)?;
let block_size = as_i32("block_size", self.block_size)?;
let k_blocks = as_i32("K block count", k_blocks)?;
let blob_size = as_i32("block blob size", blob_size)?;
let zp_row_bytes = as_i32("zero-point row byte count", zp_row_bytes)?;
let scales_fp16_flag: i32 = scales_fp16 as i32;
let bias_post_round_flag: i32 = (self.fold_bias_post_round && bias.is_some()) as i32;
let bits = as_i32("bits", self.bits)?;
let (threads, columns_per_block, shared_mem_bytes) = match selection.variant {
F16GemvVariant::DownProjection => (GEMV_F16_DOWN_THREADS, down_choice.0, 0),
F16GemvVariant::General => {
let threads = if self.n <= GEMV_F16_SMALL_N_MAX
&& self.k <= GEMV_F16_SMALL_N_MAX
&& !use_general_splitk
{
GEMV_F16_SMALL_THREADS
} else {
GEMV_F16_LARGE_THREADS
};
let columns_per_block = if use_general_splitk {
if entry == GEMV_F16_GENERAL_BS_SPLITK_WIDE_MULTICOL_ENTRY
|| entry == GEMV_F16_GENERAL_BS_SPLITK_WIDE_MULTICOL_INTERLEAVED_ENTRY
{
(threads / 32) as usize / GENERAL_BS_SPLITK_MULTICOL
* GEMV_F16_WIDE_MULTICOL_NC
} else {
(threads / 32) as usize / GENERAL_BS_SPLITK
}
} else if use_scales_f16_splitk {
let factor = if use_scales_f16_zp_splitk {
scales_f16_zp_factor
} else {
GEMV_F16_SCALES_F16_ZP_SPLITK
};
(threads / 32) as usize / factor
} else if entry == GEMV_F16_GENERAL_BS_WIDE_MULTICOL_ENTRY
|| entry == GEMV_F16_GENERAL_BS_WIDE_MULTICOL_FP16_ENTRY
|| entry == GEMV_F16_GENERAL_BS_WIDE_MULTICOL_INTERLEAVED_ENTRY
{
(threads / 32) as usize * GEMV_F16_WIDE_MULTICOL_NC
} else {
(threads / 32) as usize
};
(threads, columns_per_block, 0)
}
};
let mut builder = self.runtime.stream().launch_builder(&function);
builder
.arg(&activation_ptr)
.arg(&packed_ptr)
.arg(&scales_ptr)
.arg(&zero_points_ptr)
.arg(&bias_ptr)
.arg(&output_ptr)
.arg(&k)
.arg(&n)
.arg(&block_size)
.arg(&k_blocks)
.arg(&blob_size)
.arg(&zp_row_bytes)
.arg(&scales_fp16_flag)
.arg(&bias_post_round_flag);
if entry == GEMV_F16_GENERAL_BS_ENTRY
|| entry == GEMV_F16_GENERAL_BS_SPLITK_ENTRY
|| entry == GEMV_F16_GENERAL_BS_WIDE_ENTRY
|| entry == GEMV_F16_GENERAL_BS_SPLITK_WIDE_ENTRY
|| entry == GEMV_F16_GENERAL_BS_WIDE_MULTICOL_ENTRY
|| entry == GEMV_F16_GENERAL_BS_WIDE_MULTICOL_FP16_ENTRY
|| entry == GEMV_F16_GENERAL_BS_WIDE_MULTICOL_INTERLEAVED_ENTRY
|| entry == GEMV_F16_GENERAL_BS_SPLITK_WIDE_INTERLEAVED_ENTRY
|| entry == GEMV_F16_GENERAL_BS_SPLITK_WIDE_MULTICOL_ENTRY
|| entry == GEMV_F16_GENERAL_BS_SPLITK_WIDE_MULTICOL_INTERLEAVED_ENTRY
{
builder.arg(&bits);
}
if bf16_direct_capable {
builder.arg(&out_bf16_flag);
}
unsafe {
builder.launch(LaunchConfig {
grid_dim: (self.n.div_ceil(columns_per_block) as u32, 1, 1),
block_dim: (threads, 1, 1),
shared_mem_bytes,
})
}
.map(|_| ())
.map_err(|err| {
driver_err(
&format!("launch MatMulNBits fp16 GEMV ({})", selection.reason),
err,
)
})
}
#[allow(clippy::too_many_arguments)]
fn launch_f16_gemv_rmsnorm(
&self,
activation: &TensorView,
packed: &TensorView,
scales: &TensorView,
zero_points: Option<&TensorView>,
gamma: &TensorView,
bias: Option<&TensorView>,
output: &mut TensorMut,
k_blocks: usize,
blob_size: usize,
zp_row_bytes: usize,
) -> Result<()> {
self.runtime
.require_nvrtc_half_headers("MatMulNBits fp16 RMS-norm-prologue GEMV")?;
if bias.is_some() {
if self.fold_bias_post_round {
onnx_runtime_ep_api::record_kernel_variant_stage!(
"bias",
"qkv_bias_fused",
"folded standalone Add(MatMulNBits, bias) into GEMV epilogue with \
fp16-after-round semantics fp16(fp16(acc)+bias) (token-identity preserved)"
);
} else {
onnx_runtime_ep_api::record_kernel_variant_stage!(
"bias",
"bias_native",
"native MatMulNBits bias: single-round epilogue fp16(acc+bias)"
);
}
}
let capabilities = self.runtime.capabilities();
let use_splitk = zero_points.is_none()
&& use_f16_symmetric_splitk(
self.k,
self.n,
capabilities.multiprocessor_count(),
capabilities.max_threads_per_block(),
);
let entry = if zero_points.is_some() {
GEMV_F16_SCALES_F16_RMSNORM_ZP_ENTRY
} else if use_splitk {
GEMV_F16_SCALES_F16_RMSNORM_SPLITK_ENTRY
} else {
GEMV_F16_SCALES_F16_RMSNORM_ENTRY
};
let function = self
.runtime
.nvrtc_function(GEMV_F16_MODULE, GEMV_F16_SRC, entry)?;
let activation_ptr = cuptr(activation.data_ptr::<u8>() as *const c_void);
let packed_ptr = cuptr(packed.data_ptr::<u8>() as *const c_void);
let scales_ptr = cuptr(scales.data_ptr::<u8>() as *const c_void);
let zero_points_ptr = zero_points
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let gamma_ptr = cuptr(gamma.data_ptr::<u8>() as *const c_void);
let bias_ptr = bias
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let output_ptr = cuptr(output.data_ptr_mut::<u8>() as *const c_void);
let k = as_i32("K", self.k)?;
let n = as_i32("N", self.n)?;
let k_blocks = as_i32("K block count", k_blocks)?;
let blob_size = as_i32("block blob size", blob_size)?;
let zp_row_bytes = as_i32("zero-point row byte count", zp_row_bytes)?;
let bias_post_round_flag: i32 = (self.fold_bias_post_round && bias.is_some()) as i32;
let gamma_is_half: i32 = (gamma.dtype == DataType::Float16) as i32;
let epsilon = self.rmsnorm_epsilon;
let threads = if self.n <= GEMV_F16_SMALL_N_MAX && self.k <= GEMV_F16_SMALL_N_MAX {
GEMV_F16_SMALL_THREADS
} else {
GEMV_F16_LARGE_THREADS
};
let columns_per_block = (threads / 32) as usize
/ if use_splitk {
GEMV_F16_SCALES_F16_ZP_SPLITK
} else {
1
};
let shared_mem_bytes = (self.k * std::mem::size_of::<half::f16>()) as u32;
let mut builder = self.runtime.stream().launch_builder(&function);
builder
.arg(&activation_ptr)
.arg(&packed_ptr)
.arg(&scales_ptr)
.arg(&zero_points_ptr)
.arg(&gamma_ptr)
.arg(&bias_ptr)
.arg(&output_ptr)
.arg(&k)
.arg(&n)
.arg(&k_blocks)
.arg(&blob_size)
.arg(&zp_row_bytes)
.arg(&bias_post_round_flag)
.arg(&gamma_is_half)
.arg(&epsilon);
self.runtime
.configure_dynamic_shared_memory(&function, shared_mem_bytes)?;
unsafe {
builder.launch(LaunchConfig {
grid_dim: (self.n.div_ceil(columns_per_block) as u32, 1, 1),
block_dim: (threads, 1, 1),
shared_mem_bytes,
})
}
.map(|_| ())
.map_err(|err| driver_err("launch MatMulNBits fp16 RMS-norm-prologue GEMV", err))
}
#[allow(clippy::too_many_arguments)]
fn launch_int8_f32_gemv(
&self,
activation: &TensorView,
packed: &TensorView,
scales: &TensorView,
zero_points: Option<&TensorView>,
bias: Option<&TensorView>,
output: &mut TensorMut,
k_blocks: usize,
) -> Result<()> {
let function = self
.runtime
.nvrtc_function(GEMV_MODULE, GEMV_SRC, GEMV_INT8_F32_ENTRY)?;
let activation_ptr = cuptr(activation.data_ptr::<u8>() as *const c_void);
let packed_ptr = cuptr(packed.data_ptr::<u8>() as *const c_void);
let scales_ptr = cuptr(scales.data_ptr::<u8>() as *const c_void);
let zero_points_ptr = zero_points
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let bias_ptr = bias
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let output_ptr = cuptr(output.data_ptr_mut::<u8>() as *const c_void);
let k = as_i32("K", self.k)?;
let n = as_i32("N", self.n)?;
let k_blocks = as_i32("K block count", k_blocks)?;
let mut builder = self.runtime.stream().launch_builder(&function);
builder
.arg(&activation_ptr)
.arg(&packed_ptr)
.arg(&scales_ptr)
.arg(&zero_points_ptr)
.arg(&bias_ptr)
.arg(&output_ptr)
.arg(&k)
.arg(&n)
.arg(&k_blocks);
unsafe {
builder.launch(LaunchConfig {
grid_dim: (self.n as u32, 1, 1),
block_dim: (BLOCK_THREADS, 1, 1),
shared_mem_bytes: 0,
})
}
.map(|_| ())
.map_err(|err| driver_err("launch MatMulNBits int8 f32 GEMV", err))
}
#[allow(clippy::too_many_arguments)]
fn launch_int4_f32_gemv_block128(
&self,
activation: &TensorView,
packed: &TensorView,
scales: &TensorView,
zero_points: &TensorView,
bias: Option<&TensorView>,
output: &mut TensorMut,
k_blocks: usize,
) -> Result<()> {
let function =
self.runtime
.nvrtc_function(GEMV_MODULE, GEMV_SRC, GEMV_INT4_F32_BLOCK128_ENTRY)?;
let activation_ptr = cuptr(activation.data_ptr::<u8>() as *const c_void);
let packed_ptr = cuptr(packed.data_ptr::<u8>() as *const c_void);
let scales_ptr = cuptr(scales.data_ptr::<u8>() as *const c_void);
let zero_points_ptr = cuptr(zero_points.data_ptr::<u8>() as *const c_void);
let bias_ptr = bias
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let output_ptr = cuptr(output.data_ptr_mut::<u8>() as *const c_void);
let k = as_i32("K", self.k)?;
let n = as_i32("N", self.n)?;
let k_blocks = as_i32("K block count", k_blocks)?;
let mut builder = self.runtime.stream().launch_builder(&function);
builder
.arg(&activation_ptr)
.arg(&packed_ptr)
.arg(&scales_ptr)
.arg(&zero_points_ptr)
.arg(&bias_ptr)
.arg(&output_ptr)
.arg(&k)
.arg(&n)
.arg(&k_blocks);
unsafe {
builder.launch(LaunchConfig {
grid_dim: (self.n as u32, 1, 1),
block_dim: (BLOCK_THREADS, 1, 1),
shared_mem_bytes: 0,
})
}
.map(|_| ())
.map_err(|err| driver_err("launch MatMulNBits int4 f32 block-128 GEMV", err))
}
#[allow(clippy::too_many_arguments)]
fn launch_int8_f32_gemv_block128(
&self,
activation: &TensorView,
packed: &TensorView,
scales: &TensorView,
zero_points: Option<&TensorView>,
bias: Option<&TensorView>,
output: &mut TensorMut,
k_blocks: usize,
) -> Result<()> {
let function =
self.runtime
.nvrtc_function(GEMV_MODULE, GEMV_SRC, GEMV_INT8_F32_BLOCK128_ENTRY)?;
let activation_ptr = cuptr(activation.data_ptr::<u8>() as *const c_void);
let packed_ptr = cuptr(packed.data_ptr::<u8>() as *const c_void);
let scales_ptr = cuptr(scales.data_ptr::<u8>() as *const c_void);
let zero_points_ptr = zero_points
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let bias_ptr = bias
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let output_ptr = cuptr(output.data_ptr_mut::<u8>() as *const c_void);
let k = as_i32("K", self.k)?;
let n = as_i32("N", self.n)?;
let k_blocks = as_i32("K block count", k_blocks)?;
let mut builder = self.runtime.stream().launch_builder(&function);
builder
.arg(&activation_ptr)
.arg(&packed_ptr)
.arg(&scales_ptr)
.arg(&zero_points_ptr)
.arg(&bias_ptr)
.arg(&output_ptr)
.arg(&k)
.arg(&n)
.arg(&k_blocks);
unsafe {
builder.launch(LaunchConfig {
grid_dim: (self.n as u32, 1, 1),
block_dim: (BLOCK_THREADS, 1, 1),
shared_mem_bytes: 0,
})
}
.map(|_| ())
.map_err(|err| driver_err("launch MatMulNBits int8 f32 block-128 GEMV", err))
}
#[allow(clippy::too_many_arguments)]
fn launch_f32_gemv(
&self,
activation: &TensorView,
packed: &TensorView,
scales: &TensorView,
zero_points: Option<&TensorView>,
bias: Option<&TensorView>,
output: &mut TensorMut,
k_blocks: usize,
blob_size: usize,
zp_row_bytes: usize,
) -> Result<()> {
let function = self
.runtime
.nvrtc_function(GEMV_MODULE, GEMV_SRC, GEMV_F32_ENTRY)?;
let activation_ptr = cuptr(activation.data_ptr::<u8>() as *const c_void);
let packed_ptr = cuptr(packed.data_ptr::<u8>() as *const c_void);
let scales_ptr = cuptr(scales.data_ptr::<u8>() as *const c_void);
let zero_points_ptr = zero_points
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let bias_ptr = bias
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let output_ptr = cuptr(output.data_ptr_mut::<u8>() as *const c_void);
let k = as_i32("K", self.k)?;
let n = as_i32("N", self.n)?;
let block_size = as_i32("block_size", self.block_size)?;
let k_blocks = as_i32("K block count", k_blocks)?;
let blob_size = as_i32("block blob size", blob_size)?;
let zp_row_bytes = as_i32("zero-point row size", zp_row_bytes)?;
let bits = as_i32("bits", self.bits)?;
let mut builder = self.runtime.stream().launch_builder(&function);
builder
.arg(&activation_ptr)
.arg(&packed_ptr)
.arg(&scales_ptr)
.arg(&zero_points_ptr)
.arg(&bias_ptr)
.arg(&output_ptr)
.arg(&k)
.arg(&n)
.arg(&block_size)
.arg(&k_blocks)
.arg(&blob_size)
.arg(&zp_row_bytes)
.arg(&bits);
unsafe {
builder.launch(LaunchConfig {
grid_dim: (self.n as u32, 1, 1),
block_dim: (BLOCK_THREADS, 1, 1),
shared_mem_bytes: 0,
})
}
.map(|_| ())
.map_err(|err| driver_err("launch MatMulNBits f32 GEMV", err))
}
#[allow(clippy::too_many_arguments)]
fn launch_accuracy4_gemv(
&self,
activation: &TensorView,
packed: &TensorView,
scales: &TensorView,
bias: Option<&TensorView>,
output: &mut TensorMut,
k_blocks: usize,
) -> Result<()> {
let workspace = self
.accuracy4_workspace
.as_ref()
.ok_or_else(|| error("accuracy_level=4 GEMV workspace is unavailable"))?
.lock()
.map_err(|_| error("accuracy_level=4 GEMV workspace lock poisoned"))?;
let quantize_function =
self.runtime
.nvrtc_function(GEMV_MODULE, GEMV_SRC, QUANTIZE_ACCURACY4_ENTRY)?;
let capabilities = self.runtime.capabilities();
let stage64 = use_accuracy4_stage64(
self.n,
capabilities.multiprocessor_count(),
capabilities.compute_capability(),
capabilities.max_shared_memory_per_block_optin(),
);
let gemv_entry = if stage64 {
GEMV_ACCURACY4_STAGE64_ENTRY
} else {
GEMV_ACCURACY4_ENTRY
};
let gemv_function = self
.runtime
.nvrtc_function(GEMV_MODULE, GEMV_SRC, gemv_entry)?;
let activation_ptr = cuptr(activation.data_ptr::<u8>() as *const c_void);
let packed_ptr = cuptr(packed.data_ptr::<u8>() as *const c_void);
let scales_ptr = cuptr(scales.data_ptr::<u8>() as *const c_void);
let bias_ptr = bias
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let output_ptr = cuptr(output.data_ptr_mut::<u8>() as *const c_void);
let k = as_i32("K", self.k)?;
let n = as_i32("N", self.n)?;
let k_blocks = as_i32("K block count", k_blocks)?;
let padded_k = as_i32("padded K", workspace.padded_k)?;
let mut quantize_builder = self.runtime.stream().launch_builder(&quantize_function);
quantize_builder
.arg(&activation_ptr)
.arg(&workspace.quantized_activation)
.arg(&workspace.activation_scale)
.arg(&k)
.arg(&padded_k);
unsafe {
quantize_builder.launch(LaunchConfig {
grid_dim: ((workspace.padded_k / 32) as u32, 1, 1),
block_dim: (32, 1, 1),
shared_mem_bytes: 0,
})
}
.map_err(|err| driver_err("launch MatMulNBits accuracy_level=4 quantization", err))?;
let mut gemv_builder = self.runtime.stream().launch_builder(&gemv_function);
gemv_builder
.arg(&workspace.quantized_activation)
.arg(&workspace.activation_scale)
.arg(&packed_ptr)
.arg(&scales_ptr)
.arg(&bias_ptr)
.arg(&output_ptr)
.arg(&k)
.arg(&n)
.arg(&k_blocks);
unsafe {
gemv_builder.launch(LaunchConfig {
grid_dim: (
self.n.div_ceil(GEMV_ACCURACY4_COLUMNS_PER_BLOCK) as u32,
1,
1,
),
block_dim: (GEMV_ACCURACY4_THREADS, 1, 1),
shared_mem_bytes: if stage64 {
GEMV_ACCURACY4_STAGE64_SHARED_BYTES
} else {
GEMV_ACCURACY4_SHARED_BYTES
},
})
}
.map(|_| ())
.map_err(|err| driver_err("launch MatMulNBits accuracy_level=4 GEMV", err))
}
#[allow(clippy::too_many_arguments)]
fn launch_accuracy4_gemv_blockwise(
&self,
activation: &TensorView,
packed: &TensorView,
scales: &TensorView,
zero_points: Option<&TensorView>,
bias: Option<&TensorView>,
output: &mut TensorMut,
k_blocks: usize,
blob_size: usize,
zp_row_bytes: usize,
) -> Result<()> {
let workspace = self
.accuracy4_workspace
.as_ref()
.ok_or_else(|| error("accuracy_level=4 GEMV workspace is unavailable"))?
.lock()
.map_err(|_| error("accuracy_level=4 GEMV workspace lock poisoned"))?;
let quantize_function = self.runtime.nvrtc_function(
GEMV_MODULE,
GEMV_SRC,
QUANTIZE_ACCURACY4_BLOCKWISE_ENTRY,
)?;
let gemv_function =
self.runtime
.nvrtc_function(GEMV_MODULE, GEMV_SRC, GEMV_ACCURACY4_BLOCKWISE_ENTRY)?;
let activation_ptr = cuptr(activation.data_ptr::<u8>() as *const c_void);
let packed_ptr = cuptr(packed.data_ptr::<u8>() as *const c_void);
let scales_ptr = cuptr(scales.data_ptr::<u8>() as *const c_void);
let zero_points_ptr = zero_points
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let bias_ptr = bias
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let output_ptr = cuptr(output.data_ptr_mut::<u8>() as *const c_void);
let k = as_i32("K", self.k)?;
let n = as_i32("N", self.n)?;
let block_size = as_i32("block_size", self.block_size)?;
let k_blocks_arg = as_i32("K block count", k_blocks)?;
let blob_size_arg = as_i32("block blob size", blob_size)?;
let zp_row_bytes_arg = as_i32("zero-point row size", zp_row_bytes)?;
let padded_k = as_i32("padded K", workspace.padded_k)?;
let mut quantize_builder = self.runtime.stream().launch_builder(&quantize_function);
quantize_builder
.arg(&activation_ptr)
.arg(&workspace.quantized_activation)
.arg(&workspace.activation_scale)
.arg(&k)
.arg(&block_size)
.arg(&padded_k);
unsafe {
quantize_builder.launch(LaunchConfig {
grid_dim: (k_blocks as u32, 1, 1),
block_dim: (32, 1, 1),
shared_mem_bytes: 0,
})
}
.map_err(|err| {
driver_err(
"launch MatMulNBits accuracy_level=4 blockwise quantization",
err,
)
})?;
let warps =
select_accuracy4_gemv_warps(self.n, self.runtime.capabilities().multiprocessor_count());
let grid = self.n.div_ceil(warps as usize) as u32;
let mut gemv_builder = self.runtime.stream().launch_builder(&gemv_function);
gemv_builder
.arg(&workspace.quantized_activation)
.arg(&workspace.activation_scale)
.arg(&packed_ptr)
.arg(&scales_ptr)
.arg(&zero_points_ptr)
.arg(&bias_ptr)
.arg(&output_ptr)
.arg(&k)
.arg(&n)
.arg(&k_blocks_arg)
.arg(&block_size)
.arg(&blob_size_arg)
.arg(&zp_row_bytes_arg);
unsafe {
gemv_builder.launch(LaunchConfig {
grid_dim: (grid, 1, 1),
block_dim: (warps * 32, 1, 1),
shared_mem_bytes: 0,
})
}
.map(|_| ())
.map_err(|err| driver_err("launch MatMulNBits accuracy_level=4 blockwise GEMV", err))
}
#[allow(clippy::too_many_arguments)]
fn launch_accuracy4(
&self,
activation: &TensorView,
packed: &TensorView,
scales: &TensorView,
zero_points: Option<&TensorView>,
bias: Option<&TensorView>,
output: &mut TensorMut,
m: usize,
k_blocks: usize,
blob_size: usize,
zp_row_bytes: usize,
) -> Result<()> {
let total = m.checked_mul(self.n).ok_or_else(|| {
error(format!(
"accuracy_level=4 output size {m} * {} overflows usize",
self.n
))
})?;
let blocks = total.div_ceil(BLOCK_THREADS as usize).clamp(1, 65_535) as u32;
let function =
self.runtime
.nvrtc_function(ACCURACY4_MODULE, ACCURACY4_SRC, ACCURACY4_ENTRY)?;
let activation_ptr = cuptr(activation.data_ptr::<u8>() as *const c_void);
let packed_ptr = cuptr(packed.data_ptr::<u8>() as *const c_void);
let scales_ptr = cuptr(scales.data_ptr::<u8>() as *const c_void);
let zero_points_ptr = zero_points
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let bias_ptr = bias
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let output_ptr = cuptr(output.data_ptr_mut::<u8>() as *const c_void);
let m = as_i32("M", m)?;
let k = as_i32("K", self.k)?;
let n = as_i32("N", self.n)?;
let block_size = as_i32("block_size", self.block_size)?;
let k_blocks = as_i32("K block count", k_blocks)?;
let blob_size = as_i32("block blob size", blob_size)?;
let zp_row_bytes = as_i32("zero-point row size", zp_row_bytes)?;
let mut builder = self.runtime.stream().launch_builder(&function);
builder
.arg(&activation_ptr)
.arg(&packed_ptr)
.arg(&scales_ptr)
.arg(&zero_points_ptr)
.arg(&bias_ptr)
.arg(&output_ptr)
.arg(&m)
.arg(&k)
.arg(&n)
.arg(&block_size)
.arg(&k_blocks)
.arg(&blob_size)
.arg(&zp_row_bytes);
unsafe {
builder.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (BLOCK_THREADS, 1, 1),
shared_mem_bytes: 0,
})
}
.map(|_| ())
.map_err(|err| driver_err("launch MatMulNBits accuracy_level=4", err))
}
#[allow(clippy::too_many_arguments)]
#[allow(clippy::too_many_arguments)]
fn try_dequant_f16_cublas_gemm(
&self,
activation: &TensorView,
packed: &TensorView,
scales: &TensorView,
scales_fp16: bool,
zero_points: Option<&TensorView>,
bias: Option<&TensorView>,
output: &mut TensorMut,
m: usize,
bias_row_stride: usize,
k_blocks: usize,
blob_size: usize,
zp_row_bytes: usize,
workspace: Option<WorkspaceView>,
) -> Result<bool> {
if bias.is_some() && (bias_row_stride != 0 || self.fold_bias_post_round) {
return Ok(false);
}
let Some(block_shift) = self.dequant_f16_block_shift() else {
return Ok(false);
};
let weight_bytes = self
.k
.checked_mul(self.n)
.and_then(|elems| elems.checked_mul(2))
.ok_or_else(|| error("dequantized f16 weight size overflowed"))?;
if weight_bytes > dequant_f16_gemm_max_scratch_bytes() {
return Ok(false);
}
let (weight, _warm) = self.marlin_repack_cache.ensure_scratch(5, weight_bytes)?;
let result = self
.launch_dequant_f16(
packed,
scales,
scales_fp16,
zero_points,
weight,
k_blocks,
blob_size,
zp_row_bytes,
block_shift,
)
.and_then(|()| {
let params = self.dequant_f16_gemm_ex(
cuptr(activation.data_ptr::<u8>() as *const c_void),
weight,
cuptr(output.data_ptr_mut::<u8>() as *const c_void),
m,
bias.map(|bias| cuptr(bias.data_ptr::<u8>() as *const c_void)),
);
unsafe {
blas::governed_gemm_ex(
self.runtime.blas(),
self.runtime.stream_ptr(),
¶ms,
workspace,
"MatMulNBits",
)
}
});
result.map(|()| true)
}
fn dequant_f16_gemm_ex(
&self,
activation: CUdeviceptr,
weight_nk: CUdeviceptr,
output: CUdeviceptr,
m: usize,
bias: Option<CUdeviceptr>,
) -> GemmEx {
GemmEx {
dtype: GemmDtype::F16,
transa: true,
transb: false,
m: self.n,
n: m,
k: self.k,
alpha: 1.0,
beta: 0.0,
a: weight_nk,
lda: self.k,
b: activation,
ldb: self.k,
c: output,
ldc: self.n,
epilogue: bias.map(|bias| GemmEpilogue {
kind: GemmEpilogueKind::Bias,
bias,
}),
}
}
fn dequant_f16_block_shift(&self) -> Option<i32> {
if self.bits != 4 || !self.k.is_multiple_of(8) {
return None;
}
if !self.block_size.is_multiple_of(8) || !self.block_size.is_power_of_two() {
return None;
}
Some(self.block_size.trailing_zeros() as i32)
}
#[allow(clippy::too_many_arguments)]
fn launch_dequant_f16(
&self,
packed: &TensorView,
scales: &TensorView,
scales_fp16: bool,
zero_points: Option<&TensorView>,
weight: cudarc::driver::sys::CUdeviceptr,
k_blocks: usize,
blob_size: usize,
zp_row_bytes: usize,
block_shift: i32,
) -> Result<()> {
let packed_ptr = cuptr(packed.data_ptr::<u8>() as *const c_void);
let scales_ptr = cuptr(scales.data_ptr::<u8>() as *const c_void);
let zero_points_ptr = zero_points
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let k = as_i32("K", self.k)?;
let k_blocks_i = as_i32("K block count", k_blocks)?;
let blob_size_i = as_i32("block blob size", blob_size)?;
let zp_row_bytes_i = as_i32("zero-point row size", zp_row_bytes)?;
let scales_fp16_flag: i32 = scales_fp16 as i32;
let words = self.k / 8;
let grid_x = words.div_ceil(BLOCK_THREADS as usize) as u32;
let grid_y = as_i32("N", self.n)? as u32;
let function =
self.runtime
.nvrtc_function(DEQUANT_F16_MODULE, DEQUANT_F16_SRC, DEQUANT_F16_ENTRY)?;
let mut builder = self.runtime.stream().launch_builder(&function);
builder
.arg(&packed_ptr)
.arg(&scales_ptr)
.arg(&zero_points_ptr)
.arg(&weight)
.arg(&k)
.arg(&k_blocks_i)
.arg(&blob_size_i)
.arg(&zp_row_bytes_i)
.arg(&block_shift)
.arg(&scales_fp16_flag);
unsafe {
builder.launch(LaunchConfig {
grid_dim: (grid_x.max(1), grid_y, 1),
block_dim: (BLOCK_THREADS, 1, 1),
shared_mem_bytes: 0,
})
}
.map(|_| ())
.map_err(|err| driver_err("launch MatMulNBits f16 dequant", err))
}
#[allow(clippy::too_many_arguments)]
fn launch_dequant(
&self,
packed: &TensorView,
scales: &TensorView,
zero_points: Option<&TensorView>,
group_indices: Option<&TensorView>,
weight: cudarc::driver::sys::CUdeviceptr,
k_blocks: usize,
blob_size: usize,
zp_row_bytes: usize,
) -> Result<()> {
let packed_ptr = cuptr(packed.data_ptr::<u8>() as *const c_void);
let scales_ptr = cuptr(scales.data_ptr::<u8>() as *const c_void);
let zero_points_ptr = zero_points
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let group_indices_ptr = group_indices
.map(|tensor| cuptr(tensor.data_ptr::<u8>() as *const c_void))
.unwrap_or(0);
let k = as_i32("K", self.k)?;
let n = as_i32("N", self.n)?;
let block_size = as_i32("block_size", self.block_size)?;
let k_blocks = as_i32("K block count", k_blocks)?;
let blob_size = as_i32("block blob size", blob_size)?;
let zp_row_bytes = as_i32("zero-point row size", zp_row_bytes)?;
let bits = as_i32("bits", self.bits)?;
let total = self.k * self.n;
let blocks = total.div_ceil(BLOCK_THREADS as usize).clamp(1, 65_535) as u32;
let function = self
.runtime
.nvrtc_function(DEQUANT_MODULE, DEQUANT_SRC, DEQUANT_ENTRY)?;
let mut builder = self.runtime.stream().launch_builder(&function);
builder
.arg(&packed_ptr)
.arg(&scales_ptr)
.arg(&zero_points_ptr)
.arg(&group_indices_ptr)
.arg(&weight)
.arg(&k)
.arg(&n)
.arg(&block_size)
.arg(&k_blocks)
.arg(&blob_size)
.arg(&zp_row_bytes)
.arg(&bits);
unsafe {
builder.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (BLOCK_THREADS, 1, 1),
shared_mem_bytes: 0,
})
}
.map(|_| ())
.map_err(|err| driver_err("launch MatMulNBits dequant", err))
}
}
impl Kernel for MatMulNBitsKernel {
fn set_constant_inputs(&mut self, constant_inputs: &[bool]) {
for (index, is_constant) in self.constant_inputs.iter_mut().enumerate() {
*is_constant = constant_inputs.get(index).copied().unwrap_or(false);
}
}
fn execute(&self, inputs: &[TensorView], outputs: &mut [TensorMut]) -> Result<()> {
self.run(inputs, outputs, None)
}
fn workspace_requirement(&self, inputs: &[TensorMetadata<'_>]) -> Result<WorkspaceRequirement> {
self.workspace_requirement_for(inputs)
}
fn execute_with_workspace(
&self,
inputs: &[TensorView],
outputs: &mut [TensorMut],
workspace: Option<WorkspaceView>,
) -> Result<()> {
self.run(inputs, outputs, workspace)
}
fn supports_strided_input(&self, _input_idx: usize) -> bool {
false
}
fn device_graph_resources(&self) -> Vec<DeviceGraphResource> {
let mut resources = self.marlin_repack_cache.device_graph_resources();
if let Some(workspace) = &self.accuracy4_workspace
&& let Ok(workspace) = workspace.lock()
{
resources.push(workspace.device_graph_resource());
}
if let Ok(scratch) = self.bf16_scratch.lock()
&& let Some(resource) = scratch.device_graph_resource()
{
resources.push(resource);
}
if let Ok(cache) = self.bf16_const_cache.lock()
&& let Some(resource) = cache.device_graph_resource()
{
resources.push(resource);
}
resources
}
fn capture_support(&self) -> onnx_runtime_ep_api::CaptureSupport {
if self.last_call_capture_safe.load(Ordering::Relaxed) {
onnx_runtime_ep_api::CaptureSupport::Supported
} else {
onnx_runtime_ep_api::CaptureSupport::unsupported(
"requires M==1 decode GEMV without group_indices, or a warm-cache Marlin M>1 GEMM; a cold Marlin repack and portable prefill are outside the advertised capture contract and group_indices validation reads D2H",
)
}
}
}
fn optional_input<'a>(inputs: &'a [TensorView<'a>], index: usize) -> Option<&'a TensorView<'a>> {
inputs.get(index).filter(|input| !input.is_absent())
}
fn required_positive_attr(node: &Node, name: &str) -> Result<usize> {
let value = optional_int_attr(node, name)?
.ok_or_else(|| error(format!("missing required integer attribute '{name}'")))?;
if value <= 0 {
return Err(error(format!(
"attribute '{name}' must be positive, got {value}"
)));
}
Ok(value as usize)
}
fn optional_int_attr(node: &Node, name: &str) -> Result<Option<i64>> {
match node.attr(name) {
Some(attribute) => attribute
.as_int()
.map(Some)
.ok_or_else(|| error(format!("attribute '{name}' must be an integer"))),
None => Ok(None),
}
}
fn require_dtype(name: &str, got: DataType, expected: DataType) -> Result<()> {
if got != expected {
return Err(error(format!(
"{name} must have dtype {expected:?}, got {got:?}"
)));
}
Ok(())
}
fn require_gamma_dtype(got: DataType) -> Result<()> {
if got != DataType::Float16 && got != DataType::Float32 {
return Err(error(format!(
"gamma must have dtype Float16 or Float, got {got:?}"
)));
}
Ok(())
}
fn require_shape(name: &str, got: &[usize], expected: &[usize]) -> Result<()> {
if got != expected {
return Err(error(format!(
"{name} must have shape {expected:?}, got {got:?}"
)));
}
Ok(())
}
fn require_flat_or_matrix_shape(
name: &str,
got: &[usize],
rows: usize,
columns: usize,
) -> Result<()> {
if got != [rows * columns] && got != [rows, columns] {
return Err(error(format!(
"{name} must have shape [{}] or [{rows}, {columns}], got {got:?}",
rows * columns
)));
}
Ok(())
}
fn as_i32(name: &str, value: usize) -> Result<i32> {
i32::try_from(value).map_err(|_| error(format!("{name}={value} exceeds i32")))
}
fn error(message: impl Into<String>) -> EpError {
EpError::KernelFailed(format!("cuda_ep MatMulNBits: {}", message.into()))
}
#[cfg(test)]
mod tests {
use half::f16;
use onnx_runtime_ep_api::{DevicePtr, DevicePtrMut, TensorMut, TensorView};
use onnx_runtime_ir::{DataType, DeviceId};
use super::*;
#[test]
fn marlin_repack_requires_constant_weight_inputs() {
let mut constant_inputs = [false; 8];
assert!(!marlin_weight_inputs_are_constant(&constant_inputs, false));
constant_inputs[1] = true;
assert!(marlin_weight_inputs_are_constant(&constant_inputs, false));
assert!(!marlin_weight_inputs_are_constant(&constant_inputs, true));
constant_inputs[3] = true;
assert!(marlin_weight_inputs_are_constant(&constant_inputs, true));
}
static INTERLEAVE_TEST_ENV_LOCK: std::sync::OnceLock<std::sync::RwLock<()>> =
std::sync::OnceLock::new();
struct LeverEnvGuard(#[allow(dead_code)] std::sync::RwLockWriteGuard<'static, ()>);
impl LeverEnvGuard {
fn acquire() -> Self {
Self(
INTERLEAVE_TEST_ENV_LOCK
.get_or_init(Default::default)
.write()
.unwrap_or_else(|e| e.into_inner()),
)
}
}
impl Drop for LeverEnvGuard {
fn drop(&mut self) {
unsafe {
std::env::remove_var("ONNX_GENAI_INTERLEAVE_DEQUANT");
std::env::remove_var("ONNX_GENAI_GENERAL_SPLITK");
std::env::remove_var("ONNX_GENAI_GEMV_PIPELINE");
std::env::remove_var("ONNX_GENAI_GATEUP_VEC");
std::env::remove_var("ONNX_GENAI_GATEUP_OCC");
}
}
}
fn default_levers_guard() -> std::sync::RwLockReadGuard<'static, ()> {
INTERLEAVE_TEST_ENV_LOCK
.get_or_init(Default::default)
.read()
.unwrap_or_else(|e| e.into_inner())
}
const QWEN_DOWN_K: usize = 4864;
const QWEN_DOWN_N: usize = 896;
const STAGED_DOWN_REFERENCE_ENTRY: &str =
"matmul_nbits_gemv_f16_scales_f16_down_staged_reference";
const STAGED_DOWN_REFERENCE_SRC: &str = r#"
extern "C" __global__ void matmul_nbits_gemv_f16_scales_f16_down_staged_reference(
const __half* __restrict__ activation,
const unsigned char* __restrict__ packed,
const void* __restrict__ scales_raw,
const unsigned char* __restrict__ zero_points,
const __half* __restrict__ bias,
__half* __restrict__ output,
const int k,
const int n,
const int block_size,
const int k_blocks,
const int blob_size,
const int zp_row_bytes,
const int scales_fp16,
const int bias_post_round)
{
(void)block_size;
(void)zero_points;
(void)zp_row_bytes;
(void)scales_fp16;
extern __shared__ uint4 activation_shared[];
__shared__ float warp_sums[8][8];
const __half* __restrict__ scales =
reinterpret_cast<const __half*>(scales_raw);
const int tid = (int)threadIdx.x;
const int lane = tid & 31;
const int warp = tid >> 5;
const int column_base = (int)blockIdx.x * 8;
for (int vector = tid; vector * 8 < k; vector += (int)blockDim.x) {
activation_shared[vector] =
permute_activation_f16x8(activation + vector * 8);
}
__syncthreads();
float values[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
for (int block = tid; block < k_blocks; block += (int)blockDim.x) {
const uint4 activation0 = activation_shared[block * 4];
const uint4 activation1 = activation_shared[block * 4 + 1];
const uint4 activation2 = activation_shared[block * 4 + 2];
const uint4 activation3 = activation_shared[block * 4 + 3];
#pragma unroll
for (int tile_column = 0; tile_column < 8; ++tile_column) {
const int column = column_base + tile_column;
if (column < n) {
const long packed_start =
((long)column * k_blocks + block) * blob_size;
const uint4 packed_weights =
*reinterpret_cast<const uint4*>(packed + packed_start);
const __half scale = scales[(long)column * k_blocks + block];
values[tile_column] += dot_int4x32_f16_permuted_scaled(
packed_weights,
activation0,
activation1,
activation2,
activation3,
scale);
}
}
}
#pragma unroll
for (int tile_column = 0; tile_column < 8; ++tile_column) {
const float value = warp_sum(values[tile_column]);
if (lane == 0) {
warp_sums[warp][tile_column] = value;
}
}
__syncthreads();
if (warp == 0 && lane < 8) {
const int column = column_base + lane;
float value = warp_sums[0][lane];
value += warp_sums[1][lane];
value += warp_sums[2][lane];
value += warp_sums[3][lane];
value += warp_sums[4][lane];
value += warp_sums[5][lane];
value += warp_sums[6][lane];
value += warp_sums[7][lane];
output[column] = fold_bias_f16(value, bias, column, bias_post_round);
}
}
"#;
fn runtime() -> Option<Arc<CudaRuntime>> {
crate::test_support::maybe_runtime()
}
fn as_bytes<T: Copy>(values: &[T]) -> &[u8] {
unsafe {
std::slice::from_raw_parts(values.as_ptr().cast::<u8>(), std::mem::size_of_val(values))
}
}
fn as_bytes_mut<T: Copy>(values: &mut [T]) -> &mut [u8] {
unsafe {
std::slice::from_raw_parts_mut(
values.as_mut_ptr().cast::<u8>(),
std::mem::size_of_val(values),
)
}
}
#[test]
#[ignore = "perf microbench; requires a dedicated idle SM80+ CUDA device"]
fn marlin_m_gt_1_op_wall_vs_tiled() {
let Some(runtime) = runtime() else {
eprintln!("skipping Marlin M>1 wall bench: CUDA runtime unavailable");
return;
};
if runtime
.require_nvrtc_half_headers("marlin_m_gt_1_wall")
.is_err()
|| !marlin_gemm::device_supports_marlin(runtime.capabilities().compute_capability())
{
eprintln!("skipping Marlin M>1 wall bench: headers unavailable or pre-SM80");
return;
}
let k = 5120usize;
let n = 13824usize;
let block_size = 128usize;
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let zp_row_bytes = k_blocks.div_ceil(2);
let mut state = 0xcafe_f00d_1234_5678_u64;
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let max_m = 128usize;
let mut activation_f16 = vec![f16::ZERO; max_m * k];
for h in activation_f16.iter_mut() {
*h = f16::from_f32(next());
}
let mut quant = vec![0u8; n * k];
for value in quant.iter_mut() {
*value = ((next() * 0.5 + 0.5) * 15.0).round().clamp(0.0, 15.0) as u8;
}
let mut packed = vec![0u8; n * k_blocks * blob_size];
for col in 0..n {
for block in 0..k_blocks {
for pair in 0..blob_size {
let low = quant[col * k + block * block_size + pair * 2] & 15;
let high = quant[col * k + block * block_size + pair * 2 + 1] & 15;
packed[(col * k_blocks + block) * blob_size + pair] = low | (high << 4);
}
}
}
let mut zp_packed = vec![0u8; n * zp_row_bytes];
for byte in zp_packed.iter_mut() {
*byte = ((next() * 0.5 + 0.5) * 255.0) as u8;
}
let mut scale_f16 = vec![f16::ZERO; n * k_blocks];
for h in scale_f16.iter_mut() {
*h = f16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5));
}
let activation_dev = runtime.alloc_raw(activation_f16.len() * 2).unwrap();
let packed_dev = runtime.alloc_raw(packed.len()).unwrap();
let scales_dev = runtime.alloc_raw(scale_f16.len() * 2).unwrap();
let zp_dev = runtime.alloc_raw(zp_packed.len()).unwrap();
let output_dev = runtime.alloc_raw(max_m * n * 2).unwrap();
unsafe {
runtime
.htod(as_bytes(&activation_f16), activation_dev)
.unwrap();
runtime.htod(&packed, packed_dev).unwrap();
runtime.htod(as_bytes(&scale_f16), scales_dev).unwrap();
runtime.htod(&zp_packed, zp_dev).unwrap();
}
let device = DeviceId::cuda(0);
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let zp_shape = [n, zp_row_bytes];
let zp_strides = [zp_row_bytes as i64, 1];
let kernel = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: false,
decomposed_silu: false,
rmsnorm_prologue: false,
rmsnorm_epsilon: 1e-5,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
let time_path = |kernel: &MatMulNBitsKernel, m: usize| -> f64 {
let a_shape = [m, k];
let a_strides = [k as i64, 1];
let y_shape = [m, n];
let y_strides = [n as i64, 1];
let inputs = vec![
TensorView::new(
device_ptr(activation_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
),
TensorView::new(
device_ptr(packed_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
),
TensorView::new(
device_ptr(scales_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
),
TensorView::new(
device_ptr(zp_dev),
DataType::Uint8,
&zp_shape,
&zp_strides,
device,
),
];
let run = || {
let mut outputs = [TensorMut::new(
device_ptr_mut(output_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
)];
kernel.run(&inputs, &mut outputs, None).unwrap();
};
run();
runtime.synchronize().unwrap();
let iters = 30;
let mut samples = Vec::with_capacity(iters);
for _ in 0..iters {
let start = std::time::Instant::now();
run();
runtime.synchronize().unwrap();
samples.push(start.elapsed().as_secs_f64() * 1e6);
}
samples.sort_by(|a, b| a.partial_cmp(b).unwrap());
samples[samples.len() / 2]
};
eprintln!("Marlin M>1 wall vs tiled @ K={k} N={n} block={block_size} (median us):");
eprintln!(" M tiled_us marlin_us speedup");
for &m in &[1usize, 2, 4, 8, 16, 32, 64, 128] {
unsafe {
std::env::set_var("ONNX_GENAI_MARLIN_M_GT_1", "0");
}
let tiled = time_path(&kernel, m);
unsafe {
std::env::set_var("ONNX_GENAI_MARLIN_M_GT_1", "1");
}
let marlin = time_path(&kernel, m);
let marlin_used = m > 1;
eprintln!(
" {m:<5} {tiled:>8.1} {marlin:>8.1} {:>6.2}x{}",
tiled / marlin,
if marlin_used {
""
} else {
" (M=1 stays on GEMV)"
}
);
}
unsafe {
std::env::remove_var("ONNX_GENAI_MARLIN_M_GT_1");
}
unsafe {
runtime.free_raw(activation_dev).unwrap();
runtime.free_raw(packed_dev).unwrap();
runtime.free_raw(scales_dev).unwrap();
runtime.free_raw(zp_dev).unwrap();
runtime.free_raw(output_dev).unwrap();
}
}
#[test]
#[ignore = "requires an SM80+ CUDA device"]
fn marlin_m_gt_1_rmsnorm_op_parity() {
let Some(runtime) = runtime() else {
eprintln!("skipping Marlin M>1 rmsnorm parity: CUDA runtime unavailable");
return;
};
if runtime
.require_nvrtc_half_headers("marlin_m_gt_1_rmsnorm")
.is_err()
|| !marlin_gemm::device_supports_marlin(runtime.capabilities().compute_capability())
{
eprintln!("skipping Marlin M>1 rmsnorm parity: headers unavailable or pre-SM80");
return;
}
let m = 8usize;
let k = 2048usize;
let n = 64usize;
let block_size = 128usize;
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let zp_row_bytes = k_blocks.div_ceil(2);
let mut state = 0x0bad_c0de_feed_face_u64;
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let mut activation = vec![f16::ZERO; m * k];
for h in activation.iter_mut() {
*h = f16::from_f32(next());
}
let mut packed = vec![0u8; n * k_blocks * blob_size];
for byte in packed.iter_mut() {
*byte = ((next() * 0.5 + 0.5) * 255.0) as u8;
}
let mut zp_packed = vec![0u8; n * zp_row_bytes];
for byte in zp_packed.iter_mut() {
*byte = ((next() * 0.5 + 0.5) * 255.0) as u8;
}
let mut scale_f16 = vec![f16::ZERO; n * k_blocks];
for h in scale_f16.iter_mut() {
*h = f16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5));
}
let mut gamma = vec![f16::ZERO; k];
for h in gamma.iter_mut() {
*h = f16::from_f32(0.5 + 0.5 * (next() * 0.5 + 0.5));
}
let activation_dev = runtime.alloc_raw(activation.len() * 2).unwrap();
let packed_dev = runtime.alloc_raw(packed.len()).unwrap();
let scales_dev = runtime.alloc_raw(scale_f16.len() * 2).unwrap();
let zp_dev = runtime.alloc_raw(zp_packed.len()).unwrap();
let gamma_dev = runtime.alloc_raw(gamma.len() * 2).unwrap();
let output_dev = runtime.alloc_raw(m * n * 2).unwrap();
unsafe {
runtime.htod(as_bytes(&activation), activation_dev).unwrap();
runtime.htod(&packed, packed_dev).unwrap();
runtime.htod(as_bytes(&scale_f16), scales_dev).unwrap();
runtime.htod(&zp_packed, zp_dev).unwrap();
runtime.htod(as_bytes(&gamma), gamma_dev).unwrap();
}
let device = DeviceId::cuda(0);
let a_shape = [m, k];
let a_strides = [k as i64, 1];
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let zp_shape = [n, zp_row_bytes];
let zp_strides = [zp_row_bytes as i64, 1];
let gamma_shape = [k];
let gamma_strides = [1i64];
let y_shape = [m, n];
let y_strides = [n as i64, 1];
let inputs = vec![
TensorView::new(
device_ptr(activation_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
),
TensorView::new(
device_ptr(packed_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
),
TensorView::new(
device_ptr(scales_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
),
TensorView::new(
device_ptr(zp_dev),
DataType::Uint8,
&zp_shape,
&zp_strides,
device,
),
TensorView::absent(DataType::Int32),
TensorView::absent(DataType::Float16),
TensorView::new(
device_ptr(gamma_dev),
DataType::Float16,
&gamma_shape,
&gamma_strides,
device,
),
];
let kernel = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: false,
decomposed_silu: false,
rmsnorm_prologue: true,
rmsnorm_epsilon: 1e-5,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
let run = |want_marlin: bool| -> Vec<f16> {
unsafe {
if want_marlin {
std::env::set_var("ONNX_GENAI_MARLIN_M_GT_1", "1");
} else {
std::env::set_var("ONNX_GENAI_MARLIN_M_GT_1", "0");
}
}
let mut outputs = [TensorMut::new(
device_ptr_mut(output_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
)];
kernel.run(&inputs, &mut outputs, None).unwrap();
runtime.synchronize().unwrap();
let mut got = vec![f16::ZERO; m * n];
unsafe {
runtime.dtoh(as_bytes_mut(&mut got), output_dev).unwrap();
}
got
};
let tiled = run(false);
let marlin = run(true);
unsafe {
std::env::remove_var("ONNX_GENAI_MARLIN_M_GT_1");
}
unsafe {
runtime.free_raw(activation_dev).unwrap();
runtime.free_raw(packed_dev).unwrap();
runtime.free_raw(scales_dev).unwrap();
runtime.free_raw(zp_dev).unwrap();
runtime.free_raw(gamma_dev).unwrap();
runtime.free_raw(output_dev).unwrap();
}
let mut worst_abs = 0.0f32;
let mut max_out = 0.0f32;
for (mv, tv) in marlin.iter().zip(tiled.iter()) {
let mf = mv.to_f32();
let tf = tv.to_f32();
assert!(mf.is_finite(), "Marlin rmsnorm output must be finite");
worst_abs = worst_abs.max((mf - tf).abs());
max_out = max_out.max(tf.abs());
}
let tol = 2e-2 * max_out.max(1e-3);
eprintln!(
"Marlin M>1 rmsnorm parity: worst_abs={worst_abs:.5}, max_out={max_out:.5}, tol={tol:.5}"
);
assert!(
worst_abs <= tol,
"fused rmsnorm Marlin output diverges from tiled: worst_abs={worst_abs} > tol={tol}"
);
}
fn run_marlin_gate_up_parity(with_gamma: bool, decomposed: bool, check_capture: bool) {
let Some(runtime) = runtime() else {
eprintln!("skipping Marlin gate/up parity: CUDA runtime unavailable");
return;
};
if runtime
.require_nvrtc_half_headers("marlin_gate_up")
.is_err()
|| !marlin_gemm::device_supports_marlin(runtime.capabilities().compute_capability())
{
eprintln!("skipping Marlin gate/up parity: headers unavailable or pre-SM80");
return;
}
let m = 8usize;
let k = 1024usize;
let n = 256usize;
let block_size = 32usize;
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let zp_row_bytes = (k_blocks * 4).div_ceil(8);
let mut state = 0xa5a5_1234_dead_beef_u64;
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let mut activation = vec![f16::ZERO; m * k];
for h in activation.iter_mut() {
*h = f16::from_f32(next() * 0.5);
}
let mut make_weights = |scale: f32| {
let mut packed = vec![0u8; n * k_blocks * blob_size];
for byte in packed.iter_mut() {
*byte = ((next() * 0.5 + 0.5) * 255.0) as u8;
}
let mut scales = vec![f16::ZERO; n * k_blocks];
for h in scales.iter_mut() {
*h = f16::from_f32(scale * (0.5 + 0.5 * (next() * 0.5 + 0.5)));
}
let mut zp = vec![0u8; n * zp_row_bytes];
for byte in zp.iter_mut() {
*byte = ((next() * 0.5 + 0.5) * 255.0) as u8;
}
(packed, scales, zp)
};
let (packed_gate, scales_gate, zp_gate) = make_weights(0.02);
let (packed_up, scales_up, zp_up) = make_weights(0.02);
let mut gamma = vec![f16::ZERO; k];
for h in gamma.iter_mut() {
*h = f16::from_f32(0.5 + 0.5 * (next() * 0.5 + 0.5));
}
let activation_dev = runtime.alloc_raw(activation.len() * 2).unwrap();
let packed_gate_dev = runtime.alloc_raw(packed_gate.len()).unwrap();
let scales_gate_dev = runtime.alloc_raw(scales_gate.len() * 2).unwrap();
let packed_up_dev = runtime.alloc_raw(packed_up.len()).unwrap();
let scales_up_dev = runtime.alloc_raw(scales_up.len() * 2).unwrap();
let zp_gate_dev = runtime.alloc_raw(zp_gate.len()).unwrap();
let zp_up_dev = runtime.alloc_raw(zp_up.len()).unwrap();
let gamma_dev = runtime.alloc_raw(gamma.len() * 2).unwrap();
let output_dev = runtime.alloc_raw(m * n * 2).unwrap();
unsafe {
runtime.htod(as_bytes(&activation), activation_dev).unwrap();
runtime.htod(&packed_gate, packed_gate_dev).unwrap();
runtime
.htod(as_bytes(&scales_gate), scales_gate_dev)
.unwrap();
runtime.htod(&packed_up, packed_up_dev).unwrap();
runtime.htod(as_bytes(&scales_up), scales_up_dev).unwrap();
runtime.htod(&zp_gate, zp_gate_dev).unwrap();
runtime.htod(&zp_up, zp_up_dev).unwrap();
runtime.htod(as_bytes(&gamma), gamma_dev).unwrap();
}
let device = DeviceId::cuda(0);
let a_shape = [m, k];
let a_strides = [k as i64, 1];
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let zp_shape = [n, zp_row_bytes];
let zp_strides = [zp_row_bytes as i64, 1];
let gamma_shape = [k];
let gamma_strides = [1i64];
let y_shape = [m, n];
let y_strides = [n as i64, 1];
let mut inputs = vec![
TensorView::new(
device_ptr(activation_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
),
TensorView::new(
device_ptr(packed_gate_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
),
TensorView::new(
device_ptr(scales_gate_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
),
TensorView::new(
device_ptr(packed_up_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
),
TensorView::new(
device_ptr(scales_up_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
),
];
if with_gamma {
inputs.push(TensorView::new(
device_ptr(gamma_dev),
DataType::Float16,
&gamma_shape,
&gamma_strides,
device,
));
} else {
inputs.push(TensorView::absent(DataType::Float16));
}
inputs.push(TensorView::new(
device_ptr(zp_gate_dev),
DataType::Uint8,
&zp_shape,
&zp_strides,
device,
));
inputs.push(TensorView::new(
device_ptr(zp_up_dev),
DataType::Uint8,
&zp_shape,
&zp_strides,
device,
));
let kernel = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: true,
decomposed_silu: decomposed,
rmsnorm_prologue: with_gamma,
rmsnorm_epsilon: 1e-5,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
let run = |want_marlin: bool| -> Vec<f16> {
unsafe {
if want_marlin {
std::env::set_var("ONNX_GENAI_MARLIN_M_GT_1", "1");
} else {
std::env::set_var("ONNX_GENAI_MARLIN_M_GT_1", "0");
}
}
let mut outputs = [TensorMut::new(
device_ptr_mut(output_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
)];
kernel.run(&inputs, &mut outputs, None).unwrap();
runtime.synchronize().unwrap();
let mut got = vec![f16::ZERO; m * n];
unsafe {
runtime.dtoh(as_bytes_mut(&mut got), output_dev).unwrap();
}
got
};
let tiled = run(false);
let marlin_cold = run(true);
let cold_safe = kernel.last_call_capture_safe.load(Ordering::Relaxed);
let marlin_warm = run(true);
let warm_safe = kernel.last_call_capture_safe.load(Ordering::Relaxed);
unsafe {
std::env::remove_var("ONNX_GENAI_MARLIN_M_GT_1");
}
unsafe {
runtime.free_raw(activation_dev).unwrap();
runtime.free_raw(packed_gate_dev).unwrap();
runtime.free_raw(scales_gate_dev).unwrap();
runtime.free_raw(packed_up_dev).unwrap();
runtime.free_raw(scales_up_dev).unwrap();
runtime.free_raw(zp_gate_dev).unwrap();
runtime.free_raw(zp_up_dev).unwrap();
runtime.free_raw(gamma_dev).unwrap();
runtime.free_raw(output_dev).unwrap();
}
let mut worst_abs = 0.0f32;
let mut max_out = 0.0f32;
for (mv, tv) in marlin_warm.iter().zip(tiled.iter()) {
let mf = mv.to_f32();
let tf = tv.to_f32();
assert!(mf.is_finite(), "Marlin gate/up output must be finite");
worst_abs = worst_abs.max((mf - tf).abs());
max_out = max_out.max(tf.abs());
}
let tol = 3e-2 * max_out.max(1e-3);
eprintln!(
"Marlin gate/up parity (gamma={with_gamma}, decomposed={decomposed}): \
worst_abs={worst_abs:.5}, max_out={max_out:.5}, tol={tol:.5}, \
cold_safe={cold_safe}, warm_safe={warm_safe}"
);
assert!(
worst_abs <= tol,
"Marlin gate/up output diverges from tiled: worst_abs={worst_abs} > tol={tol}"
);
assert_eq!(
marlin_cold, marlin_warm,
"warm gate/up Marlin replay must be byte-identical to the cold run"
);
if check_capture {
assert!(warm_safe, "warm gate/up call must remain capture-safe");
}
}
#[test]
#[ignore = "requires an SM80+ CUDA device"]
fn marlin_gate_up_swiglu_matches_tiled_plain() {
run_marlin_gate_up_parity(false, false, true);
}
#[test]
#[ignore = "requires an SM80+ CUDA device"]
fn marlin_gate_up_swiglu_matches_tiled_rmsnorm() {
run_marlin_gate_up_parity(true, false, true);
}
#[test]
#[ignore = "requires an SM80+ CUDA device"]
fn marlin_gate_up_decomposed_swiglu_matches_tiled_rmsnorm() {
run_marlin_gate_up_parity(true, true, false);
}
fn device_ptr(raw: CUdeviceptr) -> DevicePtr {
DevicePtr(raw as usize as *const c_void)
}
fn device_ptr_mut(raw: CUdeviceptr) -> DevicePtrMut {
DevicePtrMut(raw as usize as *mut c_void)
}
fn run_parity(scales_fp16: bool, with_bias: bool) -> (f32, f32, f32, bool) {
run_parity_dims(4096, 70, scales_fp16, with_bias, false)
}
fn run_parity_dims(
k: usize,
n: usize,
scales_fp16: bool,
with_bias: bool,
explicit_zp: bool,
) -> (f32, f32, f32, bool) {
run_parity_dims_block(k, n, 32, scales_fp16, with_bias, explicit_zp)
}
fn run_parity_dims_block(
k: usize,
n: usize,
block_size: usize,
scales_fp16: bool,
with_bias: bool,
explicit_zp: bool,
) -> (f32, f32, f32, bool) {
let Some(runtime) = runtime() else {
eprintln!("skipping MatMulNBits fp16 GEMV parity test: CUDA runtime unavailable");
return (0.0, 0.0, 0.0, true);
};
if runtime
.require_nvrtc_half_headers("matmul_nbits_gemv_f16")
.is_err()
{
eprintln!("skipping MatMulNBits fp16 GEMV parity test: fp16 NVRTC headers unavailable");
return (0.0, 0.0, 0.0, true);
}
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let zp_row_bytes = k_blocks.div_ceil(2);
let mut state = 0x9e37_79b9_7f4a_7c15u64;
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let mut activation_f16 = vec![f16::ZERO; k];
let mut activation_ref = vec![0.0f32; k];
for (dst_h, dst_f) in activation_f16.iter_mut().zip(activation_ref.iter_mut()) {
let h = f16::from_f32(next());
*dst_h = h;
*dst_f = h.to_f32();
}
let mut quant = vec![0u8; n * k];
for value in quant.iter_mut() {
*value = ((next() * 0.5 + 0.5) * 15.0).round().clamp(0.0, 15.0) as u8;
}
let mut packed = vec![0u8; n * k_blocks * blob_size];
for col in 0..n {
for block in 0..k_blocks {
for pair in 0..blob_size {
let low = quant[col * k + block * block_size + pair * 2] & 15;
let high = quant[col * k + block * block_size + pair * 2 + 1] & 15;
packed[(col * k_blocks + block) * blob_size + pair] = low | (high << 4);
}
}
}
let mut zp_codes = vec![8i32; n * k_blocks];
let mut zp_packed = vec![0u8; n * zp_row_bytes];
if explicit_zp {
for code in zp_codes.iter_mut().take(n * k_blocks) {
*code = ((next() * 0.5 + 0.5) * 15.0).round().clamp(0.0, 15.0) as i32;
}
for col in 0..n {
for block in 0..k_blocks {
let code = (zp_codes[col * k_blocks + block] & 15) as u8;
let byte = &mut zp_packed[col * zp_row_bytes + block / 2];
if block & 1 == 0 {
*byte = (*byte & 0xf0) | code;
} else {
*byte = (*byte & 0x0f) | (code << 4);
}
}
}
}
let mut scale_ref = vec![0.0f32; n * k_blocks];
let mut scale_f16 = vec![f16::ZERO; n * k_blocks];
let mut scale_f32 = vec![0.0f32; n * k_blocks];
for i in 0..n * k_blocks {
let raw = 0.015 + 0.01 * (next() * 0.5 + 0.5);
if scales_fp16 {
let h = f16::from_f32(raw);
scale_f16[i] = h;
scale_ref[i] = h.to_f32();
} else {
scale_f32[i] = raw;
scale_ref[i] = raw;
}
}
let mut bias_f16 = vec![f16::ZERO; n];
let mut bias_ref = vec![0.0f32; n];
if with_bias {
for (h, f) in bias_f16.iter_mut().zip(bias_ref.iter_mut()) {
let value = f16::from_f32(next());
*h = value;
*f = value.to_f32();
}
}
let mut expected = vec![0.0f32; n];
for col in 0..n {
let mut acc = 0.0f64;
for block in 0..k_blocks {
let scale = scale_ref[col * k_blocks + block] as f64;
let zero_point = zp_codes[col * k_blocks + block];
for within in 0..block_size {
let depth = block * block_size + within;
let q = quant[col * k + depth] as i32 - zero_point;
acc += activation_ref[depth] as f64 * q as f64 * scale;
}
}
if with_bias {
acc += bias_ref[col] as f64;
}
expected[col] = acc as f32;
}
let activation_dev = runtime.alloc_raw(activation_f16.len() * 2).unwrap();
let packed_dev = runtime.alloc_raw(packed.len()).unwrap();
let scales_dev = runtime
.alloc_raw(n * k_blocks * if scales_fp16 { 2 } else { 4 })
.unwrap();
let zp_dev = runtime.alloc_raw(zp_packed.len().max(1)).unwrap();
let bias_dev = runtime.alloc_raw(n * 2).unwrap();
let output_dev = runtime.alloc_raw(n * 2).unwrap();
unsafe {
runtime
.htod(as_bytes(&activation_f16), activation_dev)
.unwrap();
runtime.htod(&packed, packed_dev).unwrap();
if scales_fp16 {
runtime.htod(as_bytes(&scale_f16), scales_dev).unwrap();
} else {
runtime.htod(as_bytes(&scale_f32), scales_dev).unwrap();
}
if explicit_zp {
runtime.htod(&zp_packed, zp_dev).unwrap();
}
if with_bias {
runtime.htod(as_bytes(&bias_f16), bias_dev).unwrap();
}
}
let a_shape = [1usize, k];
let a_strides = [k as i64, 1];
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let bias_shape = [n];
let bias_strides = [1i64];
let y_shape = [1usize, n];
let y_strides = [n as i64, 1];
let scales_dtype = if scales_fp16 {
DataType::Float16
} else {
DataType::Float32
};
let device = DeviceId::cuda(0);
let mut inputs = vec![
TensorView::new(
device_ptr(activation_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
),
TensorView::new(
device_ptr(packed_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
),
TensorView::new(
device_ptr(scales_dev),
scales_dtype,
&scales_shape,
&scales_strides,
device,
),
];
let zp_shape = [n, zp_row_bytes];
let zp_strides = [zp_row_bytes as i64, 1];
let zp_view = TensorView::new(
device_ptr(zp_dev),
DataType::Uint8,
&zp_shape,
&zp_strides,
device,
);
if explicit_zp {
inputs.push(zp_view);
} else if with_bias {
inputs.push(TensorView::absent(DataType::Uint8));
}
if with_bias {
inputs.push(TensorView::absent(DataType::Int32));
inputs.push(TensorView::new(
device_ptr(bias_dev),
DataType::Float16,
&bias_shape,
&bias_strides,
device,
));
}
let mut outputs = [TensorMut::new(
device_ptr_mut(output_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
)];
let kernel = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: false,
decomposed_silu: false,
rmsnorm_prologue: false,
rmsnorm_epsilon: 1e-5,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
let _levers = default_levers_guard();
kernel.run(&inputs, &mut outputs, None).unwrap();
runtime.synchronize().unwrap();
assert!(
kernel.last_call_capture_safe.load(Ordering::Relaxed),
"fp16 decode GEMV must report capture-safe"
);
drop(_levers);
let mut got_f16 = vec![f16::ZERO; n];
unsafe {
runtime
.dtoh(as_bytes_mut(&mut got_f16), output_dev)
.unwrap();
}
unsafe {
runtime.free_raw(activation_dev).unwrap();
runtime.free_raw(packed_dev).unwrap();
runtime.free_raw(scales_dev).unwrap();
runtime.free_raw(zp_dev).unwrap();
runtime.free_raw(bias_dev).unwrap();
runtime.free_raw(output_dev).unwrap();
}
let mut worst_abs = 0.0f32;
let mut worst_rel = 0.0f32;
let mut max_out = 0.0f32;
let mut all_finite = true;
for (g16, e) in got_f16.iter().zip(expected.iter()) {
let g = g16.to_f32();
if !g.is_finite() {
all_finite = false;
}
let abs = (g - e).abs();
let rel = abs / e.abs().max(1e-1);
worst_abs = worst_abs.max(abs);
worst_rel = worst_rel.max(rel);
max_out = max_out.max(e.abs());
}
(worst_abs, worst_rel, max_out, all_finite)
}
#[allow(clippy::too_many_arguments)]
fn run_symmetric_block_raw(
k: usize,
n: usize,
block_size: usize,
scales_fp16: bool,
interleave: bool,
general_splitk: Option<bool>,
pipeline: Option<bool>,
prefetch: Option<bool>,
) -> Option<Vec<f16>> {
let runtime = runtime()?;
if runtime
.require_nvrtc_half_headers("matmul_nbits_gemv_f16")
.is_err()
{
return None;
}
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let mut state = 0x1234_5678_9abc_def0u64;
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let mut activation_f16 = vec![f16::ZERO; k];
for dst_h in activation_f16.iter_mut() {
*dst_h = f16::from_f32(next());
}
let mut quant = vec![0u8; n * k];
for value in quant.iter_mut() {
*value = ((next() * 0.5 + 0.5) * 15.0).round().clamp(0.0, 15.0) as u8;
}
let mut packed = vec![0u8; n * k_blocks * blob_size];
for col in 0..n {
for block in 0..k_blocks {
for pair in 0..blob_size {
let low = quant[col * k + block * block_size + pair * 2] & 15;
let high = quant[col * k + block * block_size + pair * 2 + 1] & 15;
packed[(col * k_blocks + block) * blob_size + pair] = low | (high << 4);
}
}
}
let mut scale_f16 = vec![f16::ZERO; n * k_blocks];
let mut scale_f32 = vec![0.0f32; n * k_blocks];
for i in 0..n * k_blocks {
let raw = 0.015 + 0.01 * (next() * 0.5 + 0.5);
if scales_fp16 {
scale_f16[i] = f16::from_f32(raw);
} else {
scale_f32[i] = raw;
}
}
let activation_dev = runtime.alloc_raw(activation_f16.len() * 2).unwrap();
let packed_dev = runtime.alloc_raw(packed.len()).unwrap();
let scales_dev = runtime
.alloc_raw(n * k_blocks * if scales_fp16 { 2 } else { 4 })
.unwrap();
let output_dev = runtime.alloc_raw(n * 2).unwrap();
unsafe {
runtime
.htod(as_bytes(&activation_f16), activation_dev)
.unwrap();
runtime.htod(&packed, packed_dev).unwrap();
if scales_fp16 {
runtime.htod(as_bytes(&scale_f16), scales_dev).unwrap();
} else {
runtime.htod(as_bytes(&scale_f32), scales_dev).unwrap();
}
}
let a_shape = [1usize, k];
let a_strides = [k as i64, 1];
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let y_shape = [1usize, n];
let y_strides = [n as i64, 1];
let scales_dtype = if scales_fp16 {
DataType::Float16
} else {
DataType::Float32
};
let device = DeviceId::cuda(0);
let inputs = vec![
TensorView::new(
device_ptr(activation_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
),
TensorView::new(
device_ptr(packed_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
),
TensorView::new(
device_ptr(scales_dev),
scales_dtype,
&scales_shape,
&scales_strides,
device,
),
];
let mut outputs = [TensorMut::new(
device_ptr_mut(output_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
)];
let kernel = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: false,
decomposed_silu: false,
rmsnorm_prologue: false,
rmsnorm_epsilon: 1e-5,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
let _guard = LeverEnvGuard::acquire();
unsafe {
match general_splitk {
Some(true) => std::env::set_var("ONNX_GENAI_GENERAL_SPLITK", "on"),
Some(false) => std::env::set_var("ONNX_GENAI_GENERAL_SPLITK", "off"),
None => std::env::remove_var("ONNX_GENAI_GENERAL_SPLITK"),
}
if interleave {
std::env::set_var("ONNX_GENAI_INTERLEAVE_DEQUANT", "1");
} else {
std::env::remove_var("ONNX_GENAI_INTERLEAVE_DEQUANT");
}
match pipeline {
Some(true) => std::env::set_var("ONNX_GENAI_GEMV_PIPELINE", "1"),
Some(false) => std::env::set_var("ONNX_GENAI_GEMV_PIPELINE", "0"),
None => std::env::remove_var("ONNX_GENAI_GEMV_PIPELINE"),
}
match prefetch {
Some(true) => std::env::set_var("ONNX_GENAI_ZP_SPLITK_PREFETCH", "1"),
Some(false) => std::env::set_var("ONNX_GENAI_ZP_SPLITK_PREFETCH", "0"),
None => std::env::remove_var("ONNX_GENAI_ZP_SPLITK_PREFETCH"),
}
}
kernel.run(&inputs, &mut outputs, None).unwrap();
runtime.synchronize().unwrap();
drop(_guard);
let mut got_f16 = vec![f16::ZERO; n];
unsafe {
runtime
.dtoh(as_bytes_mut(&mut got_f16), output_dev)
.unwrap();
runtime.free_raw(activation_dev).unwrap();
runtime.free_raw(scales_dev).unwrap();
runtime.free_raw(output_dev).unwrap();
}
Some(got_f16)
}
#[test]
fn int4_interleaved_dequant_is_bit_identical_to_multicol() {
let mut ran = false;
for (k, n, block_size, scales_fp16) in [
(4096usize, 256usize, 128usize, true),
(4096, 256, 128, false),
(4096, 70, 128, true),
(8192, 512, 128, true),
(4096, 256, 64, true),
] {
let Some(base) = run_symmetric_block_raw(
k,
n,
block_size,
scales_fp16,
false,
Some(false),
None,
None,
) else {
eprintln!(
"skipping interleaved dequant byte-identity test: CUDA runtime/headers \
unavailable"
);
return;
};
let inter = run_symmetric_block_raw(
k,
n,
block_size,
scales_fp16,
true,
Some(false),
None,
None,
)
.unwrap();
ran = true;
let mismatches = base
.iter()
.zip(inter.iter())
.filter(|(b, i)| b.to_bits() != i.to_bits())
.count();
assert_eq!(
mismatches, 0,
"interleaved int4 dequant diverged from multicol at block-{block_size} K={k} \
N={n} scales_fp16={scales_fp16}: {mismatches}/{n} fp16 outputs differ"
);
}
assert!(
ran,
"interleaved dequant byte-identity test did not execute any case"
);
}
#[test]
fn int4_interleaved_dequant_is_bit_identical_to_splitk_wide() {
let mut ran = false;
for (k, n, block_size, scales_fp16) in [
(4096usize, 256usize, 128usize, true),
(4096, 256, 128, false),
(4096, 70, 128, true),
(8192, 128, 128, true),
] {
let Some(base) = run_symmetric_block_raw(
k,
n,
block_size,
scales_fp16,
false,
Some(true),
None,
None,
) else {
eprintln!(
"skipping split-K wide interleaved byte-identity test: CUDA runtime/headers \
unavailable"
);
return;
};
let inter = run_symmetric_block_raw(
k,
n,
block_size,
scales_fp16,
true,
Some(true),
None,
None,
)
.unwrap();
ran = true;
let mismatches = base
.iter()
.zip(inter.iter())
.filter(|(b, i)| b.to_bits() != i.to_bits())
.count();
assert_eq!(
mismatches, 0,
"interleaved int4 dequant diverged from split-K wide at block-{block_size} K={k} \
N={n} scales_fp16={scales_fp16}: {mismatches}/{n} fp16 outputs differ"
);
}
assert!(
ran,
"split-K wide interleaved byte-identity test did not execute any case"
);
}
#[test]
fn interleaved_weights_are_not_shared_between_runtimes() {
const BYTES: usize = 4096;
let Some(first) = crate::test_support::maybe_runtime() else {
eprintln!("skipping interleave runtime-scoping test: no CUDA device");
return;
};
let Some(second) = crate::test_support::maybe_runtime() else {
return;
};
let packed = first.alloc_raw(BYTES).unwrap();
let (built, warm) = ensure_interleaved(&first, packed, BYTES).unwrap();
assert!(!warm, "the first sight of a weight must build it");
assert_eq!(first.interleaved_weight_count(), 1);
assert_eq!(
second.interleaved_weight_count(),
0,
"a second runtime must not see the first runtime's entries"
);
let (other, warm) = ensure_interleaved(&second, packed, BYTES).unwrap();
assert!(
!warm,
"a runtime that has never interleaved this weight must build it"
);
assert_ne!(
other, built,
"two runtimes were served the same interleaved buffer for one address"
);
unsafe { first.free_raw(packed) }.unwrap();
}
#[test]
fn tearing_down_a_runtime_releases_its_interleaved_weights() {
const BYTES: usize = 2048;
let Some(runtime) = crate::test_support::maybe_runtime() else {
eprintln!("skipping interleave teardown test: no CUDA device");
return;
};
let packed = runtime.alloc_raw(BYTES).unwrap();
let (_, warm) = ensure_interleaved(&runtime, packed, BYTES).unwrap();
assert!(!warm);
assert_eq!(runtime.interleaved_weight_count(), 1);
assert!(
ensure_interleaved(&runtime, packed, BYTES).unwrap().1,
"a cached weight must report warm before teardown"
);
runtime.release_interleaved_weights();
assert_eq!(
runtime.interleaved_weight_count(),
0,
"teardown left interleaved weights behind"
);
let (_, warm) = ensure_interleaved(&runtime, packed, BYTES).unwrap();
assert!(
!warm,
"teardown dropped the buffer but left the entry, so a later weight \
at this address would be served a freed pointer"
);
runtime.release_interleaved_weights();
unsafe { runtime.free_raw(packed).unwrap() };
}
#[test]
fn scales_f16_pipeline_is_bit_identical_to_scalar() {
let mut ran = false;
for (k, n) in [
(5120usize, 5120usize),
(5120, 13824),
(896, 896),
(896, 4870),
] {
for scales_fp16 in [true, false] {
assert_eq!(
select_f16_gemv_variant(k, n, 32, scales_fp16, false).variant,
F16GemvVariant::General,
"K={k} N={n} scales_fp16={scales_fp16} must select the General variant"
);
let Some(scalar) =
run_symmetric_block_raw(k, n, 32, scales_fp16, false, None, Some(false), None)
else {
eprintln!(
"skipping scales-fp16 pipeline byte-identity test: CUDA runtime/headers \
unavailable"
);
return;
};
let pipe =
run_symmetric_block_raw(k, n, 32, scales_fp16, false, None, Some(true), None)
.unwrap();
ran = true;
let mismatches = scalar
.iter()
.zip(pipe.iter())
.filter(|(s, p)| s.to_bits() != p.to_bits())
.count();
assert_eq!(
mismatches, 0,
"prefetch-pipelined scales-fp16 GEMV diverged from the scalar entry at \
block-32 K={k} N={n} scales_fp16={scales_fp16}: {mismatches}/{n} fp16 \
outputs differ"
);
}
}
assert!(
ran,
"scales-fp16 pipeline byte-identity test did not execute any case"
);
}
#[test]
fn scales_f16_zp_splitk_prefetch_is_bit_identical_to_splitk() {
let Some(runtime) = runtime() else {
eprintln!("skipping split-K prefetch byte-identity test: CUDA runtime unavailable");
return;
};
let sm = runtime.capabilities().multiprocessor_count();
let max_threads = runtime.capabilities().max_threads_per_block();
let mut ran = false;
for (k, n) in [(5120usize, 8usize), (13824, 16), (5120, 64)] {
for scales_fp16 in [true, false] {
assert!(
use_f16_symmetric_splitk(k, n, sm, max_threads),
"K={k} N={n} must select the symmetric split-K entry (sm={sm})"
);
let Some(plain) =
run_symmetric_block_raw(k, n, 32, scales_fp16, false, None, None, Some(false))
else {
eprintln!(
"skipping split-K prefetch byte-identity test: CUDA runtime/headers \
unavailable"
);
return;
};
let pf =
run_symmetric_block_raw(k, n, 32, scales_fp16, false, None, None, Some(true))
.unwrap();
ran = true;
let mismatches = plain
.iter()
.zip(pf.iter())
.filter(|(s, p)| s.to_bits() != p.to_bits())
.count();
assert_eq!(
mismatches, 0,
"prefetch-pipelined split-K GEMV diverged from the plain split-K entry at \
block-32 K={k} N={n} scales_fp16={scales_fp16}: {mismatches}/{n} fp16 \
outputs differ"
);
}
}
assert!(
ran,
"split-K prefetch byte-identity test did not execute any case"
);
}
#[test]
#[ignore = "requires an SM80+ CUDA device"]
fn marlin_m_gt_1_op_parity_and_capture_safety() {
let Some(runtime) = runtime() else {
eprintln!("skipping Marlin M>1 op parity: CUDA runtime unavailable");
return;
};
if runtime
.require_nvrtc_half_headers("marlin_m_gt_1_op_parity")
.is_err()
{
eprintln!("skipping Marlin M>1 op parity: fp16 NVRTC headers unavailable");
return;
}
if !marlin_gemm::device_supports_marlin(runtime.capabilities().compute_capability()) {
eprintln!("skipping Marlin M>1 op parity: device is pre-SM80");
return;
}
let m = 16usize;
let k = 4096usize;
let n = 70usize;
let block_size = 128usize;
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let zp_row_bytes = k_blocks.div_ceil(2);
let mut state = 0x1234_5678_9abc_def0u64;
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let mut activation_f16 = vec![f16::ZERO; m * k];
let mut activation_ref = vec![0.0f32; m * k];
for (h, f) in activation_f16.iter_mut().zip(activation_ref.iter_mut()) {
let value = f16::from_f32(next());
*h = value;
*f = value.to_f32();
}
let mut quant = vec![0u8; n * k];
for value in quant.iter_mut() {
*value = ((next() * 0.5 + 0.5) * 15.0).round().clamp(0.0, 15.0) as u8;
}
let mut packed = vec![0u8; n * k_blocks * blob_size];
for col in 0..n {
for block in 0..k_blocks {
for pair in 0..blob_size {
let low = quant[col * k + block * block_size + pair * 2] & 15;
let high = quant[col * k + block * block_size + pair * 2 + 1] & 15;
packed[(col * k_blocks + block) * blob_size + pair] = low | (high << 4);
}
}
}
let mut zp_codes = vec![0i32; n * k_blocks];
let mut zp_packed = vec![0u8; n * zp_row_bytes];
for code in zp_codes.iter_mut() {
*code = ((next() * 0.5 + 0.5) * 15.0).round().clamp(0.0, 15.0) as i32;
}
for col in 0..n {
for block in 0..k_blocks {
let code = (zp_codes[col * k_blocks + block] & 15) as u8;
let byte = &mut zp_packed[col * zp_row_bytes + block / 2];
if block & 1 == 0 {
*byte = (*byte & 0xf0) | code;
} else {
*byte = (*byte & 0x0f) | (code << 4);
}
}
}
let mut scale_ref = vec![0.0f32; n * k_blocks];
let mut scale_f16 = vec![f16::ZERO; n * k_blocks];
for i in 0..n * k_blocks {
let h = f16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5));
scale_f16[i] = h;
scale_ref[i] = h.to_f32();
}
let mut expected = vec![0.0f32; m * n];
for row in 0..m {
for col in 0..n {
let mut acc = 0.0f64;
for block in 0..k_blocks {
let scale = scale_ref[col * k_blocks + block] as f64;
let zp = zp_codes[col * k_blocks + block];
for within in 0..block_size {
let depth = block * block_size + within;
let q = quant[col * k + depth] as i32 - zp;
acc += activation_ref[row * k + depth] as f64 * q as f64 * scale;
}
}
expected[row * n + col] = acc as f32;
}
}
let activation_dev = runtime.alloc_raw(activation_f16.len() * 2).unwrap();
let packed_dev = runtime.alloc_raw(packed.len()).unwrap();
let scales_dev = runtime.alloc_raw(scale_f16.len() * 2).unwrap();
let zp_dev = runtime.alloc_raw(zp_packed.len()).unwrap();
let output_dev = runtime.alloc_raw(m * n * 2).unwrap();
unsafe {
runtime
.htod(as_bytes(&activation_f16), activation_dev)
.unwrap();
runtime.htod(&packed, packed_dev).unwrap();
runtime.htod(as_bytes(&scale_f16), scales_dev).unwrap();
runtime.htod(&zp_packed, zp_dev).unwrap();
}
let device = DeviceId::cuda(0);
let a_shape = [m, k];
let a_strides = [k as i64, 1];
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let zp_shape = [n, zp_row_bytes];
let zp_strides = [zp_row_bytes as i64, 1];
let y_shape = [m, n];
let y_strides = [n as i64, 1];
let inputs = vec![
TensorView::new(
device_ptr(activation_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
),
TensorView::new(
device_ptr(packed_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
),
TensorView::new(
device_ptr(scales_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
),
TensorView::new(
device_ptr(zp_dev),
DataType::Uint8,
&zp_shape,
&zp_strides,
device,
),
];
let kernel = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: false,
decomposed_silu: false,
rmsnorm_prologue: false,
rmsnorm_epsilon: 1e-5,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
unsafe {
std::env::set_var("ONNX_GENAI_MARLIN_M_GT_1", "1");
}
let run_once = |capture_expectation: &str| -> Vec<f16> {
let mut outputs = [TensorMut::new(
device_ptr_mut(output_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
)];
kernel.run(&inputs, &mut outputs, None).unwrap();
runtime.synchronize().unwrap();
let mut got = vec![f16::ZERO; m * n];
unsafe {
runtime.dtoh(as_bytes_mut(&mut got), output_dev).unwrap();
}
eprintln!(
"Marlin M>1 op {capture_expectation}: capture_safe={}",
kernel.last_call_capture_safe.load(Ordering::Relaxed)
);
got
};
let got_cold = run_once("cold");
assert!(
!kernel.last_call_capture_safe.load(Ordering::Relaxed),
"cold Marlin call performs the repack allocation and must report NOT capture-safe"
);
let got_warm = run_once("warm");
assert!(
kernel.last_call_capture_safe.load(Ordering::Relaxed),
"warm Marlin call reuses the cached repack and must report capture-safe"
);
assert_eq!(
got_cold, got_warm,
"warm replay of the static-grid Marlin kernel must be byte-identical to the cold run"
);
unsafe {
std::env::remove_var("ONNX_GENAI_MARLIN_M_GT_1");
}
unsafe {
runtime.free_raw(activation_dev).unwrap();
runtime.free_raw(packed_dev).unwrap();
runtime.free_raw(scales_dev).unwrap();
runtime.free_raw(zp_dev).unwrap();
runtime.free_raw(output_dev).unwrap();
}
let mut worst_abs = 0.0f32;
let mut max_out = 0.0f32;
for (g16, e) in got_warm.iter().zip(expected.iter()) {
let g = g16.to_f32();
assert!(g.is_finite(), "Marlin M>1 output must be finite");
worst_abs = worst_abs.max((g - e).abs());
max_out = max_out.max(e.abs());
}
let tol = 2e-2 * max_out.max(1e-3);
eprintln!(
"Marlin M>1 op parity: worst_abs={worst_abs:.5}, max_out={max_out:.5}, tol={tol:.5}"
);
assert!(
worst_abs <= tol,
"Marlin M>1 op output exceeds tolerance: worst_abs={worst_abs} > tol={tol}"
);
}
#[test]
#[ignore = "requires an SM80+ CUDA device"]
fn dequant_f16_cublas_m_gt_1_op_parity() {
let Some(runtime) = runtime() else {
eprintln!("skipping dequant-f16 cuBLASLt parity: CUDA runtime unavailable");
return;
};
if runtime
.require_nvrtc_half_headers("dequant_f16_cublas_m_gt_1_op_parity")
.is_err()
{
eprintln!("skipping dequant-f16 cuBLASLt parity: fp16 NVRTC headers unavailable");
return;
}
let m = 32usize;
let k = 1024usize;
let n = 96usize;
let block_size = 32usize;
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let zp_row_bytes = k_blocks.div_ceil(2);
let mut state = 0x0f1e_2d3c_4b5a_6978u64;
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let mut activation_f16 = vec![f16::ZERO; m * k];
let mut activation_ref = vec![0.0f32; m * k];
for (h, f) in activation_f16.iter_mut().zip(activation_ref.iter_mut()) {
let value = f16::from_f32(next());
*h = value;
*f = value.to_f32();
}
let mut quant = vec![0u8; n * k];
for value in quant.iter_mut() {
*value = ((next() * 0.5 + 0.5) * 15.0).round().clamp(0.0, 15.0) as u8;
}
let mut packed = vec![0u8; n * k_blocks * blob_size];
for col in 0..n {
for block in 0..k_blocks {
for pair in 0..blob_size {
let low = quant[col * k + block * block_size + pair * 2] & 15;
let high = quant[col * k + block * block_size + pair * 2 + 1] & 15;
packed[(col * k_blocks + block) * blob_size + pair] = low | (high << 4);
}
}
}
let mut zp_codes = vec![0i32; n * k_blocks];
let mut zp_packed = vec![0u8; n * zp_row_bytes];
for code in zp_codes.iter_mut() {
*code = ((next() * 0.5 + 0.5) * 15.0).round().clamp(0.0, 15.0) as i32;
}
for col in 0..n {
for block in 0..k_blocks {
let code = (zp_codes[col * k_blocks + block] & 15) as u8;
let byte = &mut zp_packed[col * zp_row_bytes + block / 2];
if block & 1 == 0 {
*byte = (*byte & 0xf0) | code;
} else {
*byte = (*byte & 0x0f) | (code << 4);
}
}
}
let mut scale_ref = vec![0.0f32; n * k_blocks];
let mut scale_f16 = vec![f16::ZERO; n * k_blocks];
for i in 0..n * k_blocks {
let h = f16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5));
scale_f16[i] = h;
scale_ref[i] = h.to_f32();
}
let mut expected = vec![0.0f32; m * n];
for row in 0..m {
for col in 0..n {
let mut acc = 0.0f64;
for block in 0..k_blocks {
let scale = scale_ref[col * k_blocks + block] as f64;
let zp = zp_codes[col * k_blocks + block];
for within in 0..block_size {
let depth = block * block_size + within;
let q = quant[col * k + depth] as i32 - zp;
acc += activation_ref[row * k + depth] as f64 * q as f64 * scale;
}
}
expected[row * n + col] = acc as f32;
}
}
let activation_dev = runtime.alloc_raw(activation_f16.len() * 2).unwrap();
let packed_dev = runtime.alloc_raw(packed.len()).unwrap();
let scales_dev = runtime.alloc_raw(scale_f16.len() * 2).unwrap();
let zp_dev = runtime.alloc_raw(zp_packed.len()).unwrap();
let output_dev = runtime.alloc_raw(m * n * 2).unwrap();
unsafe {
runtime
.htod(as_bytes(&activation_f16), activation_dev)
.unwrap();
runtime.htod(&packed, packed_dev).unwrap();
runtime.htod(as_bytes(&scale_f16), scales_dev).unwrap();
runtime.htod(&zp_packed, zp_dev).unwrap();
}
let device = DeviceId::cuda(0);
let a_shape = [m, k];
let a_strides = [k as i64, 1];
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let zp_shape = [n, zp_row_bytes];
let zp_strides = [zp_row_bytes as i64, 1];
let y_shape = [m, n];
let y_strides = [n as i64, 1];
let inputs = vec![
TensorView::new(
device_ptr(activation_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
),
TensorView::new(
device_ptr(packed_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
),
TensorView::new(
device_ptr(scales_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
),
TensorView::new(
device_ptr(zp_dev),
DataType::Uint8,
&zp_shape,
&zp_strides,
device,
),
];
let kernel = zc_kernel(&runtime, k, n, block_size);
let metadata = [TensorMetadata {
dtype: DataType::Float16,
shape: &a_shape,
present: true,
}];
let requirement = kernel.workspace_requirement_for(&metadata).unwrap();
let workspace_bytes = requirement.bytes.max(1) as usize;
let workspace_dev = runtime.alloc_raw(workspace_bytes).unwrap();
let run_once = |dequant_f16: bool| -> Vec<f16> {
unsafe {
std::env::set_var(
"ONNX_GENAI_DEQUANT_F16_GEMM",
if dequant_f16 { "1" } else { "0" },
);
std::env::set_var(
"ONNX_GENAI_MARLIN_M_GT_1",
if dequant_f16 { "0" } else { "1" },
);
}
let mut outputs = [TensorMut::new(
device_ptr_mut(output_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
)];
let workspace = Some(WorkspaceView::new(
device_ptr_mut(workspace_dev),
workspace_bytes,
));
kernel.run(&inputs, &mut outputs, workspace).unwrap();
runtime.synchronize().unwrap();
let mut got = vec![f16::ZERO; m * n];
unsafe {
runtime.dtoh(as_bytes_mut(&mut got), output_dev).unwrap();
}
got
};
assert!(
requirement.bytes > 0,
"the dequant-f16 path must declare a cuBLASLt workspace; a zero requirement means \
the launch would fail and silently fall back to Marlin"
);
let got_dequant = run_once(true);
let got_marlin = run_once(false);
unsafe {
std::env::remove_var("ONNX_GENAI_DEQUANT_F16_GEMM");
std::env::remove_var("ONNX_GENAI_MARLIN_M_GT_1");
}
unsafe {
runtime.free_raw(activation_dev).unwrap();
runtime.free_raw(packed_dev).unwrap();
runtime.free_raw(scales_dev).unwrap();
runtime.free_raw(zp_dev).unwrap();
runtime.free_raw(output_dev).unwrap();
runtime.free_raw(workspace_dev).unwrap();
}
let mut worst_abs = 0.0f32;
let mut worst_vs_marlin = 0.0f32;
let mut max_out = 0.0f32;
for ((d16, m16), e) in got_dequant
.iter()
.zip(got_marlin.iter())
.zip(expected.iter())
{
let d = d16.to_f32();
assert!(d.is_finite(), "dequant-f16 output must be finite");
worst_abs = worst_abs.max((d - e).abs());
worst_vs_marlin = worst_vs_marlin.max((d - m16.to_f32()).abs());
max_out = max_out.max(e.abs());
}
let tol = 2e-2 * max_out.max(1e-3);
eprintln!(
"dequant-f16 cuBLASLt parity: worst_abs={worst_abs:.5} vs_marlin={worst_vs_marlin:.5} \
max_out={max_out:.5} tol={tol:.5} workspace={workspace_bytes}B"
);
assert!(
worst_abs <= tol,
"dequant-f16 cuBLASLt output exceeds oracle tolerance: {worst_abs} > {tol}"
);
assert!(
worst_vs_marlin <= tol,
"dequant-f16 cuBLASLt output diverged from Marlin: {worst_vs_marlin} > {tol}"
);
}
fn zc_env_usize(key: &str, default: usize) -> usize {
std::env::var(key)
.ok()
.and_then(|v| v.trim().parse().ok())
.unwrap_or(default)
}
#[allow(clippy::type_complexity)]
fn zc_make_inputs(
k: usize,
n: usize,
block_size: usize,
) -> (Vec<f16>, Vec<u8>, Vec<f16>, Vec<u8>) {
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let zp_row_bytes = k_blocks.div_ceil(2);
let mut state = 0x9e37_79b9_7f4a_7c15u64;
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let mut activation = vec![f16::ZERO; k];
for slot in activation.iter_mut() {
*slot = f16::from_f32(next());
}
let mut packed = vec![0u8; n * k_blocks * blob_size];
for byte in packed.iter_mut() {
let low = ((next() * 0.5 + 0.5) * 15.0).round().clamp(0.0, 15.0) as u8;
let high = ((next() * 0.5 + 0.5) * 15.0).round().clamp(0.0, 15.0) as u8;
*byte = low | (high << 4);
}
let mut scales = vec![f16::ZERO; n * k_blocks];
for slot in scales.iter_mut() {
*slot = f16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5));
}
let mut zp_packed = vec![0u8; n * zp_row_bytes];
for col in 0..n {
for block in 0..k_blocks {
let code = ((next() * 0.5 + 0.5) * 15.0).round().clamp(0.0, 15.0) as u8;
let byte = &mut zp_packed[col * zp_row_bytes + block / 2];
if block & 1 == 0 {
*byte = (*byte & 0xf0) | code;
} else {
*byte = (*byte & 0x0f) | (code << 4);
}
}
}
(activation, packed, scales, zp_packed)
}
fn zc_kernel(
runtime: &Arc<CudaRuntime>,
k: usize,
n: usize,
block_size: usize,
) -> MatMulNBitsKernel {
MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: false,
decomposed_silu: false,
rmsnorm_prologue: false,
rmsnorm_epsilon: 1e-5,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
}
}
#[test]
#[ignore]
fn zerocopy_kernel_host_mapped_vs_vram_bandwidth() {
use cudarc::driver::result::event;
use cudarc::driver::sys;
let Some(runtime) = runtime() else {
eprintln!("skipping zero-copy kernel probe: CUDA runtime unavailable");
return;
};
if runtime
.require_nvrtc_half_headers("zero-copy kernel probe")
.is_err()
{
eprintln!("skipping zero-copy kernel probe: fp16 NVRTC headers unavailable");
return;
}
runtime.bind().unwrap();
const CU_MEMHOSTREGISTER_DEVICEMAP: u32 = 0x02;
const CU_MEMHOSTREGISTER_READ_ONLY: u32 = 0x08;
let k = zc_env_usize("ZC_K", 5120);
let block_size = zc_env_usize("ZC_BLOCK", 32);
let reps = zc_env_usize("ZC_REPS", 7).max(3);
assert!(
k.is_multiple_of(block_size),
"K must be a multiple of block_size"
);
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let zp_row_bytes = k_blocks.div_ceil(2);
let row_bytes = k_blocks * blob_size;
let sizes_mib: Vec<usize> = std::env::var("ZC_SIZES_MIB")
.ok()
.map(|s| s.split(',').filter_map(|x| x.trim().parse().ok()).collect())
.unwrap_or_else(|| vec![12, 64, 128, 256, 384]);
let caps = runtime.capabilities();
println!(
"\n=== #864 zero-copy REAL-kernel bandwidth probe ===\n\
SMs={} K={k} block_size={block_size} k_blocks={k_blocks} row_bytes/col={row_bytes} best-of-{reps}\n\
(packed = int4 weights; only the packed pointer differs VRAM vs host-mapped)\n",
caps.multiprocessor_count(),
);
println!(
"{:>10} {:>9} {:>12} {:>12} {:>8} {:>10}",
"packedMiB", "N", "VRAM GB/s", "host GB/s", "ratio", "match"
);
for mib in sizes_mib {
let target = mib * (1usize << 20);
let n = (target / row_bytes).max(1);
let packed_bytes = n * row_bytes;
let packed_mib = packed_bytes as f64 / (1u64 << 20) as f64;
let (activation, packed, scales, zp_packed) = zc_make_inputs(k, n, block_size);
let activation_dev = runtime.alloc_raw(activation.len() * 2).unwrap();
let scales_dev = runtime.alloc_raw(scales.len() * 2).unwrap();
let zp_dev = runtime.alloc_raw(zp_packed.len().max(1)).unwrap();
let output_dev = runtime.alloc_raw(n * 2).unwrap();
let packed_vram = runtime.alloc_raw(packed_bytes).unwrap();
unsafe {
runtime.htod(as_bytes(&activation), activation_dev).unwrap();
runtime.htod(as_bytes(&scales), scales_dev).unwrap();
runtime.htod(&zp_packed, zp_dev).unwrap();
runtime.htod(&packed, packed_vram).unwrap();
}
let host_ptr = packed.as_ptr() as *mut c_void;
unsafe {
sys::cuMemHostRegister_v2(
host_ptr,
packed_bytes,
CU_MEMHOSTREGISTER_DEVICEMAP | CU_MEMHOSTREGISTER_READ_ONLY,
)
.result()
.expect("cuMemHostRegister(DEVICEMAP|READ_ONLY)");
}
let mut packed_host_dptr: CUdeviceptr = 0;
unsafe {
sys::cuMemHostGetDevicePointer_v2(&mut packed_host_dptr, host_ptr, 0)
.result()
.expect("cuMemHostGetDevicePointer");
}
let kernel = zc_kernel(&runtime, k, n, block_size);
let a_shape = [1usize, k];
let a_strides = [k as i64, 1];
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let s_shape = [n, k_blocks];
let s_strides = [k_blocks as i64, 1];
let zp_shape = [n, zp_row_bytes];
let zp_strides = [zp_row_bytes as i64, 1];
let y_shape = [1usize, n];
let y_strides = [n as i64, 1];
let device = DeviceId::cuda(0);
let make_inputs = |packed_ptr: CUdeviceptr| {
vec![
TensorView::new(
device_ptr(activation_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
),
TensorView::new(
device_ptr(packed_ptr),
DataType::Uint8,
&b_shape,
&b_strides,
device,
),
TensorView::new(
device_ptr(scales_dev),
DataType::Float16,
&s_shape,
&s_strides,
device,
),
TensorView::new(
device_ptr(zp_dev),
DataType::Uint8,
&zp_shape,
&zp_strides,
device,
),
]
};
let time_arm = |packed_ptr: CUdeviceptr| -> Vec<f32> {
let inputs = make_inputs(packed_ptr);
let mut outputs = [TensorMut::new(
device_ptr_mut(output_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
)];
kernel.run(&inputs, &mut outputs, None).unwrap();
runtime.synchronize().unwrap();
let mut times = Vec::with_capacity(reps);
for _ in 0..reps {
let start = event::create(sys::CUevent_flags::CU_EVENT_DEFAULT).unwrap();
let end = event::create(sys::CUevent_flags::CU_EVENT_DEFAULT).unwrap();
unsafe {
event::record(start, runtime.stream_ptr()).unwrap();
kernel.run(&inputs, &mut outputs, None).unwrap();
event::record(end, runtime.stream_ptr()).unwrap();
event::synchronize(end).unwrap();
times.push(event::elapsed(start, end).unwrap());
event::destroy(start).ok();
event::destroy(end).ok();
}
}
times
};
let vram_times = time_arm(packed_vram);
let mut got_vram = vec![f16::ZERO; n];
unsafe {
runtime
.dtoh(as_bytes_mut(&mut got_vram), output_dev)
.unwrap()
};
let host_times = time_arm(packed_host_dptr);
let mut got_host = vec![f16::ZERO; n];
unsafe {
runtime
.dtoh(as_bytes_mut(&mut got_host), output_dev)
.unwrap()
};
let best_ms = |t: &[f32]| t.iter().cloned().fold(f32::INFINITY, f32::min);
let med_ms = |t: &[f32]| {
let mut s = t.to_vec();
s.sort_by(|a, b| a.partial_cmp(b).unwrap());
s[s.len() / 2]
};
let gbps = |ms: f32| (packed_bytes as f64) / (ms as f64 / 1e3) / 1e9;
let vram_gbps = gbps(best_ms(&vram_times));
let host_gbps = gbps(best_ms(&host_times));
let bit_match = got_vram
.iter()
.zip(got_host.iter())
.all(|(a, b)| a.to_bits() == b.to_bits());
println!(
"{:>10.1} {:>9} {:>12.2} {:>12.2} {:>7.2}x {:>10}",
packed_mib,
n,
vram_gbps,
host_gbps,
vram_gbps / host_gbps,
if bit_match { "yes" } else { "NO" }
);
println!(
" spread: VRAM best={:.4} med={:.4} ms host best={:.4} med={:.4} ms",
best_ms(&vram_times),
med_ms(&vram_times),
best_ms(&host_times),
med_ms(&host_times),
);
unsafe {
sys::cuMemHostUnregister(host_ptr)
.result()
.expect("cuMemHostUnregister");
}
unsafe {
runtime.free_raw(activation_dev).unwrap();
runtime.free_raw(scales_dev).unwrap();
runtime.free_raw(zp_dev).unwrap();
runtime.free_raw(output_dev).unwrap();
runtime.free_raw(packed_vram).unwrap();
}
drop(packed);
assert!(
bit_match,
"zero-copy host-mapped read produced a DIFFERENT result than the VRAM read \
(packed={packed_mib:.1} MiB, N={n}) — correctness failure"
);
}
}
#[test]
#[ignore]
fn zerocopy_kernel_capture_with_host_mapped_pointer() {
use cudarc::driver::sys;
let Some(runtime) = runtime() else {
eprintln!("skipping capture probe: CUDA runtime unavailable");
return;
};
if runtime.require_nvrtc_half_headers("capture probe").is_err() {
eprintln!("skipping capture probe: fp16 NVRTC headers unavailable");
return;
}
runtime.bind().unwrap();
const CU_MEMHOSTREGISTER_DEVICEMAP: u32 = 0x02;
const CU_MEMHOSTREGISTER_READ_ONLY: u32 = 0x08;
let k = zc_env_usize("ZC_K", 5120);
let block_size = zc_env_usize("ZC_BLOCK", 32);
let n = zc_env_usize("ZC_CAPTURE_N", 4608);
assert!(k.is_multiple_of(block_size));
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let zp_row_bytes = k_blocks.div_ceil(2);
let packed_bytes = n * k_blocks * blob_size;
let (activation, packed, scales, zp_packed) = zc_make_inputs(k, n, block_size);
let activation_dev = runtime.alloc_raw(activation.len() * 2).unwrap();
let scales_dev = runtime.alloc_raw(scales.len() * 2).unwrap();
let zp_dev = runtime.alloc_raw(zp_packed.len().max(1)).unwrap();
let output_dev = runtime.alloc_raw(n * 2).unwrap();
unsafe {
runtime.htod(as_bytes(&activation), activation_dev).unwrap();
runtime.htod(as_bytes(&scales), scales_dev).unwrap();
runtime.htod(&zp_packed, zp_dev).unwrap();
}
let host_ptr = packed.as_ptr() as *mut c_void;
unsafe {
sys::cuMemHostRegister_v2(
host_ptr,
packed_bytes,
CU_MEMHOSTREGISTER_DEVICEMAP | CU_MEMHOSTREGISTER_READ_ONLY,
)
.result()
.expect("cuMemHostRegister(DEVICEMAP|READ_ONLY)");
}
let mut packed_host_dptr: CUdeviceptr = 0;
unsafe {
sys::cuMemHostGetDevicePointer_v2(&mut packed_host_dptr, host_ptr, 0)
.result()
.expect("cuMemHostGetDevicePointer");
}
let kernel = zc_kernel(&runtime, k, n, block_size);
let a_shape = [1usize, k];
let a_strides = [k as i64, 1];
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let s_shape = [n, k_blocks];
let s_strides = [k_blocks as i64, 1];
let zp_shape = [n, zp_row_bytes];
let zp_strides = [zp_row_bytes as i64, 1];
let y_shape = [1usize, n];
let y_strides = [n as i64, 1];
let device = DeviceId::cuda(0);
let inputs = vec![
TensorView::new(
device_ptr(activation_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
),
TensorView::new(
device_ptr(packed_host_dptr),
DataType::Uint8,
&b_shape,
&b_strides,
device,
),
TensorView::new(
device_ptr(scales_dev),
DataType::Float16,
&s_shape,
&s_strides,
device,
),
TensorView::new(
device_ptr(zp_dev),
DataType::Uint8,
&zp_shape,
&zp_strides,
device,
),
];
let mut outputs = [TensorMut::new(
device_ptr_mut(output_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
)];
kernel.run(&inputs, &mut outputs, None).unwrap();
runtime.synchronize().unwrap();
let mut eager = vec![f16::ZERO; n];
unsafe { runtime.dtoh(as_bytes_mut(&mut eager), output_dev).unwrap() };
unsafe {
runtime
.htod(as_bytes(&vec![f16::ONE; n]), output_dev)
.unwrap()
};
let captured = runtime
.begin_graph_capture(&[&kernel as &dyn Kernel])
.is_ok();
let mut replay_match = false;
if captured {
kernel.run(&inputs, &mut outputs, None).unwrap();
match runtime.end_graph_capture() {
Ok(()) => {
for _ in 0..4 {
runtime.replay_graph().unwrap();
}
runtime.synchronize().unwrap();
let mut got = vec![f16::ZERO; n];
unsafe { runtime.dtoh(as_bytes_mut(&mut got), output_dev).unwrap() };
replay_match = got
.iter()
.zip(eager.iter())
.all(|(a, b)| a.to_bits() == b.to_bits());
runtime.reset_graph().ok();
println!(
"\n=== #864 capture-with-host-mapped-pointer probe ===\n\
capture: SUPPORTED; replay output bit-identical to eager: {}",
if replay_match { "YES" } else { "NO" }
);
}
Err(e) => {
runtime.abort_graph_capture().ok();
println!(
"\n=== #864 capture-with-host-mapped-pointer probe ===\n\
capture: end_graph_capture FAILED: {e:?}"
);
}
}
} else {
println!(
"\n=== #864 capture-with-host-mapped-pointer probe ===\n\
capture: begin_graph_capture declined (kernel/audit not capturable in this build)"
);
}
unsafe {
sys::cuMemHostUnregister(host_ptr)
.result()
.expect("cuMemHostUnregister");
}
unsafe {
runtime.free_raw(activation_dev).unwrap();
runtime.free_raw(scales_dev).unwrap();
runtime.free_raw(zp_dev).unwrap();
runtime.free_raw(output_dev).unwrap();
}
drop(packed);
if captured {
assert!(
replay_match,
"graph replay with a host-mapped weight pointer produced a different result \
than the eager host-mapped run"
);
}
}
fn run_int8_parity_dims(
k: usize,
n: usize,
block_size: usize,
scales_fp16: bool,
with_bias: bool,
explicit_zp: bool,
) -> (f32, f32, f32, bool) {
let Some(runtime) = runtime() else {
eprintln!("skipping MatMulNBits int8 fp16 GEMV parity test: CUDA runtime unavailable");
return (0.0, 0.0, 0.0, true);
};
if runtime
.require_nvrtc_half_headers("matmul_nbits_gemv_f16")
.is_err()
{
eprintln!(
"skipping MatMulNBits int8 fp16 GEMV parity test: fp16 NVRTC headers unavailable"
);
return (0.0, 0.0, 0.0, true);
}
let k_blocks = k / block_size;
let blob_size = block_size;
let mut state = 0x243f_6a88_85a3_08d3u64;
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let mut activation_f16 = vec![f16::ZERO; k];
let mut activation_ref = vec![0.0f32; k];
for (dst_h, dst_f) in activation_f16.iter_mut().zip(activation_ref.iter_mut()) {
let h = f16::from_f32(next());
*dst_h = h;
*dst_f = h.to_f32();
}
let mut quant = vec![0u8; n * k];
for value in quant.iter_mut() {
*value = ((next() * 0.5 + 0.5) * 255.0).round().clamp(0.0, 255.0) as u8;
}
let mut packed = vec![0u8; n * k_blocks * blob_size];
for col in 0..n {
for block in 0..k_blocks {
for within in 0..block_size {
packed[(col * k_blocks + block) * blob_size + within] =
quant[col * k + block * block_size + within];
}
}
}
let mut zero_points = vec![0u8; n * k_blocks];
if explicit_zp {
for zp in zero_points.iter_mut() {
*zp = ((next() * 0.5 + 0.5) * 255.0).round().clamp(0.0, 255.0) as u8;
}
}
let zp_ref = |col: usize, block: usize| -> i32 {
if explicit_zp {
zero_points[col * k_blocks + block] as i32
} else {
128
}
};
let mut scale_ref = vec![0.0f32; n * k_blocks];
let mut scale_f16 = vec![f16::ZERO; n * k_blocks];
let mut scale_f32 = vec![0.0f32; n * k_blocks];
for i in 0..n * k_blocks {
let raw = 0.015 + 0.01 * (next() * 0.5 + 0.5);
if scales_fp16 {
let h = f16::from_f32(raw);
scale_f16[i] = h;
scale_ref[i] = h.to_f32();
} else {
scale_f32[i] = raw;
scale_ref[i] = raw;
}
}
let mut bias_f16 = vec![f16::ZERO; n];
let mut bias_ref = vec![0.0f32; n];
if with_bias {
for (h, f) in bias_f16.iter_mut().zip(bias_ref.iter_mut()) {
let value = f16::from_f32(next());
*h = value;
*f = value.to_f32();
}
}
let mut expected = vec![0.0f32; n];
for col in 0..n {
let mut acc = 0.0f64;
for block in 0..k_blocks {
let scale = scale_ref[col * k_blocks + block] as f64;
let zp = zp_ref(col, block);
for within in 0..block_size {
let depth = block * block_size + within;
let q = quant[col * k + depth] as i32 - zp;
acc += activation_ref[depth] as f64 * q as f64 * scale;
}
}
if with_bias {
acc += bias_ref[col] as f64;
}
expected[col] = acc as f32;
}
let activation_dev = runtime.alloc_raw(activation_f16.len() * 2).unwrap();
let packed_dev = runtime.alloc_raw(packed.len()).unwrap();
let scales_dev = runtime
.alloc_raw(n * k_blocks * if scales_fp16 { 2 } else { 4 })
.unwrap();
let zp_dev = runtime.alloc_raw(zero_points.len()).unwrap();
let bias_dev = runtime.alloc_raw(n * 2).unwrap();
let output_dev = runtime.alloc_raw(n * 2).unwrap();
unsafe {
runtime
.htod(as_bytes(&activation_f16), activation_dev)
.unwrap();
runtime.htod(&packed, packed_dev).unwrap();
if scales_fp16 {
runtime.htod(as_bytes(&scale_f16), scales_dev).unwrap();
} else {
runtime.htod(as_bytes(&scale_f32), scales_dev).unwrap();
}
if explicit_zp {
runtime.htod(&zero_points, zp_dev).unwrap();
}
if with_bias {
runtime.htod(as_bytes(&bias_f16), bias_dev).unwrap();
}
}
let a_shape = [1usize, k];
let a_strides = [k as i64, 1];
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let zp_shape = [n, k_blocks];
let zp_strides = [k_blocks as i64, 1];
let bias_shape = [n];
let bias_strides = [1i64];
let y_shape = [1usize, n];
let y_strides = [n as i64, 1];
let scales_dtype = if scales_fp16 {
DataType::Float16
} else {
DataType::Float32
};
let device = DeviceId::cuda(0);
let mut inputs = vec![
TensorView::new(
device_ptr(activation_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
),
TensorView::new(
device_ptr(packed_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
),
TensorView::new(
device_ptr(scales_dev),
scales_dtype,
&scales_shape,
&scales_strides,
device,
),
];
if explicit_zp || with_bias {
inputs.push(if explicit_zp {
TensorView::new(
device_ptr(zp_dev),
DataType::Uint8,
&zp_shape,
&zp_strides,
device,
)
} else {
TensorView::absent(DataType::Uint8)
});
}
if with_bias {
inputs.push(TensorView::absent(DataType::Int32));
inputs.push(TensorView::new(
device_ptr(bias_dev),
DataType::Float16,
&bias_shape,
&bias_strides,
device,
));
}
let mut outputs = [TensorMut::new(
device_ptr_mut(output_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
)];
let kernel = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 8,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: false,
decomposed_silu: false,
rmsnorm_prologue: false,
rmsnorm_epsilon: 1e-5,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
kernel.run(&inputs, &mut outputs, None).unwrap();
runtime.synchronize().unwrap();
assert!(
kernel.last_call_capture_safe.load(Ordering::Relaxed),
"int8 decode GEMV must report capture-safe"
);
let mut got_f16 = vec![f16::ZERO; n];
unsafe {
runtime
.dtoh(as_bytes_mut(&mut got_f16), output_dev)
.unwrap();
}
unsafe {
runtime.free_raw(activation_dev).unwrap();
runtime.free_raw(packed_dev).unwrap();
runtime.free_raw(scales_dev).unwrap();
runtime.free_raw(zp_dev).unwrap();
runtime.free_raw(bias_dev).unwrap();
runtime.free_raw(output_dev).unwrap();
}
let mut worst_abs = 0.0f32;
let mut worst_rel = 0.0f32;
let mut max_out = 0.0f32;
let mut all_finite = true;
for (g16, e) in got_f16.iter().zip(expected.iter()) {
let g = g16.to_f32();
if !g.is_finite() {
all_finite = false;
}
let abs = (g - e).abs();
let rel = abs / e.abs().max(1e-1);
worst_abs = worst_abs.max(abs);
worst_rel = worst_rel.max(rel);
max_out = max_out.max(e.abs());
}
(worst_abs, worst_rel, max_out, all_finite)
}
#[test]
fn fp16_down_projection_is_bit_exact_to_staged_kernel() {
let Some(runtime) = runtime() else {
eprintln!("skipping down-projection GEMV parity test: CUDA runtime unavailable");
return;
};
if runtime
.require_nvrtc_half_headers("matmul_nbits_gemv_f16")
.is_err()
{
eprintln!("skipping down-projection GEMV parity test: fp16 NVRTC headers unavailable");
return;
}
for (k, n) in [
(QWEN_DOWN_K, QWEN_DOWN_N),
(5632usize, 2048usize),
(16_384usize, 8192usize),
] {
assert_eq!(
select_f16_gemv_variant(k, n, 32, true, false).variant,
F16GemvVariant::DownProjection,
"shape K={k}, N={n} must select the down variant under test"
);
let block_size = 32usize;
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let activation: Vec<f16> = (0..k)
.map(|i| f16::from_f32(((i * 17 % 257) as f32 - 128.0) / 128.0))
.collect();
let packed: Vec<u8> = (0..n * k_blocks * blob_size)
.map(|i| ((i * 29 + i / 7 + 13) & 0xff) as u8)
.collect();
let scales: Vec<f16> = (0..n * k_blocks)
.map(|i| f16::from_f32(0.01 + (i % 17) as f32 * 0.0005))
.collect();
let activation_dev = runtime.alloc_raw(activation.len() * 2).unwrap();
let packed_dev = runtime.alloc_raw(packed.len()).unwrap();
let scales_dev = runtime.alloc_raw(scales.len() * 2).unwrap();
let staged_output_dev = runtime.alloc_raw(n * 2).unwrap();
let down_output_dev = runtime.alloc_raw(n * 2).unwrap();
unsafe {
runtime.htod(as_bytes(&activation), activation_dev).unwrap();
runtime.htod(&packed, packed_dev).unwrap();
runtime.htod(as_bytes(&scales), scales_dev).unwrap();
}
let device = DeviceId::cuda(0);
let a_shape = [1usize, k];
let a_strides = [k as i64, 1];
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let y_shape = [1usize, n];
let y_strides = [n as i64, 1];
let activation_view = TensorView::new(
device_ptr(activation_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
);
let packed_view = TensorView::new(
device_ptr(packed_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
);
let scales_view = TensorView::new(
device_ptr(scales_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
);
let mut down_output = TensorMut::new(
device_ptr_mut(down_output_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
);
let kernel = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: false,
decomposed_silu: false,
rmsnorm_prologue: false,
rmsnorm_epsilon: 1e-5,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
let staged_source = format!("{GEMV_F16_SRC}\n{STAGED_DOWN_REFERENCE_SRC}");
let staged_function = runtime
.nvrtc_function(
"matmul_nbits_gemv_f16_down_staged_reference",
&staged_source,
STAGED_DOWN_REFERENCE_ENTRY,
)
.unwrap();
let activation_ptr = cuptr(activation_view.data_ptr::<u8>() as *const c_void);
let packed_ptr = cuptr(packed_view.data_ptr::<u8>() as *const c_void);
let scales_ptr = cuptr(scales_view.data_ptr::<u8>() as *const c_void);
let zero_points_ptr: CUdeviceptr = 0;
let bias_ptr: CUdeviceptr = 0;
let staged_output_ptr = staged_output_dev;
let k_i32 = as_i32("K", k).unwrap();
let n_i32 = as_i32("N", n).unwrap();
let block_size_i32 = as_i32("block_size", block_size).unwrap();
let k_blocks_i32 = as_i32("K block count", k_blocks).unwrap();
let blob_size_i32 = as_i32("block blob size", blob_size).unwrap();
let zp_row_bytes_i32 =
as_i32("zero-point row byte count", k_blocks.div_ceil(2)).unwrap();
let scales_fp16_flag = 1i32;
let bias_post_round_flag = 0i32;
let mut staged_builder = runtime.stream().launch_builder(&staged_function);
staged_builder
.arg(&activation_ptr)
.arg(&packed_ptr)
.arg(&scales_ptr)
.arg(&zero_points_ptr)
.arg(&bias_ptr)
.arg(&staged_output_ptr)
.arg(&k_i32)
.arg(&n_i32)
.arg(&block_size_i32)
.arg(&k_blocks_i32)
.arg(&blob_size_i32)
.arg(&zp_row_bytes_i32)
.arg(&scales_fp16_flag)
.arg(&bias_post_round_flag);
unsafe {
staged_builder
.launch(LaunchConfig {
grid_dim: (n.div_ceil(GEMV_F16_DOWN_COLUMNS_PER_BLOCK) as u32, 1, 1),
block_dim: (GEMV_F16_DOWN_THREADS, 1, 1),
shared_mem_bytes: (k * std::mem::size_of::<f16>()) as u32,
})
.unwrap();
}
kernel
.launch_f16_gemv_variant(
&activation_view,
&packed_view,
&scales_view,
true,
None,
None,
&mut down_output,
k_blocks,
blob_size,
k_blocks.div_ceil(2),
select_f16_gemv_variant(k, n, block_size, true, false),
)
.unwrap();
runtime.synchronize().unwrap();
let mut staged = vec![f16::ZERO; n];
let mut down = vec![f16::ZERO; n];
unsafe {
runtime
.dtoh(as_bytes_mut(&mut staged), staged_output_dev)
.unwrap();
runtime
.dtoh(as_bytes_mut(&mut down), down_output_dev)
.unwrap();
runtime.free_raw(activation_dev).unwrap();
runtime.free_raw(packed_dev).unwrap();
runtime.free_raw(scales_dev).unwrap();
runtime.free_raw(staged_output_dev).unwrap();
runtime.free_raw(down_output_dev).unwrap();
}
assert_eq!(
as_bytes(&staged),
as_bytes(&down),
"register-loaded down projection must be bit-exact to the pre-change staged kernel \
at K={k}, N={n}"
);
}
}
#[test]
fn fp16_gemv_variant_selection_is_structural() {
let qwen = select_f16_gemv_variant(QWEN_DOWN_K, QWEN_DOWN_N, 32, true, false);
assert_eq!(qwen.variant, F16GemvVariant::DownProjection);
assert_eq!(
qwen.reason,
"variant=down_projection;class=tall_skinny(K>N);block_size=32;\
scales=fp16;K%32==0"
);
for (k, n) in [
(5632, 2048),
(11008, 4096),
(2048, 512),
(4096, 4096 - 32),
(32_768, 4096),
] {
let selection = select_f16_gemv_variant(k, n, 32, true, false);
assert_eq!(
selection.variant,
F16GemvVariant::DownProjection,
"tall-skinny K={k}, N={n} must select the down variant"
);
}
let general_cases = [
(896, 4864, 32, true), (896, 896, 32, true), (896, 151_936, 32, true), (4880, 896, 32, true), (4864, 896, 64, true), ];
for (k, n, block_size, scales_fp16) in general_cases {
let selection = select_f16_gemv_variant(k, n, block_size, scales_fp16, false);
assert_eq!(
selection.variant,
F16GemvVariant::General,
"K={k}, N={n}, block_size={block_size} must retain the general GEMV"
);
}
assert_eq!(
select_f16_gemv_variant(QWEN_DOWN_K, QWEN_DOWN_N, 32, false, false).variant,
F16GemvVariant::General,
);
let asymmetric = select_f16_gemv_variant(QWEN_DOWN_K, QWEN_DOWN_N, 32, true, true);
assert_eq!(asymmetric.variant, F16GemvVariant::General);
assert_eq!(
asymmetric.reason,
"variant=general;zero_points=explicit;down_projection requires symmetric zp=8"
);
}
#[test]
fn scales_f16_pipe_well_occupied_routes_lm_head_to_plain() {
assert!(
scales_f16_pipe_well_occupied(248_320, 8, 132),
"the wide LM-head projection must be treated as well-occupied (plain)"
);
assert!(
!scales_f16_pipe_well_occupied(4096, 8, 132),
"a grid-starved N=4096 projection must keep the pipe entry"
);
assert!(
scales_f16_pipe_well_occupied(8192, 8, 32),
"grid 1024 >= 32*32=1024 on a 32-SM device is well-occupied"
);
assert!(
!scales_f16_pipe_well_occupied(8192, 8, 33),
"grid 1024 < 33*32=1056 on a 33-SM device stays grid-starved"
);
assert!(scales_f16_pipe_well_occupied(1_000_000, 0, 1));
assert!(!scales_f16_pipe_well_occupied(1, 8, 0));
}
#[test]
fn asymmetric_splitk_gate_splits_for_bandwidth_as_well_as_for_occupancy() {
let a100 = 108u32;
let starved_n = 8; let saturated_n = 19968;
assert!(
use_scales_f16_zp_splitk_gate(512, starved_n, a100),
"a grid-starved projection must still split to fill idle SMs, \
whatever K is"
);
assert!(
use_scales_f16_zp_splitk_gate(6656, saturated_n, a100),
"a saturated grid must still split when K is deep enough for the \
halved load chain to pay for the extra reduction"
);
assert!(
!use_scales_f16_zp_splitk_gate(512, saturated_n, a100),
"a saturated grid with a shallow K keeps the occupancy-only rule, \
which nothing has measured against"
);
assert!(
use_scales_f16_zp_splitk_gate(ZP_SPLITK_BANDWIDTH_MIN_K, saturated_n, a100)
&& !use_scales_f16_zp_splitk_gate(ZP_SPLITK_BANDWIDTH_MIN_K - 1, saturated_n, a100),
"the K threshold is inclusive at the measured depth and does not \
extrapolate below it"
);
}
#[test]
fn symmetric_int8_splitk_gate_targets_grid_starved_only() {
let (mp, max_threads) = (132u32, GEMV_F16_LARGE_THREADS);
assert!(
use_f16_symmetric_splitk(2048, 1024, mp, max_threads),
"grid-starved N=1024 symmetric int8 must take split-K"
);
assert!(
use_f16_symmetric_splitk(1024, 2048, mp, max_threads),
"grid-starved N=2048 symmetric int8 must take split-K"
);
assert!(
!use_f16_symmetric_splitk(1024, 3584, mp, max_threads),
"occupied N=3584 symmetric int8 must keep the single-warp entry"
);
assert!(
!use_f16_symmetric_splitk(1024, 6144, mp, max_threads),
"occupied N=6144 symmetric int8 must keep the single-warp entry"
);
assert!(!use_f16_symmetric_splitk(256, 512, mp, max_threads));
assert!(use_f16_symmetric_splitk(1024, 3584, 256, max_threads));
if std::env::var_os("ONNX_GENAI_CUDA_DISABLE_INT8_SYMMETRIC_SPLITK").is_none() {
assert!(int8_symmetric_splitk_enabled());
}
}
#[test]
fn splitk_smalln_single_column_is_shape_driven_and_general() {
assert!(
splitk_smalln_prefers_single_column(256, 132),
"a narrow N=256 kv projection must fall back to the single-column entry"
);
assert!(splitk_smalln_prefers_single_column(512, 132));
assert!(
!splitk_smalln_prefers_single_column(4096, 132),
"a wide N=4096 projection must keep the multicol hybrid"
);
assert!(!splitk_smalln_prefers_single_column(27_392, 132));
assert!(
splitk_smalln_prefers_single_column(1024, 132),
"N=1024 -> grid 128 < 132 SMs still under-fills a large device"
);
assert!(
!splitk_smalln_prefers_single_column(1024, 100),
"N=1024 -> grid 128 >= 100 SMs fills a smaller device"
);
unsafe {
std::env::set_var("ONNX_GENAI_GEMV_SPLITK_SMALLN_SINGLECOL", "1");
}
assert!(splitk_smalln_prefers_single_column(4096, 132));
unsafe {
std::env::set_var("ONNX_GENAI_GEMV_SPLITK_SMALLN_SINGLECOL", "0");
}
assert!(!splitk_smalln_prefers_single_column(256, 132));
unsafe {
std::env::remove_var("ONNX_GENAI_GEMV_SPLITK_SMALLN_SINGLECOL");
}
}
#[test]
fn down_columns_fill_the_device_and_never_undersplit() {
assert_eq!(
select_down_columns(65_536, 132),
(8, GEMV_F16_DOWN_ENTRY),
"a wide down projection must keep the 8-column launch"
);
assert_eq!(
select_down_columns(8192, 132),
(4, GEMV_F16_DOWN_C4_ENTRY),
"a mid-width down projection must split to 4 columns/CTA"
);
assert_eq!(
select_down_columns(3584, 132),
(2, GEMV_F16_DOWN_C2_ENTRY),
"a grid-starved down projection must split to 2 columns/CTA"
);
assert_eq!(select_down_columns(64, 132).0, 2);
assert_eq!(select_down_columns(3584, 0), (8, GEMV_F16_DOWN_ENTRY));
}
#[test]
fn accuracy4_stage64_is_limited_to_sub_wave_block32_grids() {
assert!(use_accuracy4_stage64(4608, 132, (9, 0), 64 * 1024));
assert!(!use_accuracy4_stage64(4608, 46, (8, 9), 48 * 1024));
assert!(!use_accuracy4_stage64(4608, 28, (8, 6), 48 * 1024));
assert!(!use_accuracy4_stage64(4608, 132, (9, 0), 1024));
assert!(!use_accuracy4_stage64(65_536, 132, (9, 0), 64 * 1024));
}
#[test]
fn accuracy4_resident_warps_matches_prior_inline_ladder() {
fn prior_inline_ladder(compute_capability: (u32, u32)) -> u32 {
match compute_capability {
(8, 0) | (9.., _) => 64,
_ => 48,
}
}
for major in 0u32..=13 {
for minor in 0u32..=12 {
let cc = (major, minor);
assert_eq!(
crate::arch::decode_resident_warps_per_sm(cc),
prior_inline_ladder(cc),
"arch resident-warp ladder diverged from the frozen inline ladder at {cc:?}"
);
}
}
for &(n, sms, cc, smem) in &[
(4608usize, 132u32, (9u32, 0u32), 64 * 1024u32),
(4608, 46, (8, 9), 48 * 1024),
(4608, 28, (8, 6), 48 * 1024),
(1024, 132, (9, 0), 64 * 1024),
(65_536, 132, (9, 0), 64 * 1024),
(2048, 84, (8, 9), 64 * 1024),
(4096, 68, (8, 6), 64 * 1024),
] {
let resident = prior_inline_ladder(cc) as usize;
let resident_ctas = resident / (GEMV_ACCURACY4_THREADS as usize / 32);
let one_wave = (sms.max(1) as usize).saturating_mul(resident_ctas);
let expected = smem >= GEMV_ACCURACY4_STAGE64_SHARED_BYTES
&& n.div_ceil(GEMV_ACCURACY4_COLUMNS_PER_BLOCK) < one_wave;
assert_eq!(
use_accuracy4_stage64(n, sms, cc, smem),
expected,
"stage64 decision changed for n={n} sms={sms} cc={cc:?} smem={smem}"
);
}
}
#[test]
fn symmetric_fp16_splitk_is_device_driven_and_falls_back_on_small_gpus() {
assert!(use_f16_symmetric_splitk(896, 1152, 132, 1024));
assert!(!use_f16_symmetric_splitk(896, 1152, 46, 1024));
assert!(!use_f16_symmetric_splitk(896, 1152, 132, 128));
assert!(!use_f16_symmetric_splitk(256, 1024, 132, 1024));
}
#[test]
fn fp16_down_projection_loads_activation_directly_into_registers() {
let start = GEMV_F16_SRC
.find("void matmul_nbits_gemv_f16_scales_f16_down_tpl")
.expect("down-projection template must exist");
let body = &GEMV_F16_SRC[start..];
let end = body
.find("\n}\n\n// Default 8-column down projection")
.expect("down-projection template must have a bounded body");
let body = &body[..end];
assert!(
!body.contains("activation_shared"),
"down projection must not round-trip activations through shared memory"
);
for offset in [
"activation_block);",
"activation_block + 8);",
"activation_block + 16);",
"activation_block + 24);",
] {
assert!(
body.contains(offset),
"down projection must directly load the block-32 activation at {offset}"
);
}
}
#[test]
fn fp16_gemv_matches_dequant_reference() {
let (mut worst_abs, mut worst_rel, mut max_out, mut all_finite) =
(0.0f32, 0.0f32, 0.0f32, true);
for (scales_fp16, with_bias) in [(false, false), (true, false), (true, true)] {
let (abs, rel, out, finite) = run_parity(scales_fp16, with_bias);
worst_abs = worst_abs.max(abs);
worst_rel = worst_rel.max(rel);
max_out = max_out.max(out);
all_finite &= finite;
}
let abs_bound = (max_out * 1e-3).max(1e-3);
eprintln!(
"MatMulNBits fp16 GEMV parity: max_abs={worst_abs:.3e} max_rel={worst_rel:.3e} \
max_out={max_out:.3e} abs_bound={abs_bound:.3e}"
);
assert!(all_finite, "fp16 GEMV produced a non-finite output");
assert!(
worst_abs < abs_bound,
"fp16 GEMV diverged from dequant reference: max_abs={worst_abs:.3e} bound={abs_bound:.3e}"
);
assert!(
worst_rel < 5e-2,
"fp16 GEMV diverged from dequant reference: max_rel={worst_rel:.3e}"
);
}
#[test]
fn decode_gemv_loop_is_byte_identical_to_per_row_singlestream() {
let Some(runtime) = runtime() else {
eprintln!("skipping looped decode GEMV byte-identity test: CUDA runtime unavailable");
return;
};
if runtime
.require_nvrtc_half_headers("matmul_nbits_gemv_f16")
.is_err()
{
eprintln!(
"skipping looped decode GEMV byte-identity test: fp16 NVRTC headers unavailable"
);
return;
}
for (k, n) in [
(896usize, 896usize),
(896usize, 4864usize),
(4864usize, 896usize),
] {
let m = 4usize;
let block_size = 32usize;
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let mut state = 0x0123_4567_89ab_cdefu64 ^ ((k as u64) << 20) ^ (n as u64);
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let mut activation_f16 = vec![f16::ZERO; m * k];
for h in activation_f16.iter_mut() {
*h = f16::from_f32(next());
}
let mut quant = vec![0u8; n * k];
for value in quant.iter_mut() {
*value = ((next() * 0.5 + 0.5) * 15.0).round().clamp(0.0, 15.0) as u8;
}
let mut packed = vec![0u8; n * k_blocks * blob_size];
for col in 0..n {
for block in 0..k_blocks {
for pair in 0..blob_size {
let low = quant[col * k + block * block_size + pair * 2] & 15;
let high = quant[col * k + block * block_size + pair * 2 + 1] & 15;
packed[(col * k_blocks + block) * blob_size + pair] = low | (high << 4);
}
}
}
let mut scale_f16 = vec![f16::ZERO; n * k_blocks];
for s in scale_f16.iter_mut() {
*s = f16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5));
}
let activation_dev = runtime.alloc_raw(activation_f16.len() * 2).unwrap();
let packed_dev = runtime.alloc_raw(packed.len()).unwrap();
let scales_dev = runtime.alloc_raw(scale_f16.len() * 2).unwrap();
let loop_out_dev = runtime.alloc_raw(m * n * 2).unwrap();
let ref_out_dev = runtime.alloc_raw(n * 2).unwrap();
unsafe {
runtime
.htod(as_bytes(&activation_f16), activation_dev)
.unwrap();
runtime.htod(&packed, packed_dev).unwrap();
runtime.htod(as_bytes(&scale_f16), scales_dev).unwrap();
}
let device = DeviceId::cuda(0);
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let packed_view = TensorView::new(
device_ptr(packed_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
);
let scales_view = TensorView::new(
device_ptr(scales_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
);
let kernel = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: false,
decomposed_silu: false,
rmsnorm_prologue: false,
rmsnorm_epsilon: 1e-5,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
let a_shape_m = [m, k];
let a_strides_m = [k as i64, 1];
let y_shape_m = [m, n];
let y_strides_m = [n as i64, 1];
let inputs_m = vec![
TensorView::new(
device_ptr(activation_dev),
DataType::Float16,
&a_shape_m,
&a_strides_m,
device,
),
packed_view,
scales_view,
];
{
let mut outputs_m = [TensorMut::new(
device_ptr_mut(loop_out_dev),
DataType::Float16,
&y_shape_m,
&y_strides_m,
device,
)];
kernel.run(&inputs_m, &mut outputs_m, None).unwrap();
kernel.run(&inputs_m, &mut outputs_m, None).unwrap();
runtime.synchronize().unwrap();
}
assert!(
kernel.last_call_capture_safe.load(Ordering::Relaxed),
"warm looped decode GEMV must report capture-safe (K={k} N={n})"
);
let mut loop_out = vec![f16::ZERO; m * n];
unsafe {
runtime
.dtoh(as_bytes_mut(&mut loop_out), loop_out_dev)
.unwrap();
}
let mut ref_out = vec![f16::ZERO; m * n];
for r in 0..m {
let a_off = (r * k * 2) as CUdeviceptr;
let a_shape_1 = [1usize, k];
let a_strides_1 = [k as i64, 1];
let y_shape_1 = [1usize, n];
let y_strides_1 = [n as i64, 1];
let inputs_1 = vec![
TensorView::new(
device_ptr(activation_dev + a_off),
DataType::Float16,
&a_shape_1,
&a_strides_1,
device,
),
packed_view,
scales_view,
];
let mut outputs_1 = [TensorMut::new(
device_ptr_mut(ref_out_dev),
DataType::Float16,
&y_shape_1,
&y_strides_1,
device,
)];
kernel.run(&inputs_1, &mut outputs_1, None).unwrap();
runtime.synchronize().unwrap();
unsafe {
runtime
.dtoh(as_bytes_mut(&mut ref_out[r * n..(r + 1) * n]), ref_out_dev)
.unwrap();
}
}
unsafe {
runtime.free_raw(activation_dev).unwrap();
runtime.free_raw(packed_dev).unwrap();
runtime.free_raw(scales_dev).unwrap();
runtime.free_raw(loop_out_dev).unwrap();
runtime.free_raw(ref_out_dev).unwrap();
}
let mismatches = loop_out
.iter()
.zip(ref_out.iter())
.filter(|(a, b)| a.to_bits() != b.to_bits())
.count();
assert_eq!(
mismatches,
0,
"looped batch-{m} decode GEMV diverged from per-row single-stream decode at \
K={k} N={n}: {mismatches}/{} fp16 outputs differ",
m * n
);
}
}
#[test]
fn gate_up_swiglu_rmsnorm_loop_is_byte_identical_to_per_row_singlestream() {
let Some(runtime) = runtime() else {
eprintln!(
"skipping rmsnorm gate/up SwiGLU byte-identity test: CUDA runtime unavailable"
);
return;
};
if runtime
.require_nvrtc_half_headers("matmul_nbits_gemv_f16")
.is_err()
{
eprintln!(
"skipping rmsnorm gate/up SwiGLU byte-identity test: fp16 NVRTC headers unavailable"
);
return;
}
for (k, n) in [(896usize, 4864usize), (896usize, 896usize)] {
let m = 4usize;
let block_size = 32usize;
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let mut state = 0x51ed_2701_c0ff_ee11u64 ^ ((k as u64) << 20) ^ (n as u64);
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let mut activation_f16 = vec![f16::ZERO; m * k];
for h in activation_f16.iter_mut() {
*h = f16::from_f32(next());
}
let mut gamma_f16 = vec![f16::ZERO; k];
for g in gamma_f16.iter_mut() {
*g = f16::from_f32(0.75 + 0.5 * (next() * 0.5 + 0.5));
}
let pack = |next: &mut dyn FnMut() -> f32| {
let mut quant = vec![0u8; n * k];
for value in quant.iter_mut() {
*value = ((next() * 0.5 + 0.5) * 15.0).round().clamp(0.0, 15.0) as u8;
}
let mut packed = vec![0u8; n * k_blocks * blob_size];
for col in 0..n {
for block in 0..k_blocks {
for pair in 0..blob_size {
let low = quant[col * k + block * block_size + pair * 2] & 15;
let high = quant[col * k + block * block_size + pair * 2 + 1] & 15;
packed[(col * k_blocks + block) * blob_size + pair] = low | (high << 4);
}
}
}
packed
};
let packed_gate = pack(&mut next);
let packed_up = pack(&mut next);
let mut scales_gate = vec![f16::ZERO; n * k_blocks];
for s in scales_gate.iter_mut() {
*s = f16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5));
}
let mut scales_up = vec![f16::ZERO; n * k_blocks];
for s in scales_up.iter_mut() {
*s = f16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5));
}
let activation_dev = runtime.alloc_raw(activation_f16.len() * 2).unwrap();
let gamma_dev = runtime.alloc_raw(gamma_f16.len() * 2).unwrap();
let packed_gate_dev = runtime.alloc_raw(packed_gate.len()).unwrap();
let packed_up_dev = runtime.alloc_raw(packed_up.len()).unwrap();
let scales_gate_dev = runtime.alloc_raw(scales_gate.len() * 2).unwrap();
let scales_up_dev = runtime.alloc_raw(scales_up.len() * 2).unwrap();
let loop_out_dev = runtime.alloc_raw(m * n * 2).unwrap();
let ref_out_dev = runtime.alloc_raw(n * 2).unwrap();
unsafe {
runtime
.htod(as_bytes(&activation_f16), activation_dev)
.unwrap();
runtime.htod(as_bytes(&gamma_f16), gamma_dev).unwrap();
runtime.htod(&packed_gate, packed_gate_dev).unwrap();
runtime.htod(&packed_up, packed_up_dev).unwrap();
runtime
.htod(as_bytes(&scales_gate), scales_gate_dev)
.unwrap();
runtime.htod(as_bytes(&scales_up), scales_up_dev).unwrap();
}
let device = DeviceId::cuda(0);
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let packed_gate_view = TensorView::new(
device_ptr(packed_gate_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
);
let packed_up_view = TensorView::new(
device_ptr(packed_up_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
);
let scales_gate_view = TensorView::new(
device_ptr(scales_gate_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
);
let scales_up_view = TensorView::new(
device_ptr(scales_up_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
);
let gamma_shape = [k];
let gamma_strides = [1i64];
let gamma_view = TensorView::new(
device_ptr(gamma_dev),
DataType::Float16,
&gamma_shape,
&gamma_strides,
device,
);
let kernel = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: true,
decomposed_silu: false,
rmsnorm_prologue: true,
rmsnorm_epsilon: 1e-5,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
let a_shape_m = [m, k];
let a_strides_m = [k as i64, 1];
let y_shape_m = [m, n];
let y_strides_m = [n as i64, 1];
let inputs_m = [
TensorView::new(
device_ptr(activation_dev),
DataType::Float16,
&a_shape_m,
&a_strides_m,
device,
),
packed_gate_view,
scales_gate_view,
packed_up_view,
scales_up_view,
gamma_view,
];
{
let mut outputs_m = [TensorMut::new(
device_ptr_mut(loop_out_dev),
DataType::Float16,
&y_shape_m,
&y_strides_m,
device,
)];
kernel
.run_f16_gate_up_swiglu(&inputs_m, &mut outputs_m, None)
.unwrap();
runtime.synchronize().unwrap();
}
assert!(
kernel.last_call_capture_safe.load(Ordering::Relaxed),
"looped rmsnorm gate/up SwiGLU decode must report capture-safe (K={k} N={n})"
);
let mut loop_out = vec![f16::ZERO; m * n];
unsafe {
runtime
.dtoh(as_bytes_mut(&mut loop_out), loop_out_dev)
.unwrap();
}
let mut ref_out = vec![f16::ZERO; m * n];
for r in 0..m {
let a_off = (r * k * 2) as CUdeviceptr;
let a_shape_1 = [1usize, k];
let a_strides_1 = [k as i64, 1];
let y_shape_1 = [1usize, n];
let y_strides_1 = [n as i64, 1];
let inputs_1 = [
TensorView::new(
device_ptr(activation_dev + a_off),
DataType::Float16,
&a_shape_1,
&a_strides_1,
device,
),
packed_gate_view,
scales_gate_view,
packed_up_view,
scales_up_view,
gamma_view,
];
let mut outputs_1 = [TensorMut::new(
device_ptr_mut(ref_out_dev),
DataType::Float16,
&y_shape_1,
&y_strides_1,
device,
)];
kernel
.run_f16_gate_up_swiglu(&inputs_1, &mut outputs_1, None)
.unwrap();
runtime.synchronize().unwrap();
unsafe {
runtime
.dtoh(as_bytes_mut(&mut ref_out[r * n..(r + 1) * n]), ref_out_dev)
.unwrap();
}
}
unsafe {
runtime.free_raw(activation_dev).unwrap();
runtime.free_raw(gamma_dev).unwrap();
runtime.free_raw(packed_gate_dev).unwrap();
runtime.free_raw(packed_up_dev).unwrap();
runtime.free_raw(scales_gate_dev).unwrap();
runtime.free_raw(scales_up_dev).unwrap();
runtime.free_raw(loop_out_dev).unwrap();
runtime.free_raw(ref_out_dev).unwrap();
}
let mismatches = loop_out
.iter()
.zip(ref_out.iter())
.filter(|(a, b)| a.to_bits() != b.to_bits())
.count();
assert_eq!(
mismatches,
0,
"looped batch-{m} rmsnorm gate/up SwiGLU decode diverged from per-row \
single-stream decode at K={k} N={n}: {mismatches}/{} fp16 outputs differ",
m * n
);
}
}
#[test]
fn fp16_gemv_matches_dequant_reference_qwen_1_5b_dims() {
for (k, n) in [(1536usize, 8960usize), (8960usize, 1536usize)] {
let (mut worst_abs, mut worst_rel, mut max_out, mut all_finite) =
(0.0f32, 0.0f32, 0.0f32, true);
for (scales_fp16, with_bias) in [(false, false), (true, false), (true, true)] {
let (abs, rel, out, finite) = run_parity_dims(k, n, scales_fp16, with_bias, false);
worst_abs = worst_abs.max(abs);
worst_rel = worst_rel.max(rel);
max_out = max_out.max(out);
all_finite &= finite;
}
let abs_bound = (max_out * 1e-3).max(1e-3);
eprintln!(
"MatMulNBits fp16 GEMV parity K={k} N={n}: max_abs={worst_abs:.3e} \
max_rel={worst_rel:.3e} max_out={max_out:.3e} abs_bound={abs_bound:.3e}"
);
assert!(
all_finite,
"fp16 GEMV produced a non-finite output (K={k} N={n})"
);
assert!(
worst_abs < abs_bound,
"fp16 GEMV diverged from dequant reference at K={k} N={n}: \
max_abs={worst_abs:.3e} bound={abs_bound:.3e}"
);
assert!(
worst_rel < 5e-2,
"fp16 GEMV diverged from dequant reference at K={k} N={n}: max_rel={worst_rel:.3e}"
);
}
}
#[test]
fn fp16_gemv_matches_dequant_reference_block128() {
for (block_size, k, n) in [
(128usize, 896usize, 896usize),
(128usize, 896usize, 4864usize),
(128usize, 4864usize, 896usize),
(128usize, 896usize, 70usize),
(64usize, 896usize, 896usize),
(64usize, 896usize, 70usize),
] {
assert_eq!(
select_f16_gemv_variant(k, n, block_size, true, false).variant,
F16GemvVariant::General,
"block_size={block_size} K={k} N={n} must select the general variant"
);
let (mut worst_abs, mut worst_rel, mut max_out, mut all_finite) =
(0.0f32, 0.0f32, 0.0f32, true);
for (scales_fp16, with_bias, explicit_zp) in [
(false, false, false),
(true, false, false),
(true, true, false),
(false, false, true),
(true, true, true),
] {
let (abs, rel, out, finite) =
run_parity_dims_block(k, n, block_size, scales_fp16, with_bias, explicit_zp);
worst_abs = worst_abs.max(abs);
worst_rel = worst_rel.max(rel);
max_out = max_out.max(out);
all_finite &= finite;
}
let abs_bound = (max_out * 1e-3).max(1e-3);
eprintln!(
"MatMulNBits fp16 GEMV block-{block_size} parity K={k} N={n}: \
max_abs={worst_abs:.3e} max_rel={worst_rel:.3e} max_out={max_out:.3e} \
abs_bound={abs_bound:.3e}"
);
assert!(
all_finite,
"block-{block_size} fp16 GEMV produced a non-finite output (K={k} N={n})"
);
assert!(
worst_abs < abs_bound,
"block-{block_size} fp16 GEMV diverged from dequant reference at K={k} N={n}: \
max_abs={worst_abs:.3e} bound={abs_bound:.3e}"
);
assert!(
worst_rel < 5e-2,
"block-{block_size} fp16 GEMV diverged from dequant reference at K={k} N={n}: \
max_rel={worst_rel:.3e}"
);
}
}
#[test]
fn fp16_gemv_matches_dequant_reference_phi_int4_zp_dims() {
for (k, n) in [(3072usize, 3072usize), (3072, 8192)] {
let (mut worst_abs, mut worst_rel, mut max_out, mut all_finite) =
(0.0f32, 0.0f32, 0.0f32, true);
for (with_bias, explicit_zp) in [(false, true), (true, true)] {
let (abs, rel, out, finite) = run_parity_dims(k, n, true, with_bias, explicit_zp);
worst_abs = worst_abs.max(abs);
worst_rel = worst_rel.max(rel);
max_out = max_out.max(out);
all_finite &= finite;
}
let abs_bound = (max_out * 1e-3).max(1e-3);
eprintln!(
"MatMulNBits int4 asymmetric-zp GEMV parity K={k} N={n}: max_abs={worst_abs:.3e} \
max_rel={worst_rel:.3e} max_out={max_out:.3e} abs_bound={abs_bound:.3e}"
);
assert!(
all_finite,
"int4 asymmetric-zp GEMV produced a non-finite output (K={k} N={n})"
);
assert!(
worst_abs < abs_bound,
"int4 asymmetric-zp GEMV diverged from dequant reference at K={k} N={n}: \
max_abs={worst_abs:.3e} bound={abs_bound:.3e}"
);
assert!(
worst_rel < 5e-2,
"int4 asymmetric-zp GEMV diverged from dequant reference at K={k} N={n}: \
max_rel={worst_rel:.3e}"
);
}
}
#[test]
fn fp16_gemv_matches_dequant_reference_gqa_kv_deep_split_dims() {
let (k, n) = (6656usize, 256usize);
let (mut worst_abs, mut worst_rel, mut max_out, mut all_finite) =
(0.0f32, 0.0f32, 0.0f32, true);
for (with_bias, explicit_zp) in [(false, true), (true, true)] {
let (abs, rel, out, finite) = run_parity_dims(k, n, true, with_bias, explicit_zp);
worst_abs = worst_abs.max(abs);
worst_rel = worst_rel.max(rel);
max_out = max_out.max(out);
all_finite &= finite;
}
let abs_bound = (max_out * 1e-3).max(1e-3);
eprintln!(
"MatMulNBits GQA k/v deep-split GEMV parity K={k} N={n}: max_abs={worst_abs:.3e} \
max_rel={worst_rel:.3e} max_out={max_out:.3e} abs_bound={abs_bound:.3e}"
);
assert!(
all_finite,
"GQA k/v deep-split GEMV produced a non-finite output"
);
assert!(
worst_abs < abs_bound,
"GQA k/v deep-split GEMV diverged from dequant reference: \
max_abs={worst_abs:.3e} bound={abs_bound:.3e}"
);
assert!(
worst_rel < 5e-2,
"GQA k/v deep-split GEMV diverged from dequant reference: max_rel={worst_rel:.3e}"
);
}
#[test]
fn zp_split_factor_deepens_only_for_starved_grids_and_keeps_a_column_per_block() {
let warps = (GEMV_F16_LARGE_THREADS / 32) as usize;
assert_eq!(
scales_f16_zp_split_factor(256, 108),
GEMV_F16_SCALES_F16_ZP_SPLITK8
);
for n in [4096usize, 6656, 19968, 202048] {
assert_eq!(
scales_f16_zp_split_factor(n, 108),
GEMV_F16_SCALES_F16_ZP_SPLITK,
"N={n} fills the grid at the default split and must not deepen"
);
}
for mp in [1u32, 16, 80, 108, 132, 148, 256] {
for n in [1usize, 2, 8, 64, 256, 257, 431, 432, 433, 4096, 202048] {
let factor = scales_f16_zp_split_factor(n, mp);
assert!(
factor <= warps && warps.is_multiple_of(factor),
"N={n} mp={mp} produced factor {factor}, which does not divide {warps} warps"
);
}
}
}
#[test]
fn int8_fp16_gemv_matches_dequant_reference_phi_dims() {
let (mut worst_abs, mut worst_rel, mut max_out, mut all_finite) =
(0.0f32, 0.0f32, 0.0f32, true);
let cases = [
(3072usize, 5120usize, true, false, false), (8192, 3072, true, false, false), (3072, 5120, true, true, false), (3072, 5120, true, false, true), (8192, 3072, true, false, true), (3072, 5121, true, false, false), (8192, 3072, false, false, false), ];
for (k, n, scales_fp16, with_bias, explicit_zp) in cases {
let (abs, rel, out, finite) =
run_int8_parity_dims(k, n, 32, scales_fp16, with_bias, explicit_zp);
worst_abs = worst_abs.max(abs);
worst_rel = worst_rel.max(rel);
max_out = max_out.max(out);
all_finite &= finite;
}
let abs_bound = (max_out * 2e-3).max(1e-3);
eprintln!(
"MatMulNBits int8 fp16 GEMV parity: max_abs={worst_abs:.3e} max_rel={worst_rel:.3e} \
max_out={max_out:.3e} abs_bound={abs_bound:.3e}"
);
assert!(all_finite, "int8 fp16 GEMV produced a non-finite output");
assert!(
worst_abs < abs_bound,
"int8 fp16 GEMV diverged from dequant reference: max_abs={worst_abs:.3e} \
bound={abs_bound:.3e}"
);
assert!(
worst_rel < 5e-2,
"int8 fp16 GEMV diverged from dequant reference: max_rel={worst_rel:.3e}"
);
}
#[test]
fn int8_fp16_gemv_matches_dequant_reference_block128() {
for (block_size, k, n) in [
(128usize, 1024usize, 1024usize), (128, 3072, 1024), (128, 1024, 3072), (128, 1024, 1030), (64, 1024, 1024), ] {
assert_eq!(
select_f16_gemv_variant(k, n, block_size, true, false).variant,
F16GemvVariant::General,
"int8 block_size={block_size} K={k} N={n} must select the general variant"
);
let (mut worst_abs, mut worst_rel, mut max_out, mut all_finite) =
(0.0f32, 0.0f32, 0.0f32, true);
for (scales_fp16, with_bias, explicit_zp) in [
(true, false, false),
(false, false, false),
(true, true, false),
(true, false, true),
] {
let (abs, rel, out, finite) =
run_int8_parity_dims(k, n, block_size, scales_fp16, with_bias, explicit_zp);
worst_abs = worst_abs.max(abs);
worst_rel = worst_rel.max(rel);
max_out = max_out.max(out);
all_finite &= finite;
}
let abs_bound = (max_out * 2e-3).max(1e-3);
eprintln!(
"MatMulNBits int8 fp16 GEMV block-{block_size} parity K={k} N={n}: \
max_abs={worst_abs:.3e} max_rel={worst_rel:.3e} max_out={max_out:.3e} \
abs_bound={abs_bound:.3e}"
);
assert!(
all_finite,
"int8 block-{block_size} fp16 GEMV produced a non-finite output (K={k} N={n})"
);
assert!(
worst_abs < abs_bound,
"int8 block-{block_size} fp16 GEMV diverged from dequant reference at K={k} \
N={n}: max_abs={worst_abs:.3e} bound={abs_bound:.3e}"
);
assert!(
worst_rel < 5e-2,
"int8 block-{block_size} fp16 GEMV diverged from dequant reference at K={k} \
N={n}: max_rel={worst_rel:.3e}"
);
}
}
fn run_int8_f32_parity_dims(
k: usize,
n: usize,
block_size: usize,
with_bias: bool,
explicit_zp: bool,
) -> (f32, f32, f32, bool) {
let Some(runtime) = runtime() else {
eprintln!("skipping MatMulNBits int8 f32 GEMV parity test: CUDA runtime unavailable");
return (0.0, 0.0, 0.0, true);
};
let k_blocks = k / block_size;
let blob_size = block_size;
let mut state = 0x1234_5678_9abc_def0u64;
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let mut activation = vec![0.0f32; k];
for value in activation.iter_mut() {
*value = next();
}
let mut quant = vec![0u8; n * k];
for value in quant.iter_mut() {
*value = ((next() * 0.5 + 0.5) * 255.0).round().clamp(0.0, 255.0) as u8;
}
let mut packed = vec![0u8; n * k_blocks * blob_size];
for col in 0..n {
for block in 0..k_blocks {
for within in 0..block_size {
packed[(col * k_blocks + block) * blob_size + within] =
quant[col * k + block * block_size + within];
}
}
}
let mut zero_points = vec![0u8; n * k_blocks];
if explicit_zp {
for zp in zero_points.iter_mut() {
*zp = ((next() * 0.5 + 0.5) * 255.0).round().clamp(0.0, 255.0) as u8;
}
}
let zp_ref = |col: usize, block: usize| -> i32 {
if explicit_zp {
zero_points[col * k_blocks + block] as i32
} else {
128
}
};
let mut scale = vec![0.0f32; n * k_blocks];
for value in scale.iter_mut() {
*value = 0.015 + 0.01 * (next() * 0.5 + 0.5);
}
let mut bias = vec![0.0f32; n];
if with_bias {
for value in bias.iter_mut() {
*value = next();
}
}
let mut expected = vec![0.0f32; n];
for col in 0..n {
let mut acc = 0.0f64;
for block in 0..k_blocks {
let s = scale[col * k_blocks + block] as f64;
let zp = zp_ref(col, block);
for within in 0..block_size {
let depth = block * block_size + within;
let q = quant[col * k + depth] as i32 - zp;
acc += activation[depth] as f64 * q as f64 * s;
}
}
if with_bias {
acc += bias[col] as f64;
}
expected[col] = acc as f32;
}
let activation_dev = runtime.alloc_raw(activation.len() * 4).unwrap();
let packed_dev = runtime.alloc_raw(packed.len()).unwrap();
let scales_dev = runtime.alloc_raw(scale.len() * 4).unwrap();
let zp_dev = runtime.alloc_raw(zero_points.len().max(1)).unwrap();
let bias_dev = runtime.alloc_raw(n * 4).unwrap();
let output_dev = runtime.alloc_raw(n * 4).unwrap();
unsafe {
runtime.htod(as_bytes(&activation), activation_dev).unwrap();
runtime.htod(&packed, packed_dev).unwrap();
runtime.htod(as_bytes(&scale), scales_dev).unwrap();
if explicit_zp {
runtime.htod(&zero_points, zp_dev).unwrap();
}
if with_bias {
runtime.htod(as_bytes(&bias), bias_dev).unwrap();
}
}
let a_shape = [1usize, k];
let a_strides = [k as i64, 1];
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let zp_shape = [n, k_blocks];
let zp_strides = [k_blocks as i64, 1];
let bias_shape = [n];
let bias_strides = [1i64];
let y_shape = [1usize, n];
let y_strides = [n as i64, 1];
let device = DeviceId::cuda(0);
let mut inputs = vec![
TensorView::new(
device_ptr(activation_dev),
DataType::Float32,
&a_shape,
&a_strides,
device,
),
TensorView::new(
device_ptr(packed_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
),
TensorView::new(
device_ptr(scales_dev),
DataType::Float32,
&scales_shape,
&scales_strides,
device,
),
];
if explicit_zp || with_bias {
inputs.push(if explicit_zp {
TensorView::new(
device_ptr(zp_dev),
DataType::Uint8,
&zp_shape,
&zp_strides,
device,
)
} else {
TensorView::absent(DataType::Uint8)
});
}
if with_bias {
inputs.push(TensorView::absent(DataType::Int32));
inputs.push(TensorView::new(
device_ptr(bias_dev),
DataType::Float32,
&bias_shape,
&bias_strides,
device,
));
}
let mut outputs = [TensorMut::new(
device_ptr_mut(output_dev),
DataType::Float32,
&y_shape,
&y_strides,
device,
)];
let kernel = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 8,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: false,
decomposed_silu: false,
rmsnorm_prologue: false,
rmsnorm_epsilon: 1e-5,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
kernel.run(&inputs, &mut outputs, None).unwrap();
runtime.synchronize().unwrap();
assert!(
kernel.last_call_capture_safe.load(Ordering::Relaxed),
"int8 f32 decode GEMV must report capture-safe"
);
let mut got = vec![0.0f32; n];
unsafe {
runtime.dtoh(as_bytes_mut(&mut got), output_dev).unwrap();
}
unsafe {
runtime.free_raw(activation_dev).unwrap();
runtime.free_raw(packed_dev).unwrap();
runtime.free_raw(scales_dev).unwrap();
runtime.free_raw(zp_dev).unwrap();
runtime.free_raw(bias_dev).unwrap();
runtime.free_raw(output_dev).unwrap();
}
let mut worst_abs = 0.0f32;
let mut worst_rel = 0.0f32;
let mut max_out = 0.0f32;
let mut all_finite = true;
for (g, e) in got.iter().zip(expected.iter()) {
if !g.is_finite() {
all_finite = false;
}
let abs = (g - e).abs();
let rel = abs / e.abs().max(1e-1);
worst_abs = worst_abs.max(abs);
worst_rel = worst_rel.max(rel);
max_out = max_out.max(e.abs());
}
(worst_abs, worst_rel, max_out, all_finite)
}
#[test]
fn int8_f32_gemv_matches_dequant_reference_block128() {
for (block_size, k, n) in [
(128usize, 1024usize, 1024usize), (128, 3072, 1024), (128, 1024, 3072), (128, 1024, 1030), (64, 1024, 1024), ] {
let (mut worst_abs, mut worst_rel, mut max_out, mut all_finite) =
(0.0f32, 0.0f32, 0.0f32, true);
for (with_bias, explicit_zp) in [(false, false), (true, false), (false, true)] {
let (abs, rel, out, finite) =
run_int8_f32_parity_dims(k, n, block_size, with_bias, explicit_zp);
worst_abs = worst_abs.max(abs);
worst_rel = worst_rel.max(rel);
max_out = max_out.max(out);
all_finite &= finite;
}
let abs_bound = (max_out * 2e-3).max(1e-3);
eprintln!(
"MatMulNBits int8 f32 GEMV block-{block_size} parity K={k} N={n}: \
max_abs={worst_abs:.3e} max_rel={worst_rel:.3e} max_out={max_out:.3e} \
abs_bound={abs_bound:.3e}"
);
assert!(
all_finite,
"int8 block-{block_size} f32 GEMV produced a non-finite output (K={k} N={n})"
);
assert!(
worst_abs < abs_bound,
"int8 block-{block_size} f32 GEMV diverged from dequant reference at K={k} \
N={n}: max_abs={worst_abs:.3e} bound={abs_bound:.3e}"
);
assert!(
worst_rel < 5e-2,
"int8 block-{block_size} f32 GEMV diverged from dequant reference at K={k} \
N={n}: max_rel={worst_rel:.3e}"
);
}
}
fn run_f32_block128_byte_identity(
bits: usize,
k: usize,
n: usize,
with_bias: bool,
explicit_zp: bool,
) -> (usize, bool) {
let Some(runtime) = runtime() else {
eprintln!(
"skipping MatMulNBits int{bits} f32 block-128 byte-identity test: CUDA runtime unavailable"
);
return (0, true);
};
let block_size = 128usize;
let k_blocks = k.div_ceil(block_size);
let blob_size = block_size * bits / 8;
let zp_row_bytes = (k_blocks * bits).div_ceil(8);
let mut state = 0x0bad_c0de_1234_5678u64;
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let mut activation = vec![0.0f32; k];
for value in activation.iter_mut() {
*value = next();
}
let mut packed = vec![0u8; n * k_blocks * blob_size];
for value in packed.iter_mut() {
*value = ((next() * 0.5 + 0.5) * 255.0).round().clamp(0.0, 255.0) as u8;
}
let mut zero_points = vec![0u8; n * zp_row_bytes];
if explicit_zp {
for zp in zero_points.iter_mut() {
*zp = ((next() * 0.5 + 0.5) * 255.0).round().clamp(0.0, 255.0) as u8;
}
}
let mut scale = vec![0.0f32; n * k_blocks];
for value in scale.iter_mut() {
*value = 0.015 + 0.01 * (next() * 0.5 + 0.5);
}
let mut bias = vec![0.0f32; n];
if with_bias {
for value in bias.iter_mut() {
*value = next();
}
}
let activation_dev = runtime.alloc_raw(activation.len() * 4).unwrap();
let packed_dev = runtime.alloc_raw(packed.len()).unwrap();
let scales_dev = runtime.alloc_raw(scale.len() * 4).unwrap();
let zp_dev = runtime.alloc_raw(zero_points.len().max(1)).unwrap();
let bias_dev = runtime.alloc_raw(n * 4).unwrap();
let spec_dev = runtime.alloc_raw(n * 4).unwrap();
let generic_dev = runtime.alloc_raw(n * 4).unwrap();
unsafe {
runtime.htod(as_bytes(&activation), activation_dev).unwrap();
runtime.htod(&packed, packed_dev).unwrap();
runtime.htod(as_bytes(&scale), scales_dev).unwrap();
if explicit_zp {
runtime.htod(&zero_points, zp_dev).unwrap();
}
if with_bias {
runtime.htod(as_bytes(&bias), bias_dev).unwrap();
}
}
let device = DeviceId::cuda(0);
let a_shape = [1usize, k];
let a_strides = [k as i64, 1];
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let zp_shape = [n, zp_row_bytes];
let zp_strides = [zp_row_bytes as i64, 1];
let bias_shape = [n];
let bias_strides = [1i64];
let activation_view = TensorView::new(
device_ptr(activation_dev),
DataType::Float32,
&a_shape,
&a_strides,
device,
);
let packed_view = TensorView::new(
device_ptr(packed_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
);
let scales_view = TensorView::new(
device_ptr(scales_dev),
DataType::Float32,
&scales_shape,
&scales_strides,
device,
);
let zp_view = explicit_zp.then(|| {
TensorView::new(
device_ptr(zp_dev),
DataType::Uint8,
&zp_shape,
&zp_strides,
device,
)
});
let bias_view = with_bias.then(|| {
TensorView::new(
device_ptr(bias_dev),
DataType::Float32,
&bias_shape,
&bias_strides,
device,
)
});
let y_shape = [1usize, n];
let y_strides = [n as i64, 1];
let mut spec_out = TensorMut::new(
device_ptr_mut(spec_dev),
DataType::Float32,
&y_shape,
&y_strides,
device,
);
let mut generic_out = TensorMut::new(
device_ptr_mut(generic_dev),
DataType::Float32,
&y_shape,
&y_strides,
device,
);
let kernel = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: false,
decomposed_silu: false,
rmsnorm_prologue: false,
rmsnorm_epsilon: 1e-5,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
match bits {
4 => kernel
.launch_int4_f32_gemv_block128(
&activation_view,
&packed_view,
&scales_view,
zp_view
.as_ref()
.expect("int4 block-128 specialization requires zero points"),
bias_view.as_ref(),
&mut spec_out,
k_blocks,
)
.unwrap(),
8 => kernel
.launch_int8_f32_gemv_block128(
&activation_view,
&packed_view,
&scales_view,
zp_view.as_ref(),
bias_view.as_ref(),
&mut spec_out,
k_blocks,
)
.unwrap(),
_ => unreachable!("byte-identity harness only supports int4/int8"),
}
kernel
.launch_f32_gemv(
&activation_view,
&packed_view,
&scales_view,
zp_view.as_ref(),
bias_view.as_ref(),
&mut generic_out,
k_blocks,
blob_size,
zp_row_bytes,
)
.unwrap();
runtime.synchronize().unwrap();
let mut spec = vec![0.0f32; n];
let mut generic = vec![0.0f32; n];
unsafe {
runtime.dtoh(as_bytes_mut(&mut spec), spec_dev).unwrap();
runtime
.dtoh(as_bytes_mut(&mut generic), generic_dev)
.unwrap();
}
unsafe {
runtime.free_raw(activation_dev).unwrap();
runtime.free_raw(packed_dev).unwrap();
runtime.free_raw(scales_dev).unwrap();
runtime.free_raw(zp_dev).unwrap();
runtime.free_raw(bias_dev).unwrap();
runtime.free_raw(spec_dev).unwrap();
runtime.free_raw(generic_dev).unwrap();
}
let mut mismatches = 0usize;
let mut all_finite = true;
for (s, g) in spec.iter().zip(generic.iter()) {
if !s.is_finite() {
all_finite = false;
}
if s.to_bits() != g.to_bits() {
mismatches += 1;
}
}
(mismatches, all_finite)
}
#[test]
fn int8_f32_block128_specialization_is_bit_identical_to_generic() {
for (k, n) in [
(1024usize, 1024usize), (3072, 1024), (1024, 3072), (1024, 1030), (1000, 1024), (896, 1027), ] {
for (with_bias, explicit_zp) in [(true, true), (false, true), (true, false)] {
let (mismatches, all_finite) =
run_f32_block128_byte_identity(8, k, n, with_bias, explicit_zp);
assert!(
all_finite,
"int8 block-128 specialization produced a non-finite output \
(K={k} N={n} bias={with_bias} zp={explicit_zp})"
);
assert_eq!(
mismatches, 0,
"int8 block-128 specialization diverged from the generic GEMV in \
{mismatches} column(s) (K={k} N={n} bias={with_bias} zp={explicit_zp}) — \
output must be bit-for-bit identical"
);
}
}
}
#[test]
fn int4_f32_block128_specialization_is_bit_identical_to_generic() {
let (mismatches, all_finite) = run_f32_block128_byte_identity(4, 259, 37, true, true);
assert!(
all_finite,
"int4 block-128 specialization produced non-finite output"
);
assert_eq!(
mismatches, 0,
"int4 block-128 specialization diverged from generic GEMV in \
{mismatches} column(s)"
);
}
#[test]
fn fp16_folded_bias_is_bit_exact_to_two_op_path() {
let Some(runtime) = runtime() else {
eprintln!("skipping folded-bias bit-exactness test: CUDA runtime unavailable");
return;
};
if runtime
.require_nvrtc_half_headers("matmul_nbits_gemv_f16")
.is_err()
{
eprintln!("skipping folded-bias bit-exactness test: fp16 NVRTC headers unavailable");
return;
}
let k = 896usize;
let n = 1152usize;
let block_size = 32usize;
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let mut state = 0x1234_5678_9abc_def0u64;
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let activation: Vec<f16> = (0..k).map(|_| f16::from_f32(next())).collect();
let mut quant = vec![0u8; n * k];
for value in quant.iter_mut() {
*value = ((next() * 0.5 + 0.5) * 15.0).round().clamp(0.0, 15.0) as u8;
}
let mut packed = vec![0u8; n * k_blocks * blob_size];
for col in 0..n {
for block in 0..k_blocks {
for pair in 0..blob_size {
let low = quant[col * k + block * block_size + pair * 2] & 15;
let high = quant[col * k + block * block_size + pair * 2 + 1] & 15;
packed[(col * k_blocks + block) * blob_size + pair] = low | (high << 4);
}
}
}
let scales: Vec<f16> = (0..n * k_blocks)
.map(|_| f16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5)))
.collect();
let bias: Vec<f16> = (0..n).map(|_| f16::from_f32(next() * 4.0)).collect();
let activation_dev = runtime.alloc_raw(activation.len() * 2).unwrap();
let packed_dev = runtime.alloc_raw(packed.len()).unwrap();
let scales_dev = runtime.alloc_raw(scales.len() * 2).unwrap();
let bias_dev = runtime.alloc_raw(bias.len() * 2).unwrap();
let nobias_output_dev = runtime.alloc_raw(n * 2).unwrap();
let fused_output_dev = runtime.alloc_raw(n * 2).unwrap();
unsafe {
runtime.htod(as_bytes(&activation), activation_dev).unwrap();
runtime.htod(&packed, packed_dev).unwrap();
runtime.htod(as_bytes(&scales), scales_dev).unwrap();
runtime.htod(as_bytes(&bias), bias_dev).unwrap();
}
let device = DeviceId::cuda(0);
let a_shape = [1usize, k];
let a_strides = [k as i64, 1];
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let bias_shape = [n];
let bias_strides = [1i64];
let y_shape = [1usize, n];
let y_strides = [n as i64, 1];
let activation_view = TensorView::new(
device_ptr(activation_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
);
let packed_view = TensorView::new(
device_ptr(packed_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
);
let scales_view = TensorView::new(
device_ptr(scales_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
);
let bias_view = TensorView::new(
device_ptr(bias_dev),
DataType::Float16,
&bias_shape,
&bias_strides,
device,
);
let mut nobias_output = TensorMut::new(
device_ptr_mut(nobias_output_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
);
let mut fused_output = TensorMut::new(
device_ptr_mut(fused_output_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
);
let selection = select_f16_gemv_variant(k, n, block_size, true, false);
let kernel_nobias = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: false,
decomposed_silu: false,
rmsnorm_prologue: false,
rmsnorm_epsilon: 1e-5,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
let kernel_fold = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: true,
gate_up_swiglu: false,
decomposed_silu: false,
rmsnorm_prologue: false,
rmsnorm_epsilon: 1e-5,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
kernel_nobias
.launch_f16_gemv_variant(
&activation_view,
&packed_view,
&scales_view,
true,
None,
None,
&mut nobias_output,
k_blocks,
blob_size,
k_blocks.div_ceil(2),
selection,
)
.unwrap();
kernel_fold
.launch_f16_gemv_variant(
&activation_view,
&packed_view,
&scales_view,
true,
None,
Some(&bias_view),
&mut fused_output,
k_blocks,
blob_size,
k_blocks.div_ceil(2),
selection,
)
.unwrap();
runtime.synchronize().unwrap();
let mut gemv_out = vec![f16::ZERO; n];
let mut fused_out = vec![f16::ZERO; n];
unsafe {
runtime
.dtoh(as_bytes_mut(&mut gemv_out), nobias_output_dev)
.unwrap();
runtime
.dtoh(as_bytes_mut(&mut fused_out), fused_output_dev)
.unwrap();
runtime.free_raw(activation_dev).unwrap();
runtime.free_raw(packed_dev).unwrap();
runtime.free_raw(scales_dev).unwrap();
runtime.free_raw(bias_dev).unwrap();
runtime.free_raw(nobias_output_dev).unwrap();
runtime.free_raw(fused_output_dev).unwrap();
}
for col in 0..n {
let two_op = f16::from_f32(gemv_out[col].to_f32() + bias[col].to_f32());
assert_eq!(
fused_out[col].to_bits(),
two_op.to_bits(),
"folded bias diverged at column {col}: fused={:?} two_op={:?} (gemv={:?} bias={:?})",
fused_out[col],
two_op,
gemv_out[col],
bias[col]
);
}
}
const REF_SILU_MUL_SRC: &str = r#"
#include <cuda_fp16.h>
__device__ float ref_op_silu(float x) {
if (x >= 0.0f) {
const float denominator = __fadd_rn(1.0f, (float)exp((double)-x));
return __fdiv_rn(x, denominator);
}
const float e = (float)exp((double)x);
const float numerator = __fmul_rn(x, e);
return __fdiv_rn(numerator, __fadd_rn(1.0f, e));
}
extern "C" __global__ void ref_silu_mul_f16(
const __half* g, const __half* u, __half* y, const int n) {
const int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
y[i] = __float2half_rn(
__fmul_rn(ref_op_silu(__half2float(g[i])), __half2float(u[i])));
}
}
"#;
#[test]
fn fp16_gate_up_swiglu_is_bit_exact_to_two_op_path() {
let Some(runtime) = runtime() else {
eprintln!("skipping gate/up SwiGLU bit-exactness test: CUDA runtime unavailable");
return;
};
if runtime
.require_nvrtc_half_headers("matmul_nbits_gemv_f16")
.is_err()
{
eprintln!("skipping gate/up SwiGLU bit-exactness test: fp16 NVRTC headers unavailable");
return;
}
for (m, k, n) in [
(1usize, QWEN_DOWN_N, QWEN_DOWN_K),
(1, 2048, 5632),
(5, QWEN_DOWN_N, QWEN_DOWN_K),
(3, 96, 77),
] {
let block_size = 32usize;
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let mut state = 0x0bad_c0de_dead_beefu64;
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let pack = |next: &mut dyn FnMut() -> f32| -> Vec<u8> {
let mut quant = vec![0u8; n * k];
for value in quant.iter_mut() {
*value = ((next() * 0.5 + 0.5) * 15.0).round().clamp(0.0, 15.0) as u8;
}
let mut packed = vec![0u8; n * k_blocks * blob_size];
for col in 0..n {
for block in 0..k_blocks {
for pair in 0..blob_size {
let low = quant[col * k + block * block_size + pair * 2] & 15;
let high = quant[col * k + block * block_size + pair * 2 + 1] & 15;
packed[(col * k_blocks + block) * blob_size + pair] = low | (high << 4);
}
}
}
packed
};
let activation: Vec<f16> = (0..m * k).map(|_| f16::from_f32(next())).collect();
let packed_gate = pack(&mut next);
let scales_gate: Vec<f16> = (0..n * k_blocks)
.map(|_| f16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5)))
.collect();
let packed_up = pack(&mut next);
let scales_up: Vec<f16> = (0..n * k_blocks)
.map(|_| f16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5)))
.collect();
let activation_dev = runtime.alloc_raw(activation.len() * 2).unwrap();
let packed_gate_dev = runtime.alloc_raw(packed_gate.len()).unwrap();
let scales_gate_dev = runtime.alloc_raw(scales_gate.len() * 2).unwrap();
let packed_up_dev = runtime.alloc_raw(packed_up.len()).unwrap();
let scales_up_dev = runtime.alloc_raw(scales_up.len() * 2).unwrap();
let output_elements = m * n;
let gate_out_dev = runtime.alloc_raw(output_elements * 2).unwrap();
let up_out_dev = runtime.alloc_raw(output_elements * 2).unwrap();
let ref_out_dev = runtime.alloc_raw(output_elements * 2).unwrap();
let fused_out_dev = runtime.alloc_raw(output_elements * 2).unwrap();
unsafe {
runtime.htod(as_bytes(&activation), activation_dev).unwrap();
runtime.htod(&packed_gate, packed_gate_dev).unwrap();
runtime
.htod(as_bytes(&scales_gate), scales_gate_dev)
.unwrap();
runtime.htod(&packed_up, packed_up_dev).unwrap();
runtime.htod(as_bytes(&scales_up), scales_up_dev).unwrap();
}
let device = DeviceId::cuda(0);
let a_shape = [m, k];
let a_strides = [k as i64, 1];
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let y_shape = [m, n];
let y_strides = [n as i64, 1];
let activation_view = TensorView::new(
device_ptr(activation_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
);
let packed_gate_view = TensorView::new(
device_ptr(packed_gate_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
);
let scales_gate_view = TensorView::new(
device_ptr(scales_gate_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
);
let packed_up_view = TensorView::new(
device_ptr(packed_up_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
);
let scales_up_view = TensorView::new(
device_ptr(scales_up_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
);
let mut gate_out = TensorMut::new(
device_ptr_mut(gate_out_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
);
let mut up_out = TensorMut::new(
device_ptr_mut(up_out_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
);
let fused_out = TensorMut::new(
device_ptr_mut(fused_out_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
);
let gemv_kernel = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: false,
decomposed_silu: false,
rmsnorm_prologue: false,
rmsnorm_epsilon: 1e-5,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
if m == 1 {
let selection = select_f16_gemv_variant(k, n, block_size, true, false);
assert_eq!(
selection.variant,
F16GemvVariant::General,
"gate/up decode projections must use the general GEMV as the reference"
);
gemv_kernel
.launch_f16_gemv_variant(
&activation_view,
&packed_gate_view,
&scales_gate_view,
true,
None,
None,
&mut gate_out,
k_blocks,
blob_size,
k_blocks.div_ceil(2),
selection,
)
.unwrap();
gemv_kernel
.launch_f16_gemv_variant(
&activation_view,
&packed_up_view,
&scales_up_view,
true,
None,
None,
&mut up_out,
k_blocks,
blob_size,
k_blocks.div_ceil(2),
selection,
)
.unwrap();
} else {
gemv_kernel
.launch_f16_gemm(
&activation_view,
&packed_gate_view,
&scales_gate_view,
true,
None,
None,
&mut gate_out,
m,
k_blocks,
gemv_kernel.block_size * gemv_kernel.bits / 8,
0,
)
.unwrap();
gemv_kernel
.launch_f16_gemm(
&activation_view,
&packed_up_view,
&scales_up_view,
true,
None,
None,
&mut up_out,
m,
k_blocks,
gemv_kernel.block_size * gemv_kernel.bits / 8,
0,
)
.unwrap();
}
let ref_function = runtime
.nvrtc_function(
"matmul_nbits_ref_silu_mul",
REF_SILU_MUL_SRC,
"ref_silu_mul_f16",
)
.unwrap();
let gate_out_ptr = cuptr(device_ptr(gate_out_dev).0);
let up_out_ptr = cuptr(device_ptr(up_out_dev).0);
let ref_out_ptr = cuptr(device_ptr(ref_out_dev).0);
let output_elements_i32 = output_elements as i32;
let mut ref_builder = runtime.stream().launch_builder(&ref_function);
ref_builder
.arg(&gate_out_ptr)
.arg(&up_out_ptr)
.arg(&ref_out_ptr)
.arg(&output_elements_i32);
unsafe {
ref_builder.launch(LaunchConfig {
grid_dim: (output_elements.div_ceil(256) as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
}
.unwrap();
let inputs = [
activation_view,
packed_gate_view,
scales_gate_view,
packed_up_view,
scales_up_view,
];
let mut outputs = [fused_out];
gemv_kernel
.run_f16_gate_up_swiglu(&inputs, &mut outputs, None)
.unwrap();
assert_eq!(
gemv_kernel.last_call_capture_safe.load(Ordering::Relaxed),
m == 1,
"only M=1 decode may be advertised capture-safe"
);
runtime.synchronize().unwrap();
let mut reference = vec![f16::ZERO; output_elements];
let mut fused = vec![f16::ZERO; output_elements];
unsafe {
runtime
.dtoh(as_bytes_mut(&mut reference), ref_out_dev)
.unwrap();
runtime
.dtoh(as_bytes_mut(&mut fused), fused_out_dev)
.unwrap();
runtime.free_raw(activation_dev).unwrap();
runtime.free_raw(packed_gate_dev).unwrap();
runtime.free_raw(scales_gate_dev).unwrap();
runtime.free_raw(packed_up_dev).unwrap();
runtime.free_raw(scales_up_dev).unwrap();
runtime.free_raw(gate_out_dev).unwrap();
runtime.free_raw(up_out_dev).unwrap();
runtime.free_raw(ref_out_dev).unwrap();
runtime.free_raw(fused_out_dev).unwrap();
}
let marlin_prefill = m > 1
&& marlin_gemm::marlin_m_gt_1_enabled()
&& marlin_gemm::device_supports_marlin(runtime.capabilities().compute_capability())
&& k.is_multiple_of(16)
&& k.is_multiple_of(block_size);
let mut max_ulp = 0i64;
let mut worst = 0usize;
for index in 0..output_elements {
let ulp = f16_ulp_diff(fused[index].to_bits(), reference[index].to_bits());
if ulp > max_ulp {
max_ulp = ulp;
worst = index;
}
if !marlin_prefill {
assert_eq!(
fused[index].to_bits(),
reference[index].to_bits(),
"paired gate/up SwiGLU diverged at M={m}, K={k}, N={n}, row={}, \
column={}: fused={:?} reference={:?}",
index / n,
index % n,
fused[index],
reference[index]
);
}
}
if marlin_prefill {
assert!(
max_ulp <= 2,
"paired gate/up SwiGLU Marlin prefill exceeded the measured 2-ULP \
sentinel bound at M={m}, K={k}, N={n}, row={}, column={}: fused={:?} \
reference={:?}, max_ulp={max_ulp}",
worst / n,
worst % n,
fused[worst],
reference[worst]
);
}
}
}
fn f16_ulp_diff(a: u16, b: u16) -> i64 {
let key = |bits: u16| -> i64 {
let mag = (bits & 0x7fff) as i64;
if bits & 0x8000 != 0 { -mag } else { mag }
};
(key(a) - key(b)).abs()
}
fn gate_up_swiglu_two_op_ulp_case(
runtime: &Arc<CudaRuntime>,
m: usize,
k: usize,
n: usize,
seed: u64,
) -> (i64, usize, usize) {
const REF_SILU_MUL_SRC: &str = r#"
#include <cuda_fp16.h>
__device__ float ref_op_silu(float x) {
if (x >= 0.0f) {
const float denominator = __fadd_rn(1.0f, (float)exp((double)-x));
return __fdiv_rn(x, denominator);
}
const float e = (float)exp((double)x);
const float numerator = __fmul_rn(x, e);
return __fdiv_rn(numerator, __fadd_rn(1.0f, e));
}
extern "C" __global__ void ref_silu_mul_f16(
const __half* g, const __half* u, __half* y, const int n) {
const int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
y[i] = __float2half_rn(
__fmul_rn(ref_op_silu(__half2float(g[i])), __half2float(u[i])));
}
}
"#;
let block_size = 32usize;
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let mut state = seed;
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let pack = |next: &mut dyn FnMut() -> f32| -> Vec<u8> {
let mut quant = vec![0u8; n * k];
for value in quant.iter_mut() {
*value = ((next() * 0.5 + 0.5) * 15.0).round().clamp(0.0, 15.0) as u8;
}
let mut packed = vec![0u8; n * k_blocks * blob_size];
for col in 0..n {
for block in 0..k_blocks {
for pair in 0..blob_size {
let low = quant[col * k + block * block_size + pair * 2] & 15;
let high = quant[col * k + block * block_size + pair * 2 + 1] & 15;
packed[(col * k_blocks + block) * blob_size + pair] = low | (high << 4);
}
}
}
packed
};
let activation: Vec<f16> = (0..m * k).map(|_| f16::from_f32(next())).collect();
let packed_gate = pack(&mut next);
let scales_gate: Vec<f16> = (0..n * k_blocks)
.map(|_| f16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5)))
.collect();
let packed_up = pack(&mut next);
let scales_up: Vec<f16> = (0..n * k_blocks)
.map(|_| f16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5)))
.collect();
let activation_dev = runtime.alloc_raw(activation.len() * 2).unwrap();
let packed_gate_dev = runtime.alloc_raw(packed_gate.len()).unwrap();
let scales_gate_dev = runtime.alloc_raw(scales_gate.len() * 2).unwrap();
let packed_up_dev = runtime.alloc_raw(packed_up.len()).unwrap();
let scales_up_dev = runtime.alloc_raw(scales_up.len() * 2).unwrap();
let output_elements = m * n;
let gate_out_dev = runtime.alloc_raw(output_elements * 2).unwrap();
let up_out_dev = runtime.alloc_raw(output_elements * 2).unwrap();
let ref_out_dev = runtime.alloc_raw(output_elements * 2).unwrap();
let fused_out_dev = runtime.alloc_raw(output_elements * 2).unwrap();
unsafe {
runtime.htod(as_bytes(&activation), activation_dev).unwrap();
runtime.htod(&packed_gate, packed_gate_dev).unwrap();
runtime
.htod(as_bytes(&scales_gate), scales_gate_dev)
.unwrap();
runtime.htod(&packed_up, packed_up_dev).unwrap();
runtime.htod(as_bytes(&scales_up), scales_up_dev).unwrap();
}
let device = DeviceId::cuda(0);
let a_shape = [m, k];
let a_strides = [k as i64, 1];
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let y_shape = [m, n];
let y_strides = [n as i64, 1];
let activation_view = TensorView::new(
device_ptr(activation_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
);
let packed_gate_view = TensorView::new(
device_ptr(packed_gate_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
);
let scales_gate_view = TensorView::new(
device_ptr(scales_gate_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
);
let packed_up_view = TensorView::new(
device_ptr(packed_up_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
);
let scales_up_view = TensorView::new(
device_ptr(scales_up_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
);
let mut gate_out = TensorMut::new(
device_ptr_mut(gate_out_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
);
let mut up_out = TensorMut::new(
device_ptr_mut(up_out_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
);
let fused_out = TensorMut::new(
device_ptr_mut(fused_out_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
);
let gemv_kernel = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: false,
decomposed_silu: false,
rmsnorm_prologue: false,
rmsnorm_epsilon: 1e-5,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
if m == 1 {
let selection = select_f16_gemv_variant(k, n, block_size, true, false);
gemv_kernel
.launch_f16_gemv_variant(
&activation_view,
&packed_gate_view,
&scales_gate_view,
true,
None,
None,
&mut gate_out,
k_blocks,
blob_size,
k_blocks.div_ceil(2),
selection,
)
.unwrap();
gemv_kernel
.launch_f16_gemv_variant(
&activation_view,
&packed_up_view,
&scales_up_view,
true,
None,
None,
&mut up_out,
k_blocks,
blob_size,
k_blocks.div_ceil(2),
selection,
)
.unwrap();
} else {
gemv_kernel
.launch_f16_gemm(
&activation_view,
&packed_gate_view,
&scales_gate_view,
true,
None,
None,
&mut gate_out,
m,
k_blocks,
gemv_kernel.block_size * gemv_kernel.bits / 8,
0,
)
.unwrap();
gemv_kernel
.launch_f16_gemm(
&activation_view,
&packed_up_view,
&scales_up_view,
true,
None,
None,
&mut up_out,
m,
k_blocks,
gemv_kernel.block_size * gemv_kernel.bits / 8,
0,
)
.unwrap();
}
let ref_function = runtime
.nvrtc_function(
"matmul_nbits_ref_silu_mul",
REF_SILU_MUL_SRC,
"ref_silu_mul_f16",
)
.unwrap();
let gate_out_ptr = cuptr(device_ptr(gate_out_dev).0);
let up_out_ptr = cuptr(device_ptr(up_out_dev).0);
let ref_out_ptr = cuptr(device_ptr(ref_out_dev).0);
let output_elements_i32 = output_elements as i32;
let mut ref_builder = runtime.stream().launch_builder(&ref_function);
ref_builder
.arg(&gate_out_ptr)
.arg(&up_out_ptr)
.arg(&ref_out_ptr)
.arg(&output_elements_i32);
unsafe {
ref_builder.launch(LaunchConfig {
grid_dim: (output_elements.div_ceil(256) as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
}
.unwrap();
let inputs = [
activation_view,
packed_gate_view,
scales_gate_view,
packed_up_view,
scales_up_view,
];
let mut outputs = [fused_out];
gemv_kernel
.run_f16_gate_up_swiglu(&inputs, &mut outputs, None)
.unwrap();
runtime.synchronize().unwrap();
let mut reference = vec![f16::ZERO; output_elements];
let mut fused = vec![f16::ZERO; output_elements];
unsafe {
runtime
.dtoh(as_bytes_mut(&mut reference), ref_out_dev)
.unwrap();
runtime
.dtoh(as_bytes_mut(&mut fused), fused_out_dev)
.unwrap();
runtime.free_raw(activation_dev).unwrap();
runtime.free_raw(packed_gate_dev).unwrap();
runtime.free_raw(scales_gate_dev).unwrap();
runtime.free_raw(packed_up_dev).unwrap();
runtime.free_raw(scales_up_dev).unwrap();
runtime.free_raw(gate_out_dev).unwrap();
runtime.free_raw(up_out_dev).unwrap();
runtime.free_raw(ref_out_dev).unwrap();
runtime.free_raw(fused_out_dev).unwrap();
}
let mut max_ulp = 0i64;
let mut nonzero = 0usize;
for index in 0..output_elements {
let d = f16_ulp_diff(fused[index].to_bits(), reference[index].to_bits());
if d > 0 {
nonzero += 1;
}
if d > max_ulp {
max_ulp = d;
}
}
(max_ulp, nonzero, output_elements)
}
#[test]
fn fp16_gate_up_swiglu_two_op_ulp_bound_sweep() {
let Some(runtime) = runtime() else {
eprintln!("skipping gate/up SwiGLU ULP sweep: CUDA runtime unavailable");
return;
};
if runtime
.require_nvrtc_half_headers("matmul_nbits_gemv_f16")
.is_err()
{
eprintln!("skipping gate/up SwiGLU ULP sweep: fp16 NVRTC headers unavailable");
return;
}
let shapes = [
(QWEN_DOWN_N, QWEN_DOWN_K), (2048usize, 5632usize), (1024, 2816),
(896, 896), (512, 1536),
(128, 256), (96, 77), ];
let ms = [1usize];
let seeds: [u64; 8] = [
0x0bad_c0de_dead_beef,
0x1234_5678_9abc_def0,
0xdead_beef_cafe_babe,
0x0f0f_0f0f_f0f0_f0f0,
0xa5a5_5a5a_1111_2222,
0x9e37_79b9_7f4a_7c15,
0xc2b2_ae3d_27d4_eb4f,
0xff51_afd7_ed55_8ccd,
];
let mut global_max = 0i64;
let mut m1_max = 0i64;
let mut m1_nonzero_cases = 0usize;
let mut m1_cases = 0usize;
let mut worst = (0usize, 0usize, 0usize, 0u64); for (k, n) in shapes {
for m in ms {
for seed in seeds {
let (max_ulp, nonzero, total) =
gate_up_swiglu_two_op_ulp_case(&runtime, m, k, n, seed);
assert!(total > 0);
if m == 1 {
m1_cases += 1;
if nonzero > 0 {
m1_nonzero_cases += 1;
}
if max_ulp > m1_max {
m1_max = max_ulp;
}
}
if max_ulp > global_max {
global_max = max_ulp;
worst = (m, k, n, seed);
}
}
}
}
eprintln!(
"gate/up SwiGLU two-op ULP sweep: max_ulp={global_max} \
(worst: M={} K={} N={} seed={:#018x}); \
M==1: max_ulp={m1_max}, {m1_nonzero_cases}/{m1_cases} cases had >=1 ULP divergence",
worst.0, worst.1, worst.2, worst.3
);
assert_eq!(
global_max, m1_max,
"sweep is M==1-only; global and M==1 maxima must coincide"
);
assert!(
global_max <= 2,
"plain fused gate/up SwiGLU exceeded the measured 2 ULP bound vs the \
two-op reference: max_ulp={global_max} at M={} K={} N={} seed={:#018x}",
worst.0,
worst.1,
worst.2,
worst.3
);
}
fn gate_up_swiglu_rmsnorm_two_op_ulp_case(
runtime: &Arc<CudaRuntime>,
m: usize,
k: usize,
n: usize,
seed: u64,
gamma_dtype: DataType,
) -> (i64, usize, usize) {
const REF_SILU_MUL_SRC: &str = r#"
#include <cuda_fp16.h>
__device__ float ref_op_silu(float x) {
if (x >= 0.0f) {
const float denominator = __fadd_rn(1.0f, (float)exp((double)-x));
return __fdiv_rn(x, denominator);
}
const float e = (float)exp((double)x);
const float numerator = __fmul_rn(x, e);
return __fdiv_rn(numerator, __fadd_rn(1.0f, e));
}
extern "C" __global__ void ref_silu_mul_f16(
const __half* g, const __half* u, __half* y, const int n) {
const int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
y[i] = __float2half_rn(
__fmul_rn(ref_op_silu(__half2float(g[i])), __half2float(u[i])));
}
}
"#;
let epsilon = 1e-5f32;
let block_size = 32usize;
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let mut state = seed;
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let pack = |next: &mut dyn FnMut() -> f32| -> Vec<u8> {
let mut quant = vec![0u8; n * k];
for value in quant.iter_mut() {
*value = ((next() * 0.5 + 0.5) * 15.0).round().clamp(0.0, 15.0) as u8;
}
let mut packed = vec![0u8; n * k_blocks * blob_size];
for col in 0..n {
for block in 0..k_blocks {
for pair in 0..blob_size {
let low = quant[col * k + block * block_size + pair * 2] & 15;
let high = quant[col * k + block * block_size + pair * 2 + 1] & 15;
packed[(col * k_blocks + block) * blob_size + pair] = low | (high << 4);
}
}
}
packed
};
let activation: Vec<f16> = (0..m * k).map(|_| f16::from_f32(next())).collect();
let packed_gate = pack(&mut next);
let scales_gate: Vec<f16> = (0..n * k_blocks)
.map(|_| f16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5)))
.collect();
let packed_up = pack(&mut next);
let scales_up: Vec<f16> = (0..n * k_blocks)
.map(|_| f16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5)))
.collect();
let gamma_is_f32 = gamma_dtype == DataType::Float32;
let gamma_f32: Vec<f32> = (0..k).map(|_| 0.5 + 0.5 * (next() * 0.5 + 0.5)).collect();
let gamma_bytes: Vec<u8> = if gamma_is_f32 {
gamma_f32.iter().flat_map(|v| v.to_le_bytes()).collect()
} else {
gamma_f32
.iter()
.flat_map(|v| f16::from_f32(*v).to_le_bytes())
.collect()
};
let activation_dev = runtime.alloc_raw(activation.len() * 2).unwrap();
let packed_gate_dev = runtime.alloc_raw(packed_gate.len()).unwrap();
let scales_gate_dev = runtime.alloc_raw(scales_gate.len() * 2).unwrap();
let packed_up_dev = runtime.alloc_raw(packed_up.len()).unwrap();
let scales_up_dev = runtime.alloc_raw(scales_up.len() * 2).unwrap();
let gamma_dev = runtime.alloc_raw(gamma_bytes.len()).unwrap();
let normalized_dev = runtime.alloc_raw(m * k * 2).unwrap();
let output_elements = m * n;
let gate_out_dev = runtime.alloc_raw(output_elements * 2).unwrap();
let up_out_dev = runtime.alloc_raw(output_elements * 2).unwrap();
let ref_out_dev = runtime.alloc_raw(output_elements * 2).unwrap();
let fused_out_dev = runtime.alloc_raw(output_elements * 2).unwrap();
unsafe {
runtime.htod(as_bytes(&activation), activation_dev).unwrap();
runtime.htod(&packed_gate, packed_gate_dev).unwrap();
runtime
.htod(as_bytes(&scales_gate), scales_gate_dev)
.unwrap();
runtime.htod(&packed_up, packed_up_dev).unwrap();
runtime.htod(as_bytes(&scales_up), scales_up_dev).unwrap();
runtime.htod(&gamma_bytes, gamma_dev).unwrap();
}
let device = DeviceId::cuda(0);
let a_shape = [m, k];
let a_strides = [k as i64, 1];
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let gamma_shape = [k];
let gamma_strides = [1i64];
let y_shape = [m, n];
let y_strides = [n as i64, 1];
let activation_view = TensorView::new(
device_ptr(activation_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
);
let normalized_view = TensorView::new(
device_ptr(normalized_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
);
let packed_gate_view = TensorView::new(
device_ptr(packed_gate_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
);
let scales_gate_view = TensorView::new(
device_ptr(scales_gate_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
);
let packed_up_view = TensorView::new(
device_ptr(packed_up_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
);
let scales_up_view = TensorView::new(
device_ptr(scales_up_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
);
let gamma_view = TensorView::new(
device_ptr(gamma_dev),
gamma_dtype,
&gamma_shape,
&gamma_strides,
device,
);
let mut gate_out = TensorMut::new(
device_ptr_mut(gate_out_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
);
let mut up_out = TensorMut::new(
device_ptr_mut(up_out_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
);
let fused_out = TensorMut::new(
device_ptr_mut(fused_out_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
);
let plain_kernel = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: false,
decomposed_silu: false,
rmsnorm_prologue: false,
rmsnorm_epsilon: epsilon,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
let fused_kernel = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: true,
decomposed_silu: false,
rmsnorm_prologue: true,
rmsnorm_epsilon: epsilon,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
plain_kernel
.launch_rmsnorm_prefill(
&activation_view,
&gamma_view,
cuptr(device_ptr(normalized_dev).0),
m,
)
.unwrap();
if m == 1 {
let selection = select_f16_gemv_variant(k, n, block_size, true, false);
plain_kernel
.launch_f16_gemv_variant(
&normalized_view,
&packed_gate_view,
&scales_gate_view,
true,
None,
None,
&mut gate_out,
k_blocks,
blob_size,
k_blocks.div_ceil(2),
selection,
)
.unwrap();
plain_kernel
.launch_f16_gemv_variant(
&normalized_view,
&packed_up_view,
&scales_up_view,
true,
None,
None,
&mut up_out,
k_blocks,
blob_size,
k_blocks.div_ceil(2),
selection,
)
.unwrap();
} else {
plain_kernel
.launch_f16_gemm(
&normalized_view,
&packed_gate_view,
&scales_gate_view,
true,
None,
None,
&mut gate_out,
m,
k_blocks,
plain_kernel.block_size * plain_kernel.bits / 8,
0,
)
.unwrap();
plain_kernel
.launch_f16_gemm(
&normalized_view,
&packed_up_view,
&scales_up_view,
true,
None,
None,
&mut up_out,
m,
k_blocks,
plain_kernel.block_size * plain_kernel.bits / 8,
0,
)
.unwrap();
}
let ref_function = runtime
.nvrtc_function(
"matmul_nbits_ref_silu_mul",
REF_SILU_MUL_SRC,
"ref_silu_mul_f16",
)
.unwrap();
let gate_out_ptr = cuptr(device_ptr(gate_out_dev).0);
let up_out_ptr = cuptr(device_ptr(up_out_dev).0);
let ref_out_ptr = cuptr(device_ptr(ref_out_dev).0);
let output_elements_i32 = output_elements as i32;
let mut ref_builder = runtime.stream().launch_builder(&ref_function);
ref_builder
.arg(&gate_out_ptr)
.arg(&up_out_ptr)
.arg(&ref_out_ptr)
.arg(&output_elements_i32);
unsafe {
ref_builder.launch(LaunchConfig {
grid_dim: (output_elements.div_ceil(256) as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
}
.unwrap();
let inputs = [
activation_view,
packed_gate_view,
scales_gate_view,
packed_up_view,
scales_up_view,
gamma_view,
];
let mut outputs = [fused_out];
fused_kernel
.run_f16_gate_up_swiglu(&inputs, &mut outputs, None)
.unwrap();
runtime.synchronize().unwrap();
let mut reference = vec![f16::ZERO; output_elements];
let mut fused = vec![f16::ZERO; output_elements];
unsafe {
runtime
.dtoh(as_bytes_mut(&mut reference), ref_out_dev)
.unwrap();
runtime
.dtoh(as_bytes_mut(&mut fused), fused_out_dev)
.unwrap();
runtime.free_raw(activation_dev).unwrap();
runtime.free_raw(packed_gate_dev).unwrap();
runtime.free_raw(scales_gate_dev).unwrap();
runtime.free_raw(packed_up_dev).unwrap();
runtime.free_raw(scales_up_dev).unwrap();
runtime.free_raw(gamma_dev).unwrap();
runtime.free_raw(normalized_dev).unwrap();
runtime.free_raw(gate_out_dev).unwrap();
runtime.free_raw(up_out_dev).unwrap();
runtime.free_raw(ref_out_dev).unwrap();
runtime.free_raw(fused_out_dev).unwrap();
}
let mut max_ulp = 0i64;
let mut nonzero = 0usize;
for index in 0..output_elements {
let d = f16_ulp_diff(fused[index].to_bits(), reference[index].to_bits());
if d > 0 {
nonzero += 1;
}
if d > max_ulp {
max_ulp = d;
}
}
(max_ulp, nonzero, output_elements)
}
#[test]
fn fp16_gate_up_swiglu_rmsnorm_two_op_ulp_bound_sweep() {
let Some(runtime) = runtime() else {
eprintln!("skipping gate/up SwiGLU rmsnorm ULP sweep: CUDA runtime unavailable");
return;
};
if runtime
.require_nvrtc_half_headers("matmul_nbits_gemv_f16")
.is_err()
{
eprintln!("skipping gate/up SwiGLU rmsnorm ULP sweep: fp16 NVRTC headers unavailable");
return;
}
let shapes = [
(QWEN_DOWN_N, QWEN_DOWN_K), (2048usize, 5632usize),
(1024, 2816),
(896, 896),
(512, 1536),
];
let ms = [1usize, 2, 5, 8];
let seeds: [u64; 6] = [
0x0bad_c0de_dead_beef,
0x1234_5678_9abc_def0,
0xdead_beef_cafe_babe,
0xa5a5_5a5a_1111_2222,
0x9e37_79b9_7f4a_7c15,
0xff51_afd7_ed55_8ccd,
];
let gammas = [DataType::Float16, DataType::Float32];
let mut global_max = 0i64;
let mut m1_max = 0i64;
let mut m1_single_warp_max = 0i64;
let mut m1_splitk_max = 0i64;
let mut m1_nonzero_cases = 0usize;
let mut m1_cases = 0usize;
let mut m1_worst = (0usize, 0usize, 0u64, DataType::Float16);
let mut worst = (0usize, 0usize, 0usize, 0u64, DataType::Float16);
let caps = runtime.capabilities();
for gamma in gammas {
for (k, n) in shapes {
let reference_uses_splitk = use_f16_symmetric_splitk(
k,
n,
caps.multiprocessor_count(),
caps.max_threads_per_block(),
);
for m in ms {
for seed in seeds {
let (max_ulp, nonzero, total) =
gate_up_swiglu_rmsnorm_two_op_ulp_case(&runtime, m, k, n, seed, gamma);
assert!(total > 0);
if m == 1 {
m1_cases += 1;
if nonzero > 0 {
m1_nonzero_cases += 1;
}
if max_ulp > m1_max {
m1_max = max_ulp;
m1_worst = (k, n, seed, gamma);
}
if reference_uses_splitk {
m1_splitk_max = m1_splitk_max.max(max_ulp);
} else {
m1_single_warp_max = m1_single_warp_max.max(max_ulp);
}
}
if max_ulp > global_max {
global_max = max_ulp;
worst = (m, k, n, seed, gamma);
}
}
}
}
}
eprintln!(
"gate/up SwiGLU rmsnorm two-op ULP sweep: max_ulp={global_max} \
(worst: M={} K={} N={} seed={:#018x} gamma={:?}); \
M==1: max_ulp={m1_max} (worst: K={} N={} seed={:#018x} gamma={:?}), \
{m1_nonzero_cases}/{m1_cases} cases had >=1 ULP divergence",
worst.0,
worst.1,
worst.2,
worst.3,
worst.4,
m1_worst.0,
m1_worst.1,
m1_worst.2,
m1_worst.3
);
assert_eq!(
m1_single_warp_max, 0,
"rmsnorm fused gate/up SwiGLU decode GEMV must be byte-identical to \
the two-op reference at M==1 when both use the single-warp reduction, \
but observed {m1_single_warp_max} ULP"
);
assert!(
m1_splitk_max <= 3,
"rmsnorm fused gate/up SwiGLU decode GEMV exceeded the measured 3-ULP \
bound vs the standalone split-K two-op reference at M==1: \
max_ulp={m1_splitk_max}"
);
assert!(
global_max <= 4,
"rmsnorm fused gate/up SwiGLU ULP vs two-op unexpectedly large: \
max_ulp={global_max}"
);
}
#[test]
fn fused_gate_up_swiglu_rmsnorm_is_bit_exact_to_two_step_path() {
run_fused_gate_up_swiglu_rmsnorm_parity(DataType::Float16, false);
}
#[test]
fn fused_gate_up_swiglu_rmsnorm_fp32_gamma_is_bit_exact_to_two_step_path() {
run_fused_gate_up_swiglu_rmsnorm_parity(DataType::Float32, false);
}
#[test]
fn fused_gate_up_swiglu_rmsnorm_zero_points_is_bit_exact_to_two_step_path() {
run_fused_gate_up_swiglu_rmsnorm_parity(DataType::Float32, true);
}
fn run_fused_gate_up_swiglu_rmsnorm_parity(gamma_dtype: DataType, explicit_zp: bool) {
let Some(runtime) = runtime() else {
eprintln!("skipping gate/up SwiGLU RMS-norm parity test: CUDA runtime unavailable");
return;
};
if runtime
.require_nvrtc_half_headers("matmul_nbits_gemv_f16")
.is_err()
{
eprintln!(
"skipping gate/up SwiGLU RMS-norm parity test: fp16 NVRTC headers unavailable"
);
return;
}
let epsilon = 1e-5f32;
for (m, k, n) in [
(1usize, 896usize, 2432usize),
(1, 3584, 4864),
(5, 896, 2432),
] {
let block_size = 32usize;
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let mut state = 0xf00d_1ceb_00da_5555u64;
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let pack = |next: &mut dyn FnMut() -> f32| -> Vec<u8> {
let mut quant = vec![0u8; n * k];
for value in quant.iter_mut() {
*value = ((next() * 0.5 + 0.5) * 15.0).round().clamp(0.0, 15.0) as u8;
}
let mut packed = vec![0u8; n * k_blocks * blob_size];
for col in 0..n {
for block in 0..k_blocks {
for pair in 0..blob_size {
let low = quant[col * k + block * block_size + pair * 2] & 15;
let high = quant[col * k + block * block_size + pair * 2 + 1] & 15;
packed[(col * k_blocks + block) * blob_size + pair] = low | (high << 4);
}
}
}
packed
};
let activation: Vec<f16> = (0..m * k).map(|_| f16::from_f32(next())).collect();
let packed_gate = pack(&mut next);
let scales_gate: Vec<f16> = (0..n * k_blocks)
.map(|_| f16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5)))
.collect();
let packed_up = pack(&mut next);
let scales_up: Vec<f16> = (0..n * k_blocks)
.map(|_| f16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5)))
.collect();
let zp_row_bytes = k_blocks.div_ceil(2);
let pack_zp = |next: &mut dyn FnMut() -> f32| -> Vec<u8> {
let mut zp = vec![0u8; n * zp_row_bytes];
for col in 0..n {
for block in 0..k_blocks {
let nibble =
(((next() * 0.5 + 0.5) * 15.0).round().clamp(0.0, 15.0) as u8) & 15;
let byte = &mut zp[col * zp_row_bytes + (block >> 1)];
if block & 1 == 1 {
*byte = (*byte & 0x0f) | (nibble << 4);
} else {
*byte = (*byte & 0xf0) | nibble;
}
}
}
zp
};
let zp_gate = pack_zp(&mut next);
let zp_up = pack_zp(&mut next);
let gamma_is_f32 = gamma_dtype == DataType::Float32;
let gamma_f32: Vec<f32> = (0..k).map(|_| 0.5 + 0.5 * (next() * 0.5 + 0.5)).collect();
let gamma_bytes: Vec<u8> = if gamma_is_f32 {
gamma_f32.iter().flat_map(|v| v.to_le_bytes()).collect()
} else {
gamma_f32
.iter()
.flat_map(|v| f16::from_f32(*v).to_le_bytes())
.collect()
};
let activation_dev = runtime.alloc_raw(activation.len() * 2).unwrap();
let packed_gate_dev = runtime.alloc_raw(packed_gate.len()).unwrap();
let scales_gate_dev = runtime.alloc_raw(scales_gate.len() * 2).unwrap();
let packed_up_dev = runtime.alloc_raw(packed_up.len()).unwrap();
let scales_up_dev = runtime.alloc_raw(scales_up.len() * 2).unwrap();
let gamma_dev = runtime.alloc_raw(gamma_bytes.len()).unwrap();
let zp_gate_dev = runtime.alloc_raw(zp_gate.len()).unwrap();
let zp_up_dev = runtime.alloc_raw(zp_up.len()).unwrap();
let normalized_dev = runtime.alloc_raw(m * k * 2).unwrap();
let output_elements = m * n;
let ref_out_dev = runtime.alloc_raw(output_elements * 2).unwrap();
let fused_out_dev = runtime.alloc_raw(output_elements * 2).unwrap();
unsafe {
runtime.htod(as_bytes(&activation), activation_dev).unwrap();
runtime.htod(&packed_gate, packed_gate_dev).unwrap();
runtime
.htod(as_bytes(&scales_gate), scales_gate_dev)
.unwrap();
runtime.htod(&packed_up, packed_up_dev).unwrap();
runtime.htod(as_bytes(&scales_up), scales_up_dev).unwrap();
runtime.htod(&gamma_bytes, gamma_dev).unwrap();
if explicit_zp {
runtime.htod(&zp_gate, zp_gate_dev).unwrap();
runtime.htod(&zp_up, zp_up_dev).unwrap();
}
}
let device = DeviceId::cuda(0);
let a_shape = [m, k];
let a_strides = [k as i64, 1];
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let gamma_shape = [k];
let gamma_strides = [1i64];
let y_shape = [m, n];
let y_strides = [n as i64, 1];
let activation_view = TensorView::new(
device_ptr(activation_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
);
let _normalized_view = TensorView::new(
device_ptr(normalized_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
);
let packed_gate_view = TensorView::new(
device_ptr(packed_gate_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
);
let scales_gate_view = TensorView::new(
device_ptr(scales_gate_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
);
let packed_up_view = TensorView::new(
device_ptr(packed_up_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
);
let scales_up_view = TensorView::new(
device_ptr(scales_up_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
);
let gamma_view = TensorView::new(
device_ptr(gamma_dev),
gamma_dtype,
&gamma_shape,
&gamma_strides,
device,
);
let zp_shape = [n, zp_row_bytes];
let zp_strides = [zp_row_bytes as i64, 1];
let zp_gate_view = TensorView::new(
device_ptr(zp_gate_dev),
DataType::Uint8,
&zp_shape,
&zp_strides,
device,
);
let zp_up_view = TensorView::new(
device_ptr(zp_up_dev),
DataType::Uint8,
&zp_shape,
&zp_strides,
device,
);
let fused_out = TensorMut::new(
device_ptr_mut(fused_out_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
);
let plain_swiglu = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: true,
decomposed_silu: false,
rmsnorm_prologue: false,
rmsnorm_epsilon: epsilon,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
let fused_swiglu = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: true,
decomposed_silu: false,
rmsnorm_prologue: true,
rmsnorm_epsilon: epsilon,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
plain_swiglu
.launch_rmsnorm_prefill(
&activation_view,
&gamma_view,
cuptr(device_ptr(normalized_dev).0),
m,
)
.unwrap();
{
let norm_row_shape = [1usize, k];
let norm_row_strides = [k as i64, 1];
let out_row_shape = [1usize, n];
let out_row_strides = [n as i64, 1];
let norm_row_bytes = (k * 2) as CUdeviceptr; let out_row_bytes = (n * 2) as CUdeviceptr; for row in 0..m {
let norm_row = TensorView::new(
device_ptr(normalized_dev + row as CUdeviceptr * norm_row_bytes),
DataType::Float16,
&norm_row_shape,
&norm_row_strides,
device,
);
let out_row = TensorMut::new(
device_ptr_mut(ref_out_dev + row as CUdeviceptr * out_row_bytes),
DataType::Float16,
&out_row_shape,
&out_row_strides,
device,
);
let mut ref_outputs = [out_row];
let ref_inputs_base = [
norm_row,
packed_gate_view,
scales_gate_view,
packed_up_view,
scales_up_view,
];
if explicit_zp {
let ref_inputs = [
ref_inputs_base[0],
ref_inputs_base[1],
ref_inputs_base[2],
ref_inputs_base[3],
ref_inputs_base[4],
TensorView::absent(DataType::Float16),
zp_gate_view,
zp_up_view,
];
plain_swiglu
.run_f16_gate_up_swiglu(&ref_inputs, &mut ref_outputs, None)
.unwrap();
} else {
plain_swiglu
.run_f16_gate_up_swiglu(&ref_inputs_base, &mut ref_outputs, None)
.unwrap();
}
}
}
{
let mut fused_outputs = [fused_out];
let fused_inputs_base = [
activation_view,
packed_gate_view,
scales_gate_view,
packed_up_view,
scales_up_view,
gamma_view,
];
if explicit_zp {
let fused_inputs = [
fused_inputs_base[0],
fused_inputs_base[1],
fused_inputs_base[2],
fused_inputs_base[3],
fused_inputs_base[4],
fused_inputs_base[5],
zp_gate_view,
zp_up_view,
];
fused_swiglu
.run_f16_gate_up_swiglu(&fused_inputs, &mut fused_outputs, None)
.unwrap();
} else {
fused_swiglu
.run_f16_gate_up_swiglu(&fused_inputs_base, &mut fused_outputs, None)
.unwrap();
}
}
assert_eq!(
fused_swiglu.last_call_capture_safe.load(Ordering::Relaxed),
(1..=decode_gemv_loop_max_m()).contains(&m),
"capture-safe iff the decode-GEMV routing is per-row M==1 \
launches: M==1 or small-batch M<=decode_gemv_loop_max_m(); \
M>window still reaches the uncapturable prefill GEMM"
);
runtime.synchronize().unwrap();
let mut reference = vec![f16::ZERO; output_elements];
let mut fused = vec![f16::ZERO; output_elements];
unsafe {
runtime
.dtoh(as_bytes_mut(&mut reference), ref_out_dev)
.unwrap();
runtime
.dtoh(as_bytes_mut(&mut fused), fused_out_dev)
.unwrap();
runtime.free_raw(activation_dev).unwrap();
runtime.free_raw(packed_gate_dev).unwrap();
runtime.free_raw(scales_gate_dev).unwrap();
runtime.free_raw(packed_up_dev).unwrap();
runtime.free_raw(scales_up_dev).unwrap();
runtime.free_raw(gamma_dev).unwrap();
runtime.free_raw(zp_gate_dev).unwrap();
runtime.free_raw(zp_up_dev).unwrap();
runtime.free_raw(normalized_dev).unwrap();
runtime.free_raw(ref_out_dev).unwrap();
runtime.free_raw(fused_out_dev).unwrap();
}
for index in 0..output_elements {
assert_eq!(
fused[index].to_bits(),
reference[index].to_bits(),
"fused gate/up SwiGLU RMS prologue diverged at M={m}, K={k}, N={n}, \
row={}, column={}: fused={:?} reference={:?}",
index / n,
index % n,
fused[index],
reference[index]
);
}
}
}
#[test]
fn gate_up_swiglu_vec_is_bit_identical_to_scalar() {
let Some(runtime) = runtime() else {
eprintln!("skipping gate/up SwiGLU _vec bit-identity test: CUDA runtime unavailable");
return;
};
if runtime
.require_nvrtc_half_headers("matmul_nbits_gemv_f16")
.is_err()
{
eprintln!(
"skipping gate/up SwiGLU _vec bit-identity test: fp16 NVRTC headers unavailable"
);
return;
}
let epsilon = 1e-5f32;
for (m, k, n) in [
(1usize, 896usize, 2432usize),
(4, 3584, 4864),
(6, 5120, 13824),
(8, 5120, 13824),
] {
let block_size = 32usize;
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let mut state = 0x51de_face_c0de_1234u64 ^ ((m as u64) << 40);
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let pack = |next: &mut dyn FnMut() -> f32| -> Vec<u8> {
let mut quant = vec![0u8; n * k];
for value in quant.iter_mut() {
*value = ((next() * 0.5 + 0.5) * 15.0).round().clamp(0.0, 15.0) as u8;
}
let mut packed = vec![0u8; n * k_blocks * blob_size];
for col in 0..n {
for block in 0..k_blocks {
for pair in 0..blob_size {
let low = quant[col * k + block * block_size + pair * 2] & 15;
let high = quant[col * k + block * block_size + pair * 2 + 1] & 15;
packed[(col * k_blocks + block) * blob_size + pair] = low | (high << 4);
}
}
}
packed
};
let activation: Vec<f16> = (0..m * k).map(|_| f16::from_f32(next())).collect();
let packed_gate = pack(&mut next);
let scales_gate: Vec<f16> = (0..n * k_blocks)
.map(|_| f16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5)))
.collect();
let packed_up = pack(&mut next);
let scales_up: Vec<f16> = (0..n * k_blocks)
.map(|_| f16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5)))
.collect();
let gamma_f32: Vec<f32> = (0..k).map(|_| 0.5 + 0.5 * (next() * 0.5 + 0.5)).collect();
let activation_dev = runtime.alloc_raw(activation.len() * 2).unwrap();
let packed_gate_dev = runtime.alloc_raw(packed_gate.len()).unwrap();
let scales_gate_dev = runtime.alloc_raw(scales_gate.len() * 2).unwrap();
let packed_up_dev = runtime.alloc_raw(packed_up.len()).unwrap();
let scales_up_dev = runtime.alloc_raw(scales_up.len() * 2).unwrap();
unsafe {
runtime.htod(as_bytes(&activation), activation_dev).unwrap();
runtime.htod(&packed_gate, packed_gate_dev).unwrap();
runtime
.htod(as_bytes(&scales_gate), scales_gate_dev)
.unwrap();
runtime.htod(&packed_up, packed_up_dev).unwrap();
runtime.htod(as_bytes(&scales_up), scales_up_dev).unwrap();
}
let device = DeviceId::cuda(0);
let a_shape = [m, k];
let a_strides = [k as i64, 1];
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let gamma_shape = [k];
let gamma_strides = [1i64];
let y_shape = [m, n];
let y_strides = [n as i64, 1];
let output_elements = m * n;
let activation_view = TensorView::new(
device_ptr(activation_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
);
let packed_gate_view = TensorView::new(
device_ptr(packed_gate_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
);
let scales_gate_view = TensorView::new(
device_ptr(scales_gate_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
);
let packed_up_view = TensorView::new(
device_ptr(packed_up_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
);
let scales_up_view = TensorView::new(
device_ptr(scales_up_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
);
for rmsnorm in [false, true] {
for decomposed in [false, true] {
for gamma_dtype in [DataType::Float16, DataType::Float32] {
if !rmsnorm && gamma_dtype == DataType::Float32 {
continue;
}
let gamma_is_f32 = gamma_dtype == DataType::Float32;
let gamma_bytes: Vec<u8> = if gamma_is_f32 {
gamma_f32.iter().flat_map(|v| v.to_le_bytes()).collect()
} else {
gamma_f32
.iter()
.flat_map(|v| f16::from_f32(*v).to_le_bytes())
.collect()
};
let gamma_dev = runtime.alloc_raw(gamma_bytes.len()).unwrap();
unsafe {
runtime.htod(&gamma_bytes, gamma_dev).unwrap();
}
let gamma_view = TensorView::new(
device_ptr(gamma_dev),
gamma_dtype,
&gamma_shape,
&gamma_strides,
device,
);
let kernel = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: true,
decomposed_silu: decomposed,
rmsnorm_prologue: rmsnorm,
rmsnorm_epsilon: epsilon,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
let run_once = |vec_on: bool, occ_on: bool| -> Vec<u16> {
let out_dev = runtime.alloc_raw(output_elements * 2).unwrap();
let _guard = LeverEnvGuard::acquire();
unsafe {
std::env::set_var(
"ONNX_GENAI_GATEUP_VEC",
if vec_on { "1" } else { "0" },
);
std::env::set_var(
"ONNX_GENAI_GATEUP_OCC",
if occ_on { "1" } else { "0" },
);
}
let out = TensorMut::new(
device_ptr_mut(out_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
);
let mut outputs = [out];
if rmsnorm {
let inputs = [
activation_view,
packed_gate_view,
scales_gate_view,
packed_up_view,
scales_up_view,
gamma_view,
];
kernel
.run_f16_gate_up_swiglu(&inputs, &mut outputs, None)
.unwrap();
} else {
let inputs = [
activation_view,
packed_gate_view,
scales_gate_view,
packed_up_view,
scales_up_view,
];
kernel
.run_f16_gate_up_swiglu(&inputs, &mut outputs, None)
.unwrap();
}
runtime.synchronize().unwrap();
unsafe {
std::env::remove_var("ONNX_GENAI_GATEUP_VEC");
std::env::remove_var("ONNX_GENAI_GATEUP_OCC");
}
drop(_guard);
let mut host = vec![f16::ZERO; output_elements];
unsafe {
runtime.dtoh(as_bytes_mut(&mut host), out_dev).unwrap();
runtime.free_raw(out_dev).unwrap();
}
host.iter().map(|v| v.to_bits()).collect()
};
let reference = run_once(false, false);
let fused = run_once(true, false);
let occ = run_once(true, true);
unsafe {
runtime.free_raw(gamma_dev).unwrap();
}
for index in 0..output_elements {
assert_eq!(
fused[index],
reference[index],
"fused-symmetric gate/up _vec diverged at M={m}, K={k}, N={n}, \
rmsnorm={rmsnorm}, decomposed={decomposed}, \
gamma_f32={gamma_is_f32}, row={}, column={}: vec=0x{:04x} \
scalar=0x{:04x}",
index / n,
index % n,
fused[index],
reference[index]
);
assert_eq!(
occ[index],
reference[index],
"occupancy-raised gate/up _vec_occ diverged at M={m}, K={k}, \
N={n}, rmsnorm={rmsnorm}, decomposed={decomposed}, \
gamma_f32={gamma_is_f32}, row={}, column={}: occ=0x{:04x} \
scalar=0x{:04x}",
index / n,
index % n,
occ[index],
reference[index]
);
}
}
}
}
unsafe {
runtime.free_raw(activation_dev).unwrap();
runtime.free_raw(packed_gate_dev).unwrap();
runtime.free_raw(scales_gate_dev).unwrap();
runtime.free_raw(packed_up_dev).unwrap();
runtime.free_raw(scales_up_dev).unwrap();
}
}
}
#[test]
fn fused_skip_rmsnorm_is_bit_exact_to_three_op_path() {
run_fused_skip_rmsnorm_parity(DataType::Float16, 4, false);
}
#[test]
fn fused_skip_rmsnorm_fp32_gamma_is_bit_exact_to_three_op_path() {
run_fused_skip_rmsnorm_parity(DataType::Float32, 4, false);
}
#[test]
fn fused_skip_rmsnorm_int8_asymmetric_zp_is_bit_exact_to_three_op_path() {
run_fused_skip_rmsnorm_parity(DataType::Float32, 8, true);
}
fn run_fused_skip_rmsnorm_parity(gamma_dtype: DataType, bits: usize, explicit_zp: bool) {
let Some(runtime) = runtime() else {
eprintln!("skipping fused skip-rmsnorm parity test: CUDA runtime unavailable");
return;
};
if runtime
.require_nvrtc_half_headers("matmul_nbits_gemv_f16")
.is_err()
{
eprintln!("skipping fused skip-rmsnorm parity test: fp16 NVRTC headers unavailable");
return;
}
let (hidden, pre_k, post_n) = if bits == 8 {
(3072usize, 8192usize, 5120usize)
} else {
(896usize, QWEN_DOWN_K, 1152usize)
};
let epsilon = 1e-5f32;
let block_size = 32usize;
let blob_size = block_size * bits / 8;
let device = DeviceId::cuda(0);
for (m, following_bias) in [(1usize, false), (1, true), (5, true)] {
let mut state = 0x51ce_d00d_f00d_1234u64;
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let pack = |next: &mut dyn FnMut() -> f32, n: usize, k: usize| -> Vec<u8> {
let k_blocks = k / block_size;
let mut packed = vec![0u8; n * k_blocks * blob_size];
if bits == 8 {
for byte in packed.iter_mut() {
*byte = ((next() * 0.5 + 0.5) * 255.0).round().clamp(0.0, 255.0) as u8;
}
return packed;
}
let mut quant = vec![0u8; n * k];
for value in quant.iter_mut() {
*value = ((next() * 0.5 + 0.5) * 15.0).round().clamp(0.0, 15.0) as u8;
}
for col in 0..n {
for block in 0..k_blocks {
for pair in 0..blob_size {
let low = quant[col * k + block * block_size + pair * 2] & 15;
let high = quant[col * k + block * block_size + pair * 2 + 1] & 15;
packed[(col * k_blocks + block) * blob_size + pair] = low | (high << 4);
}
}
}
packed
};
let zp_bytes = |next: &mut dyn FnMut() -> f32, n: usize, k: usize| -> Vec<u8> {
let k_blocks = k / block_size;
(0..n * k_blocks)
.map(|_| (128.0 + (next() * 16.0)).round().clamp(96.0, 160.0) as u8)
.collect()
};
let pre_k_blocks = pre_k / block_size;
let post_k_blocks = hidden / block_size;
let activation: Vec<f16> = (0..m * pre_k).map(|_| f16::from_f32(next())).collect();
let packed_pre = pack(&mut next, hidden, pre_k);
let scales_pre: Vec<f16> = (0..hidden * pre_k_blocks)
.map(|_| f16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5)))
.collect();
let packed_post = pack(&mut next, post_n, hidden);
let scales_post: Vec<f16> = (0..post_n * post_k_blocks)
.map(|_| f16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5)))
.collect();
let residual: Vec<f16> = (0..m * hidden).map(|_| f16::from_f32(next())).collect();
let gamma_is_f32 = gamma_dtype == DataType::Float32;
let gamma_f32: Vec<f32> = (0..hidden)
.map(|_| 0.5 + 0.5 * (next() * 0.5 + 0.5))
.collect();
let gamma_bytes: Vec<u8> = if gamma_is_f32 {
gamma_f32.iter().flat_map(|v| v.to_le_bytes()).collect()
} else {
gamma_f32
.iter()
.flat_map(|v| f16::from_f32(*v).to_le_bytes())
.collect()
};
let bias_post: Vec<f16> = (0..post_n).map(|_| f16::from_f32(next())).collect();
let zp_pre: Vec<u8> = if explicit_zp {
zp_bytes(&mut next, hidden, pre_k)
} else {
Vec::new()
};
let zp_post: Vec<u8> = if explicit_zp {
zp_bytes(&mut next, post_n, hidden)
} else {
Vec::new()
};
let activation_dev = runtime.alloc_raw(activation.len() * 2).unwrap();
let packed_pre_dev = runtime.alloc_raw(packed_pre.len()).unwrap();
let scales_pre_dev = runtime.alloc_raw(scales_pre.len() * 2).unwrap();
let packed_post_dev = runtime.alloc_raw(packed_post.len()).unwrap();
let scales_post_dev = runtime.alloc_raw(scales_post.len() * 2).unwrap();
let residual_dev = runtime.alloc_raw(residual.len() * 2).unwrap();
let gamma_dev = runtime.alloc_raw(gamma_bytes.len()).unwrap();
let bias_post_dev = runtime.alloc_raw(bias_post.len() * 2).unwrap();
let matmul_out_dev = runtime.alloc_raw(m * hidden * 2).unwrap();
let normalized_dev = runtime.alloc_raw(m * hidden * 2).unwrap();
let sum_dev = runtime.alloc_raw(m * hidden * 2).unwrap();
let mean_dev = runtime.alloc_raw(m * 2).unwrap();
let invstd_dev = runtime.alloc_raw(m * 2).unwrap();
let y_ref_dev = runtime.alloc_raw(m * post_n * 2).unwrap();
let pre_fused_dev = runtime.alloc_raw(m * hidden * 2).unwrap();
let y_fused_dev = runtime.alloc_raw(m * post_n * 2).unwrap();
let zp_pre_dev = explicit_zp.then(|| runtime.alloc_raw(zp_pre.len()).unwrap());
let zp_post_dev = explicit_zp.then(|| runtime.alloc_raw(zp_post.len()).unwrap());
unsafe {
runtime.htod(as_bytes(&activation), activation_dev).unwrap();
runtime.htod(&packed_pre, packed_pre_dev).unwrap();
runtime.htod(as_bytes(&scales_pre), scales_pre_dev).unwrap();
runtime.htod(&packed_post, packed_post_dev).unwrap();
runtime
.htod(as_bytes(&scales_post), scales_post_dev)
.unwrap();
runtime.htod(as_bytes(&residual), residual_dev).unwrap();
runtime.htod(&gamma_bytes, gamma_dev).unwrap();
runtime.htod(as_bytes(&bias_post), bias_post_dev).unwrap();
if let Some(dev) = zp_pre_dev {
runtime.htod(&zp_pre, dev).unwrap();
}
if let Some(dev) = zp_post_dev {
runtime.htod(&zp_post, dev).unwrap();
}
}
let pre_a_shape = [m, pre_k];
let pre_a_strides = [pre_k as i64, 1];
let pre_b_shape = [hidden, pre_k_blocks, blob_size];
let pre_b_strides = [(pre_k_blocks * blob_size) as i64, blob_size as i64, 1];
let pre_scales_shape = [hidden, pre_k_blocks];
let pre_scales_strides = [pre_k_blocks as i64, 1];
let hidden_shape = [m, hidden];
let hidden_strides = [hidden as i64, 1];
let gamma_shape = [hidden];
let gamma_strides = [1i64];
let post_b_shape = [post_n, post_k_blocks, blob_size];
let post_b_strides = [(post_k_blocks * blob_size) as i64, blob_size as i64, 1];
let post_scales_shape = [post_n, post_k_blocks];
let post_scales_strides = [post_k_blocks as i64, 1];
let post_bias_shape = [post_n];
let post_bias_strides = [1i64];
let y_shape = [m, post_n];
let y_strides = [post_n as i64, 1];
let stat_shape = [m];
let stat_strides = [1i64];
let activation_view = TensorView::new(
device_ptr(activation_dev),
DataType::Float16,
&pre_a_shape,
&pre_a_strides,
device,
);
let packed_pre_view = TensorView::new(
device_ptr(packed_pre_dev),
DataType::Uint8,
&pre_b_shape,
&pre_b_strides,
device,
);
let scales_pre_view = TensorView::new(
device_ptr(scales_pre_dev),
DataType::Float16,
&pre_scales_shape,
&pre_scales_strides,
device,
);
let packed_post_view = TensorView::new(
device_ptr(packed_post_dev),
DataType::Uint8,
&post_b_shape,
&post_b_strides,
device,
);
let scales_post_view = TensorView::new(
device_ptr(scales_post_dev),
DataType::Float16,
&post_scales_shape,
&post_scales_strides,
device,
);
let residual_view = TensorView::new(
device_ptr(residual_dev),
DataType::Float16,
&hidden_shape,
&hidden_strides,
device,
);
let gamma_view = TensorView::new(
device_ptr(gamma_dev),
gamma_dtype,
&gamma_shape,
&gamma_strides,
device,
);
let bias_post_view = TensorView::new(
device_ptr(bias_post_dev),
DataType::Float16,
&post_bias_shape,
&post_bias_strides,
device,
);
let matmul_out_view = TensorView::new(
device_ptr(matmul_out_dev),
DataType::Float16,
&hidden_shape,
&hidden_strides,
device,
);
let normalized_input_view = TensorView::new(
device_ptr(normalized_dev),
DataType::Float16,
&hidden_shape,
&hidden_strides,
device,
);
let pre_fused_input_view = TensorView::new(
device_ptr(pre_fused_dev),
DataType::Float16,
&hidden_shape,
&hidden_strides,
device,
);
let pre_zp_shape = [hidden, pre_k_blocks];
let pre_zp_strides = [pre_k_blocks as i64, 1];
let post_zp_shape = [post_n, post_k_blocks];
let post_zp_strides = [post_k_blocks as i64, 1];
let zp_pre_view = zp_pre_dev.map(|dev| {
TensorView::new(
device_ptr(dev),
DataType::Uint8,
&pre_zp_shape,
&pre_zp_strides,
device,
)
});
let zp_post_view = zp_post_dev.map(|dev| {
TensorView::new(
device_ptr(dev),
DataType::Uint8,
&post_zp_shape,
&post_zp_strides,
device,
)
});
let make_kernel = |k: usize, n: usize, fold: bool, rmsnorm: bool| MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits,
block_size,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: fold,
gate_up_swiglu: false,
decomposed_silu: false,
rmsnorm_prologue: rmsnorm,
rmsnorm_epsilon: epsilon,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
let preceding_ref = make_kernel(pre_k, hidden, false, false);
{
let mut matmul_out = TensorMut::new(
device_ptr_mut(matmul_out_dev),
DataType::Float16,
&hidden_shape,
&hidden_strides,
device,
);
preceding_ref
.run(
&{
let mut inputs =
vec![activation_view, packed_pre_view, scales_pre_view];
if let Some(zp) = zp_pre_view {
inputs.push(zp);
}
inputs
},
std::slice::from_mut(&mut matmul_out),
None,
)
.unwrap();
}
let mut skip_node = Node::new(
onnx_runtime_ir::NodeId(0),
"SkipSimplifiedLayerNormalization",
Vec::new(),
Vec::new(),
);
skip_node
.attributes
.insert("epsilon".into(), onnx_runtime_ir::Attribute::Float(epsilon));
let skip_kernel = crate::kernels::normalization::SkipSimplifiedLayerNormFactory {
runtime: runtime.clone(),
}
.create(&skip_node, &[])
.unwrap();
{
let normalized = TensorMut::new(
device_ptr_mut(normalized_dev),
DataType::Float16,
&hidden_shape,
&hidden_strides,
device,
);
let mean = TensorMut::new(
device_ptr_mut(mean_dev),
DataType::Float16,
&stat_shape,
&stat_strides,
device,
);
let invstd = TensorMut::new(
device_ptr_mut(invstd_dev),
DataType::Float16,
&stat_shape,
&stat_strides,
device,
);
let sum = TensorMut::new(
device_ptr_mut(sum_dev),
DataType::Float16,
&hidden_shape,
&hidden_strides,
device,
);
skip_kernel
.execute(
&[matmul_out_view, residual_view, gamma_view],
&mut [normalized, mean, invstd, sum],
)
.unwrap();
}
let following_ref = make_kernel(hidden, post_n, false, false);
{
let mut y_ref = TensorMut::new(
device_ptr_mut(y_ref_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
);
let mut inputs = vec![normalized_input_view, packed_post_view, scales_post_view];
if zp_post_view.is_some() || following_bias {
inputs.push(zp_post_view.unwrap_or(TensorView::absent(DataType::Uint8)));
}
if following_bias {
inputs.push(TensorView::absent(DataType::Int32));
inputs.push(bias_post_view);
}
following_ref
.run(&inputs, std::slice::from_mut(&mut y_ref), None)
.unwrap();
}
let preceding_fused = make_kernel(pre_k, hidden, true, false);
{
let mut pre_fused = TensorMut::new(
device_ptr_mut(pre_fused_dev),
DataType::Float16,
&hidden_shape,
&hidden_strides,
device,
);
preceding_fused
.run(
&[
activation_view,
packed_pre_view,
scales_pre_view,
zp_pre_view.unwrap_or(TensorView::absent(DataType::Uint8)),
TensorView::absent(DataType::Int32),
residual_view,
],
std::slice::from_mut(&mut pre_fused),
None,
)
.unwrap();
}
let following_fused = make_kernel(hidden, post_n, false, true);
{
let mut y_fused = TensorMut::new(
device_ptr_mut(y_fused_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
);
let mut inputs = vec![
pre_fused_input_view,
packed_post_view,
scales_post_view,
zp_post_view.unwrap_or(TensorView::absent(DataType::Uint8)),
TensorView::absent(DataType::Int32),
];
if following_bias {
inputs.push(bias_post_view);
} else {
inputs.push(TensorView::absent(DataType::Float16));
}
inputs.push(gamma_view);
following_fused
.run(&inputs, std::slice::from_mut(&mut y_fused), None)
.unwrap();
}
runtime.synchronize().unwrap();
let mut sum_host = vec![f16::ZERO; m * hidden];
let mut pre_fused_host = vec![f16::ZERO; m * hidden];
let mut y_ref_host = vec![f16::ZERO; m * post_n];
let mut y_fused_host = vec![f16::ZERO; m * post_n];
unsafe {
runtime.dtoh(as_bytes_mut(&mut sum_host), sum_dev).unwrap();
runtime
.dtoh(as_bytes_mut(&mut pre_fused_host), pre_fused_dev)
.unwrap();
runtime
.dtoh(as_bytes_mut(&mut y_ref_host), y_ref_dev)
.unwrap();
runtime
.dtoh(as_bytes_mut(&mut y_fused_host), y_fused_dev)
.unwrap();
for buffer in [
activation_dev,
packed_pre_dev,
scales_pre_dev,
packed_post_dev,
scales_post_dev,
residual_dev,
gamma_dev,
bias_post_dev,
matmul_out_dev,
normalized_dev,
sum_dev,
mean_dev,
invstd_dev,
y_ref_dev,
pre_fused_dev,
y_fused_dev,
] {
runtime.free_raw(buffer).unwrap();
}
for buffer in [zp_pre_dev, zp_post_dev].into_iter().flatten() {
runtime.free_raw(buffer).unwrap();
}
}
for index in 0..m * hidden {
assert_eq!(
pre_fused_host[index].to_bits(),
sum_host[index].to_bits(),
"residual epilogue diverged from skip_rmsnorm sum at M={m}, \
following_bias={following_bias}, token={}, column={}",
index / hidden,
index % hidden
);
}
let splitk_path = bits == 8 && explicit_zp && m == 1;
let norm_block_reference =
m as u32 <= crate::kernels::normalization::SKIP_RMSNORM_BLOCK_MAX_GROUPS;
let near_equal = splitk_path || norm_block_reference;
if near_equal {
let mut max_abs = 0.0f32;
let mut worst = 0.0f32;
for index in 0..m * post_n {
let fused = y_fused_host[index].to_f32();
let reference = y_ref_host[index].to_f32();
assert!(
fused.is_finite(),
"near-equal fused int8-zp GEMV produced a non-finite output at M={m}, \
following_bias={following_bias}, column={}",
index % post_n
);
max_abs = max_abs.max(reference.abs());
worst = worst.max((fused - reference).abs());
}
let bound = (max_abs * 2e-3).max(1e-3);
assert!(
worst < bound,
"fused norm prologue diverged (beyond fp reassociation) from \
skip_rmsnorm + GEMV at \
M={m}, following_bias={following_bias}: \
max_abs_diff={worst:.3e} bound={bound:.3e}"
);
} else {
for index in 0..m * post_n {
assert_eq!(
y_fused_host[index].to_bits(),
y_ref_host[index].to_bits(),
"fused norm prologue diverged from skip_rmsnorm + GEMV at M={m}, \
following_bias={following_bias}, token={}, column={}",
index / post_n,
index % post_n
);
}
}
}
}
#[test]
fn bf16_direct_store_matches_staged_cast_bit_for_bit() {
use half::bf16;
let Some(runtime) = runtime() else {
eprintln!("skipping MatMulNBits bf16 direct-store test: CUDA runtime unavailable");
return;
};
let k = ZP_SPLITK_BANDWIDTH_MIN_K;
let n = 2048usize;
let block_size = 32usize;
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let mut state = 0x51ed_2701_c0ff_ee11u64;
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let activation_bf16: Vec<bf16> = (0..k).map(|_| bf16::from_f32(next())).collect();
let packed: Vec<u8> = (0..n * k_blocks * blob_size)
.map(|_| next().to_bits() as u8)
.collect();
let scales_bf16: Vec<bf16> = (0..n * k_blocks)
.map(|_| bf16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5)))
.collect();
let zero_points: Vec<u8> = (0..n * k_blocks.div_ceil(2))
.map(|_| ((next() * 0.5 + 0.5) * 255.0).round().clamp(0.0, 255.0) as u8)
.collect();
let device = DeviceId::cuda(0);
let a_shape = [1usize, k];
let a_strides = [k as i64, 1];
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let zp_shape = [n, k_blocks.div_ceil(2)];
let zp_strides = [k_blocks.div_ceil(2) as i64, 1];
let y_shape = [1usize, n];
let y_strides = [n as i64, 1];
let act_dev = runtime.alloc_raw(activation_bf16.len() * 2).unwrap();
let packed_dev = runtime.alloc_raw(packed.len()).unwrap();
let scales_dev = runtime.alloc_raw(scales_bf16.len() * 2).unwrap();
let zp_dev = runtime.alloc_raw(zero_points.len()).unwrap();
let out_dev = runtime.alloc_raw(n * 2).unwrap();
let act_f16_dev = runtime.alloc_raw(k * 2).unwrap();
let scales_f16_dev = runtime.alloc_raw(scales_bf16.len() * 2).unwrap();
let ref_f16_dev = runtime.alloc_raw(n * 2).unwrap();
let ref_bf16_dev = runtime.alloc_raw(n * 2).unwrap();
unsafe {
runtime.htod(as_bytes(&activation_bf16), act_dev).unwrap();
runtime.htod(&packed, packed_dev).unwrap();
runtime.htod(as_bytes(&scales_bf16), scales_dev).unwrap();
runtime.htod(&zero_points, zp_dev).unwrap();
}
let kernel = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 0,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: false,
decomposed_silu: false,
rmsnorm_prologue: false,
rmsnorm_epsilon: 1e-5,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
let bf16_inputs = vec![
TensorView::new(
device_ptr(act_dev),
DataType::BFloat16,
&a_shape,
&a_strides,
device,
),
TensorView::new(
device_ptr(packed_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
),
TensorView::new(
device_ptr(scales_dev),
DataType::BFloat16,
&scales_shape,
&scales_strides,
device,
),
TensorView::new(
device_ptr(zp_dev),
DataType::Uint8,
&zp_shape,
&zp_strides,
device,
),
];
let before = BF16_DIRECT_OUT_STORES.with(|count| count.get());
let mut got_out = [TensorMut::new(
device_ptr_mut(out_dev),
DataType::BFloat16,
&y_shape,
&y_strides,
device,
)];
kernel.run(&bf16_inputs, &mut got_out, None).unwrap();
runtime.synchronize().unwrap();
let direct_stores = BF16_DIRECT_OUT_STORES.with(|count| count.get()) - before;
super::super::cast::launch_cast_raw(
&runtime,
cuptr(act_dev as *const c_void),
DataType::BFloat16,
act_f16_dev,
DataType::Float16,
k,
)
.unwrap();
super::super::cast::launch_cast_raw(
&runtime,
cuptr(scales_dev as *const c_void),
DataType::BFloat16,
scales_f16_dev,
DataType::Float16,
scales_bf16.len(),
)
.unwrap();
let ref_inputs = vec![
TensorView::new(
device_ptr(act_f16_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
),
TensorView::new(
device_ptr(packed_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
),
TensorView::new(
device_ptr(scales_f16_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
),
TensorView::new(
device_ptr(zp_dev),
DataType::Uint8,
&zp_shape,
&zp_strides,
device,
),
];
let mut ref_f16 = [TensorMut::new(
device_ptr_mut(ref_f16_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
)];
kernel.run(&ref_inputs, &mut ref_f16, None).unwrap();
super::super::cast::launch_cast_raw(
&runtime,
ref_f16_dev,
DataType::Float16,
ref_bf16_dev,
DataType::BFloat16,
n,
)
.unwrap();
runtime.synchronize().unwrap();
let mut got = vec![bf16::ZERO; n];
let mut want = vec![bf16::ZERO; n];
unsafe {
runtime.dtoh(as_bytes_mut(&mut got), out_dev).unwrap();
runtime.dtoh(as_bytes_mut(&mut want), ref_bf16_dev).unwrap();
for buffer in [
act_dev,
packed_dev,
scales_dev,
zp_dev,
out_dev,
act_f16_dev,
scales_f16_dev,
ref_f16_dev,
ref_bf16_dev,
] {
runtime.free_raw(buffer).unwrap();
}
}
assert_eq!(
direct_stores, 1,
"the bf16 run did not take the direct-store path, so this test would \
compare the staged route against itself"
);
for index in 0..n {
assert_eq!(
got[index].to_bits(),
want[index].to_bits(),
"direct bf16 store diverged from the staged cast at column {index}"
);
}
}
#[test]
fn bf16_direct_store_declines_the_small_batch_row_loop() {
use half::bf16;
let Some(runtime) = runtime() else {
eprintln!(
"skipping MatMulNBits bf16 small-batch direct-store test: CUDA runtime unavailable"
);
return;
};
if decode_gemv_loop_max_m() < 2 {
eprintln!("skipping MatMulNBits bf16 small-batch direct-store test: row loop disabled");
return;
}
let m = 2usize;
let k = 2048usize;
let n = 2048usize;
let block_size = 32usize;
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let mut state = 0x7c0f_fee5_1234_9876u64;
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let activation_bf16: Vec<bf16> = (0..m * k).map(|_| bf16::from_f32(next())).collect();
let packed: Vec<u8> = (0..n * k_blocks * blob_size)
.map(|_| next().to_bits() as u8)
.collect();
let scales_bf16: Vec<bf16> = (0..n * k_blocks)
.map(|_| bf16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5)))
.collect();
let zero_points: Vec<u8> = (0..n * k_blocks.div_ceil(2))
.map(|_| ((next() * 0.5 + 0.5) * 255.0).round().clamp(0.0, 255.0) as u8)
.collect();
let device = DeviceId::cuda(0);
let a_shape = [m, k];
let a_strides = [k as i64, 1];
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let zp_shape = [n, k_blocks.div_ceil(2)];
let zp_strides = [k_blocks.div_ceil(2) as i64, 1];
let y_shape = [m, n];
let y_strides = [n as i64, 1];
let act_dev = runtime.alloc_raw(activation_bf16.len() * 2).unwrap();
let packed_dev = runtime.alloc_raw(packed.len()).unwrap();
let scales_dev = runtime.alloc_raw(scales_bf16.len() * 2).unwrap();
let zp_dev = runtime.alloc_raw(zero_points.len()).unwrap();
let out_dev = runtime.alloc_raw(m * n * 2).unwrap();
let act_f16_dev = runtime.alloc_raw(m * k * 2).unwrap();
let scales_f16_dev = runtime.alloc_raw(scales_bf16.len() * 2).unwrap();
let ref_f16_dev = runtime.alloc_raw(m * n * 2).unwrap();
let ref_bf16_dev = runtime.alloc_raw(m * n * 2).unwrap();
unsafe {
runtime.htod(as_bytes(&activation_bf16), act_dev).unwrap();
runtime.htod(&packed, packed_dev).unwrap();
runtime.htod(as_bytes(&scales_bf16), scales_dev).unwrap();
runtime.htod(&zero_points, zp_dev).unwrap();
}
let kernel = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 0,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: false,
decomposed_silu: false,
rmsnorm_prologue: false,
rmsnorm_epsilon: 1e-5,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
let bf16_inputs = vec![
TensorView::new(
device_ptr(act_dev),
DataType::BFloat16,
&a_shape,
&a_strides,
device,
),
TensorView::new(
device_ptr(packed_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
),
TensorView::new(
device_ptr(scales_dev),
DataType::BFloat16,
&scales_shape,
&scales_strides,
device,
),
TensorView::new(
device_ptr(zp_dev),
DataType::Uint8,
&zp_shape,
&zp_strides,
device,
),
];
let before = BF16_DIRECT_OUT_STORES.with(|count| count.get());
let mut got_out = [TensorMut::new(
device_ptr_mut(out_dev),
DataType::BFloat16,
&y_shape,
&y_strides,
device,
)];
kernel.run(&bf16_inputs, &mut got_out, None).unwrap();
runtime.synchronize().unwrap();
let direct_stores = BF16_DIRECT_OUT_STORES.with(|count| count.get()) - before;
super::super::cast::launch_cast_raw(
&runtime,
cuptr(act_dev as *const c_void),
DataType::BFloat16,
act_f16_dev,
DataType::Float16,
m * k,
)
.unwrap();
super::super::cast::launch_cast_raw(
&runtime,
cuptr(scales_dev as *const c_void),
DataType::BFloat16,
scales_f16_dev,
DataType::Float16,
scales_bf16.len(),
)
.unwrap();
let ref_inputs = vec![
TensorView::new(
device_ptr(act_f16_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
),
TensorView::new(
device_ptr(packed_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
),
TensorView::new(
device_ptr(scales_f16_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
),
TensorView::new(
device_ptr(zp_dev),
DataType::Uint8,
&zp_shape,
&zp_strides,
device,
),
];
let mut ref_f16 = [TensorMut::new(
device_ptr_mut(ref_f16_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
)];
kernel.run(&ref_inputs, &mut ref_f16, None).unwrap();
super::super::cast::launch_cast_raw(
&runtime,
ref_f16_dev,
DataType::Float16,
ref_bf16_dev,
DataType::BFloat16,
m * n,
)
.unwrap();
runtime.synchronize().unwrap();
let mut got = vec![bf16::ZERO; m * n];
let mut want = vec![bf16::ZERO; m * n];
unsafe {
runtime.dtoh(as_bytes_mut(&mut got), out_dev).unwrap();
runtime.dtoh(as_bytes_mut(&mut want), ref_bf16_dev).unwrap();
for buffer in [
act_dev,
packed_dev,
scales_dev,
zp_dev,
out_dev,
act_f16_dev,
scales_f16_dev,
ref_f16_dev,
ref_bf16_dev,
] {
runtime.free_raw(buffer).unwrap();
}
}
assert_eq!(
direct_stores, 0,
"a per-row launch accepted the whole-output direct-store offer; rows \
past the first would be left stale"
);
for index in 0..m * n {
assert_eq!(
got[index].to_bits(),
want[index].to_bits(),
"small-batch bf16 output diverged at row {}, column {}",
index / n,
index % n
);
}
}
#[test]
fn bf16_scale_cache_is_bit_exact_to_inline_staging() {
use half::bf16;
let Some(runtime) = runtime() else {
eprintln!("skipping MatMulNBits bf16 scale-cache test: CUDA runtime unavailable");
return;
};
let k = 256usize;
let n = 64usize;
let block_size = 32usize;
let k_blocks = k / block_size;
let blob_size = block_size / 2;
let mut state = 0x0bad_c0de_dead_beefu64;
let mut next = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
};
let activation_bf16: Vec<bf16> = (0..k).map(|_| bf16::from_f32(next())).collect();
let packed: Vec<u8> = (0..n * k_blocks * blob_size)
.map(|_| next().to_bits() as u8)
.collect();
let scales_bf16: Vec<bf16> = (0..n * k_blocks)
.map(|_| bf16::from_f32(0.015 + 0.01 * (next() * 0.5 + 0.5)))
.collect();
let zero_points: Vec<u8> = (0..n * k_blocks.div_ceil(2))
.map(|_| ((next() * 0.5 + 0.5) * 255.0).round().clamp(0.0, 255.0) as u8)
.collect();
let device = DeviceId::cuda(0);
let a_shape = [1usize, k];
let a_strides = [k as i64, 1];
let b_shape = [n, k_blocks, blob_size];
let b_strides = [(k_blocks * blob_size) as i64, blob_size as i64, 1];
let scales_shape = [n, k_blocks];
let scales_strides = [k_blocks as i64, 1];
let zp_shape = [n, k_blocks.div_ceil(2)];
let zp_strides = [k_blocks.div_ceil(2) as i64, 1];
let y_shape = [1usize, n];
let y_strides = [n as i64, 1];
let act_dev = runtime.alloc_raw(activation_bf16.len() * 2).unwrap();
let packed_dev = runtime.alloc_raw(packed.len()).unwrap();
let scales_dev = runtime.alloc_raw(scales_bf16.len() * 2).unwrap();
let zp_dev = runtime.alloc_raw(zero_points.len()).unwrap();
let out1_dev = runtime.alloc_raw(n * 2).unwrap();
let out2_dev = runtime.alloc_raw(n * 2).unwrap();
unsafe {
runtime.htod(as_bytes(&activation_bf16), act_dev).unwrap();
runtime.htod(&packed, packed_dev).unwrap();
runtime.htod(as_bytes(&scales_bf16), scales_dev).unwrap();
runtime.htod(&zero_points, zp_dev).unwrap();
}
let inputs = vec![
TensorView::new(
device_ptr(act_dev),
DataType::BFloat16,
&a_shape,
&a_strides,
device,
),
TensorView::new(
device_ptr(packed_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
),
TensorView::new(
device_ptr(scales_dev),
DataType::BFloat16,
&scales_shape,
&scales_strides,
device,
),
TensorView::new(
device_ptr(zp_dev),
DataType::Uint8,
&zp_shape,
&zp_strides,
device,
),
];
let kernel = MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size,
accuracy_level: 0,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: false,
decomposed_silu: false,
rmsnorm_prologue: false,
rmsnorm_epsilon: 1e-5,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
};
let mut out1 = [TensorMut::new(
device_ptr_mut(out1_dev),
DataType::BFloat16,
&y_shape,
&y_strides,
device,
)];
kernel.run(&inputs, &mut out1, None).unwrap();
let mut out2 = [TensorMut::new(
device_ptr_mut(out2_dev),
DataType::BFloat16,
&y_shape,
&y_strides,
device,
)];
kernel.run(&inputs, &mut out2, None).unwrap();
runtime.synchronize().unwrap();
let act_f16_dev = runtime.alloc_raw(k * 2).unwrap();
let scales_f16_dev = runtime.alloc_raw(scales_bf16.len() * 2).unwrap();
let ref_f16_dev = runtime.alloc_raw(n * 2).unwrap();
let ref_bf16_dev = runtime.alloc_raw(n * 2).unwrap();
super::super::cast::launch_cast_raw(
&runtime,
cuptr(act_dev as *const c_void),
DataType::BFloat16,
act_f16_dev,
DataType::Float16,
k,
)
.unwrap();
super::super::cast::launch_cast_raw(
&runtime,
cuptr(scales_dev as *const c_void),
DataType::BFloat16,
scales_f16_dev,
DataType::Float16,
scales_bf16.len(),
)
.unwrap();
let ref_inputs = vec![
TensorView::new(
device_ptr(act_f16_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
),
TensorView::new(
device_ptr(packed_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
),
TensorView::new(
device_ptr(scales_f16_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
),
TensorView::new(
device_ptr(zp_dev),
DataType::Uint8,
&zp_shape,
&zp_strides,
device,
),
];
let mut ref_f16 = [TensorMut::new(
device_ptr_mut(ref_f16_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
)];
kernel.run(&ref_inputs, &mut ref_f16, None).unwrap();
super::super::cast::launch_cast_raw(
&runtime,
ref_f16_dev,
DataType::Float16,
ref_bf16_dev,
DataType::BFloat16,
n,
)
.unwrap();
runtime.synchronize().unwrap();
let mut got1 = vec![bf16::ZERO; n];
let mut got2 = vec![bf16::ZERO; n];
let mut want = vec![bf16::ZERO; n];
unsafe {
runtime.dtoh(as_bytes_mut(&mut got1), out1_dev).unwrap();
runtime.dtoh(as_bytes_mut(&mut got2), out2_dev).unwrap();
runtime.dtoh(as_bytes_mut(&mut want), ref_bf16_dev).unwrap();
for buffer in [
act_dev,
packed_dev,
scales_dev,
zp_dev,
out1_dev,
out2_dev,
act_f16_dev,
scales_f16_dev,
ref_f16_dev,
ref_bf16_dev,
] {
runtime.free_raw(buffer).unwrap();
}
}
for index in 0..n {
assert_eq!(
got1[index].to_bits(),
got2[index].to_bits(),
"cached bf16 scale path is non-deterministic across steps at column {index}"
);
assert_eq!(
got1[index].to_bits(),
want[index].to_bits(),
"cached bf16 scale path diverged from inline staging at column {index}"
);
}
}
#[test]
#[ignore = "perf microbench; requires a dedicated idle SM80+ CUDA device"]
fn decode_gemv_achieved_bandwidth_by_projection_shape() {
use cudarc::driver::result::event;
use cudarc::driver::sys;
use cudarc::driver::sys::CUdeviceptr;
const A100_SXM4_80GB_PEAK_GBPS: f64 = 2039.0;
const SHAPES: [(usize, usize, usize, &str); 6] = [
(6656, 19968, 104, "gate/up"),
(19968, 6656, 52, "down"),
(6656, 4096, 104, "q"),
(4096, 6656, 52, "o"),
(6656, 202048, 1, "lm_head"),
(6656, 256, 104, "k/v (GQA)"),
];
const BLOCK: usize = 32;
let Some(runtime) = runtime() else {
eprintln!("skipping GEMV bandwidth probe: CUDA runtime unavailable");
return;
};
if runtime
.require_nvrtc_half_headers("gemv bandwidth probe")
.is_err()
{
eprintln!("skipping GEMV bandwidth probe: fp16 NVRTC headers unavailable");
return;
}
runtime.bind().unwrap();
let reps = zc_env_usize("GEMV_PROBE_REPS", 15).max(5);
struct ShapeBench {
kernel: MatMulNBitsKernel,
label: &'static str,
k: usize,
n: usize,
calls: usize,
k_blocks: usize,
blob_size: usize,
zp_row_bytes: usize,
activation_dev: CUdeviceptr,
packed_dev: CUdeviceptr,
scales_dev: CUdeviceptr,
zp_dev: CUdeviceptr,
output_dev: CUdeviceptr,
bytes: f64,
}
impl ShapeBench {
fn launch(&self) {
let a_shape = [1usize, self.k];
let a_strides = [self.k as i64, 1];
let b_shape = [self.n, self.k_blocks, self.blob_size];
let b_strides = [
(self.k_blocks * self.blob_size) as i64,
self.blob_size as i64,
1,
];
let scales_shape = [self.n, self.k_blocks];
let scales_strides = [self.k_blocks as i64, 1];
let zp_shape = [self.n, self.zp_row_bytes];
let zp_strides = [self.zp_row_bytes as i64, 1];
let y_shape = [1usize, self.n];
let y_strides = [self.n as i64, 1];
let device = DeviceId::cuda(0);
let inputs = vec![
TensorView::new(
device_ptr(self.activation_dev),
DataType::Float16,
&a_shape,
&a_strides,
device,
),
TensorView::new(
device_ptr(self.packed_dev),
DataType::Uint8,
&b_shape,
&b_strides,
device,
),
TensorView::new(
device_ptr(self.scales_dev),
DataType::Float16,
&scales_shape,
&scales_strides,
device,
),
TensorView::new(
device_ptr(self.zp_dev),
DataType::Uint8,
&zp_shape,
&zp_strides,
device,
),
];
let mut outputs = [TensorMut::new(
device_ptr_mut(self.output_dev),
DataType::Float16,
&y_shape,
&y_strides,
device,
)];
self.kernel.run(&inputs, &mut outputs, None).unwrap();
}
fn median_ms(&self, reps: usize) -> (f64, f64) {
const BATCH: usize = 32;
let runtime = &self.kernel.runtime;
let mut host = Vec::with_capacity(reps);
for _ in 0..reps {
let enqueue_begin = std::time::Instant::now();
for _ in 0..BATCH {
self.launch();
}
host.push(enqueue_begin.elapsed().as_secs_f64() * 1e3 / BATCH as f64);
runtime.synchronize().unwrap();
}
host.sort_by(|a, b| a.partial_cmp(b).unwrap());
let host_ms = host[host.len() / 2];
let captured = runtime.reset_graph().is_ok()
&& runtime.begin_graph_capture(&[&self.kernel]).is_ok()
&& {
for _ in 0..BATCH {
self.launch();
}
runtime.end_graph_capture().is_ok()
};
if !captured {
runtime.abort_graph_capture().ok();
}
let mut gpu = Vec::with_capacity(reps);
for _ in 0..reps {
let start = event::create(sys::CUevent_flags::CU_EVENT_DEFAULT).unwrap();
let end = event::create(sys::CUevent_flags::CU_EVENT_DEFAULT).unwrap();
unsafe {
event::record(start, runtime.stream_ptr()).unwrap();
if captured {
runtime.replay_graph().unwrap();
} else {
for _ in 0..BATCH {
self.launch();
}
}
event::record(end, runtime.stream_ptr()).unwrap();
event::synchronize(end).unwrap();
gpu.push(event::elapsed(start, end).unwrap() as f64 / BATCH as f64);
event::destroy(start).ok();
event::destroy(end).ok();
}
}
runtime.reset_graph().ok();
gpu.sort_by(|a, b| a.partial_cmp(b).unwrap());
(gpu[gpu.len() / 2], host_ms)
}
}
let mut benches = Vec::new();
for (k, n, calls, label) in SHAPES {
let k_blocks = k / BLOCK;
let blob_size = BLOCK / 2;
let zp_row_bytes = k_blocks.div_ceil(2);
let mut state = 0x9e37_79b9_7f4a_7c15u64;
let mut next_byte = || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
(state >> 56) as u8
};
let packed: Vec<u8> = (0..n * k_blocks * blob_size).map(|_| next_byte()).collect();
let zp: Vec<u8> = (0..n * zp_row_bytes).map(|_| next_byte()).collect();
let scales: Vec<f16> = (0..n * k_blocks).map(|_| f16::from_f32(0.02)).collect();
let activation: Vec<f16> = (0..k).map(|_| f16::from_f32(0.01)).collect();
let activation_dev = runtime.alloc_raw(activation.len() * 2).unwrap();
let packed_dev = runtime.alloc_raw(packed.len()).unwrap();
let scales_dev = runtime.alloc_raw(scales.len() * 2).unwrap();
let zp_dev = runtime.alloc_raw(zp.len()).unwrap();
let output_dev = runtime.alloc_raw(n * 2).unwrap();
unsafe {
runtime.htod(as_bytes(&activation), activation_dev).unwrap();
runtime.htod(&packed, packed_dev).unwrap();
runtime.htod(as_bytes(&scales), scales_dev).unwrap();
runtime.htod(&zp, zp_dev).unwrap();
}
let bench = ShapeBench {
kernel: MatMulNBitsKernel {
runtime: runtime.clone(),
k,
n,
bits: 4,
block_size: BLOCK,
accuracy_level: 4,
accuracy4_workspace: None,
marlin_repack_cache: marlin_gemm::RepackCache::new(runtime.clone()),
constant_inputs: [true; 8],
fold_bias_post_round: false,
gate_up_swiglu: false,
decomposed_silu: false,
rmsnorm_prologue: false,
rmsnorm_epsilon: 1e-5,
last_call_capture_safe: AtomicBool::new(false),
bf16_scratch: Mutex::new(Bf16Scratch::new(runtime.clone())),
bf16_const_cache: Mutex::new(Bf16ConstCache::new(runtime.clone())),
},
label,
k,
n,
calls,
k_blocks,
blob_size,
zp_row_bytes,
activation_dev,
packed_dev,
scales_dev,
zp_dev,
output_dev,
bytes: (n * k_blocks * blob_size + n * k_blocks * 2 + n * zp_row_bytes) as f64,
};
bench.launch();
runtime.synchronize().unwrap();
benches.push(bench);
}
let ramp = &benches[0];
let ramp_start = std::time::Instant::now();
let ramp_deadline = ramp_start + std::time::Duration::from_secs(30);
let mut previous = f64::INFINITY;
let mut ramp_trace = Vec::new();
loop {
let (now, _) = ramp.median_ms(5);
ramp_trace.push(now * 1000.0);
let settled = now > previous * 0.985;
previous = now;
if settled && ramp_start.elapsed() >= std::time::Duration::from_secs(8) {
break;
}
if std::time::Instant::now() >= ramp_deadline {
eprintln!(
"WARNING: clock had not settled after 30 s of ramping; \
absolute numbers below are not comparable across runs"
);
break;
}
}
let ramp_best = ramp_trace.iter().cloned().fold(f64::INFINITY, f64::min);
println!(
"clock ramp on {}: {:.0} -> {:.0} us over {} readings in {:.1} s \
(best {:.0} us, {:.0}% off the first)",
ramp.label,
ramp_trace[0],
ramp_trace[ramp_trace.len() - 1],
ramp_trace.len(),
ramp_start.elapsed().as_secs_f64(),
ramp_best,
100.0 * (ramp_trace[0] - ramp_best) / ramp_trace[0]
);
let mut total_ms_per_token = 0.0f64;
let mut total_bytes_per_token = 0.0f64;
println!(
"{:<10} {:>7} {:>7} {:>6} {:>10} {:>9} {:>7} {:>9} {:>7} {:>9}",
"shape",
"K",
"N",
"calls",
"median_us",
"MB/call",
"GB/s",
"%peak",
"ms/tok",
"host_us"
);
let (first_before, _) = benches[0].median_ms(reps);
let mut total_host_ms_per_token = 0.0f64;
for bench in &benches {
let (median_ms, host_ms) = bench.median_ms(reps);
total_host_ms_per_token += host_ms * bench.calls as f64;
let gbps = bench.bytes / (median_ms * 1e-3) / 1e9;
let ms_tok = median_ms * bench.calls as f64;
total_ms_per_token += ms_tok;
total_bytes_per_token += bench.bytes * bench.calls as f64;
println!(
"{:<10} {:>7} {:>7} {:>6} {:>10.1} {:>9.2} {:>7.0} {:>8.1}% {:>7.2} {:>9.1}",
bench.label,
bench.k,
bench.n,
bench.calls,
median_ms * 1000.0,
bench.bytes / 1e6,
gbps,
100.0 * gbps / A100_SXM4_80GB_PEAK_GBPS,
ms_tok,
host_ms * 1000.0
);
}
let (first_after, _) = benches[0].median_ms(reps);
let drift = (first_after - first_before).abs() / first_before;
if drift > 0.03 {
eprintln!(
"WARNING: {} drifted {:.1}% across the sweep ({:.1} -> {:.1} us); \
the device was not stable and these rows are not comparable",
benches[0].label,
100.0 * drift,
first_before * 1000.0,
first_after * 1000.0
);
} else {
println!("\ndevice stable across sweep: {:.1}% drift", 100.0 * drift);
}
for bench in &benches {
unsafe {
runtime.free_raw(bench.activation_dev).ok();
runtime.free_raw(bench.packed_dev).ok();
runtime.free_raw(bench.scales_dev).ok();
runtime.free_raw(bench.zp_dev).ok();
runtime.free_raw(bench.output_dev).ok();
}
}
let aggregate_gbps = total_bytes_per_token / (total_ms_per_token * 1e-3) / 1e9;
let implied_tps = 1000.0 / total_ms_per_token;
let roofline_tps = A100_SXM4_80GB_PEAK_GBPS * 1e9 / total_bytes_per_token;
println!(
"\nweight traffic {:.2} GB/token, GEMV time {:.2} ms/token\n\
aggregate achieved bandwidth {:.0} GB/s ({:.1}% of peak)\n\
decode ceiling from GEMV alone: {:.1} tok/s (bandwidth roofline {:.1} tok/s)\n\
host dispatch for the same launches: {:.2} ms/token \
(uncaptured ceiling {:.1} tok/s)",
total_bytes_per_token / 1e9,
total_ms_per_token,
aggregate_gbps,
100.0 * aggregate_gbps / A100_SXM4_80GB_PEAK_GBPS,
implied_tps,
roofline_tps,
total_host_ms_per_token,
1000.0 / total_host_ms_per_token
);
assert!(
implied_tps.is_finite() && implied_tps > 0.0,
"GEMV bandwidth probe produced no usable timing"
);
}
}