From 7ebdd7c552fa00fd1ba65ea6a05e42005820058d Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Fri, 29 May 2026 21:38:52 +0800 Subject: [PATCH] kernels: MXFP4 -> BF16 dequant kernel (verified vs numpy) From-scratch CUDA kernel for gpt-oss expert weights: one thread per packed byte decodes 2 FP4 (E2M1) codes, applies the per-32-block E8M0 scale (2^(e-127)), and writes BF16 transposed into [IN, OUT] (IN = nblk*32) so it drops straight into x @ W. dequant_mxfp4() wrapper takes raw GpuBuffers (uint8 is not an xserv Tensor dtype). mxfp4-check bin dequants layer-0 expert-0 on GPU and matches tools/mxfp4_probe.py exactly: [0, 0, 0, -0.0625, 0, -0, -0.015625, -0.03125] This lets experts stay MXFP4-resident on GPU (13GB, fits one 32GB card) and be dequantized to a BF16 scratch right before each expert GEMM, instead of holding 36GB of BF16 or uploading experts per token. Loader plumbing + GPU MoE decode use it next. Co-Authored-By: Claude Opus 4.8 --- crates/xserv-kernels/src/lib.rs | 1 + crates/xserv-kernels/src/quant.rs | 38 ++++++++++++ crates/xserv-model/src/bin/mxfp4-check.rs | 59 ++++++++++++++++++ csrc/quant/mxfp4.cu | 73 +++++++++++++++++++++++ 4 files changed, 171 insertions(+) create mode 100644 crates/xserv-kernels/src/quant.rs create mode 100644 crates/xserv-model/src/bin/mxfp4-check.rs create mode 100644 csrc/quant/mxfp4.cu diff --git a/crates/xserv-kernels/src/lib.rs b/crates/xserv-kernels/src/lib.rs index a5c30d6..1af95b0 100644 --- a/crates/xserv-kernels/src/lib.rs +++ b/crates/xserv-kernels/src/lib.rs @@ -15,6 +15,7 @@ pub use attention::{attention, decode_attention, flash_attention, paged_decode_a pub use embedding::embedding; pub use gemm::{batched_matmul, matmul, GemmBackend}; pub use layernorm::layernorm; +pub use quant::dequant_mxfp4; pub use rmsnorm::{add_rmsnorm, rmsnorm}; pub use rope::{rope_inplace, RopeCache}; pub use softmax::softmax; diff --git a/crates/xserv-kernels/src/quant.rs b/crates/xserv-kernels/src/quant.rs new file mode 100644 index 0000000..8076f8f --- /dev/null +++ b/crates/xserv-kernels/src/quant.rs @@ -0,0 +1,38 @@ +//! MXFP4 dequantization (gpt-oss expert weights) -> BF16. +//! +//! The packed weights are uint8 (blocks + E8M0 scales), which is not an xserv +//! `Tensor` dtype, so the input is raw `GpuBuffer`s; the output is a BF16 +//! `Tensor` in `[IN, OUT]` layout (IN = nblk*32), ready for `x @ W`. + +use std::ffi::c_void; +use xserv_cuda::GpuBuffer; +use xserv_tensor::{DType, Device, Tensor}; + +unsafe extern "C" { + fn launch_mxfp4_dequant_bf16( + blocks: *const c_void, scales: *const c_void, out: *mut c_void, + out_dim: i32, nblk: i32, stream: *mut c_void, + ); +} + +/// Dequantize one expert's MXFP4 weight to a BF16 `[IN, OUT]` tensor +/// (IN = `nblk*32`), the transpose of the natural `[OUT, IN]`, so it drops into +/// `y = x[1,IN] @ W[IN,OUT]` with no extra transpose. +/// +/// `blocks`: uint8 device buffer of `out_dim * nblk * 16` bytes. +/// `scales`: uint8 device buffer of `out_dim * nblk` bytes (E8M0 exponents). +pub fn dequant_mxfp4( + blocks: &GpuBuffer, scales: &GpuBuffer, out_dim: usize, nblk: usize, device: u32, +) -> Tensor { + let in_dim = nblk * 32; + let out = Tensor::empty(&[in_dim, out_dim], DType::BF16, Device::Cuda(device)); + unsafe { + launch_mxfp4_dequant_bf16( + blocks.as_ptr() as *const c_void, + scales.as_ptr() as *const c_void, + out.data_ptr() as *mut c_void, + out_dim as i32, nblk as i32, std::ptr::null_mut(), + ); + } + out +} diff --git a/crates/xserv-model/src/bin/mxfp4-check.rs b/crates/xserv-model/src/bin/mxfp4-check.rs new file mode 100644 index 0000000..684deb9 --- /dev/null +++ b/crates/xserv-model/src/bin/mxfp4-check.rs @@ -0,0 +1,59 @@ +//! Verify the GPU MXFP4 dequant kernel against the numpy reference. +//! Loads layer-0 expert-0 gate_up_proj from the raw MXFP4 model, dequantizes on +//! GPU, and prints out[in_idx, out_row=0] for in_idx 0..8 — which should match +//! tools/mxfp4_probe.py's "row0 first 8 vals". +//! +//! Usage: mxfp4-check +use std::fs; +use std::path::PathBuf; +use safetensors::SafeTensors; +use xserv_cuda::GpuBuffer; +use xserv_tensor::Device; + +fn main() { + let dir = PathBuf::from(std::env::args().nth(1).expect("model dir")); + xserv_cuda::device::set_device(0).unwrap(); + + let blocks_name = "model.layers.0.mlp.experts.gate_up_proj_blocks"; + let scales_name = "model.layers.0.mlp.experts.gate_up_proj_scales"; + + // Find the shard holding these tensors. + let index: serde_json::Value = + serde_json::from_str(&fs::read_to_string(dir.join("model.safetensors.index.json")).unwrap()).unwrap(); + let wm = &index["weight_map"]; + let shard = wm[blocks_name].as_str().expect("blocks in index"); + assert_eq!(shard, wm[scales_name].as_str().unwrap(), "blocks/scales in same shard assumed"); + eprintln!("[mxfp4-check] reading shard {shard}"); + let data = fs::read(dir.join(shard)).unwrap(); + let st = SafeTensors::deserialize(&data).unwrap(); + + let bv = st.tensor(blocks_name).unwrap(); + let sv = st.tensor(scales_name).unwrap(); + eprintln!("[mxfp4-check] blocks shape {:?} dtype {:?}", bv.shape(), bv.dtype()); + eprintln!("[mxfp4-check] scales shape {:?} dtype {:?}", sv.shape(), sv.dtype()); + + // shapes: blocks [E, OUT, nblk, 16], scales [E, OUT, nblk] + let (out_dim, nblk) = (bv.shape()[1], bv.shape()[2]); + let bytes_per_expert_blocks = out_dim * nblk * 16; + let bytes_per_expert_scales = out_dim * nblk; + + // Expert 0 = first slice of the contiguous [E, ...] buffer. + let blk0 = &bv.data()[..bytes_per_expert_blocks]; + let scl0 = &sv.data()[..bytes_per_expert_scales]; + + let mut blk_buf = GpuBuffer::alloc(blk0.len()).unwrap(); + let mut scl_buf = GpuBuffer::alloc(scl0.len()).unwrap(); + blk_buf.copy_from_host(blk0).unwrap(); + scl_buf.copy_from_host(scl0).unwrap(); + + let out = xserv_kernels::dequant_mxfp4(&blk_buf, &scl_buf, out_dim, nblk, 0); // BF16 [IN, OUT] + let in_dim = nblk * 32; + assert_eq!(out.shape(), &[in_dim, out_dim]); + + let host = out.to_device(Device::Cpu); + let s = host.as_slice::(); + // out[in_idx, 0] = s[in_idx*out_dim + 0] + let vals: Vec = (0..8).map(|i| s[i * out_dim].to_f32()).collect(); + println!("GPU out[0..8, row0] = {vals:?}"); + println!("numpy ref row0 first8 = [0.0, 0.0, 0.0, -0.0625, 0.0, -0.0, -0.0156, -0.0312]"); +} diff --git a/csrc/quant/mxfp4.cu b/csrc/quant/mxfp4.cu new file mode 100644 index 0000000..cf5519a --- /dev/null +++ b/csrc/quant/mxfp4.cu @@ -0,0 +1,73 @@ +// MXFP4 dequantization (gpt-oss expert weights) -> BF16. +// +// gpt-oss stores each expert MLP weight in OCP Microscaling FP4: +// blocks: uint8 [OUT, nblk, 16] — each 16-byte row packs 32 FP4 codes +// (low nibble = even elem, high = odd) +// scales: uint8 [OUT, nblk] — one E8M0 (8-bit exponent) per 32-elem block +// value = fp4_e2m1[code] * 2^(scale - 127) +// The contraction (input) dim is IN = nblk * 32. +// +// We emit BF16 in [IN, OUT] (row-major) layout — i.e. the transpose of the +// natural [OUT, IN] — so it drops straight into y = x[1,IN] @ W[IN,OUT] without +// a separate transpose. This matches the offline BF16 path (gptoss_dequant.py), +// keeping the two routes numerically identical. + +#include +#include + +// FP4 E2M1 code -> value (OCP MX). 16 codes: sign, 2-bit exp, 1-bit mantissa. +__constant__ float kFp4[16] = { + 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, + -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f, +}; + +// One thread per packed byte: OUT * nblk * 16 threads. Each decodes 2 elements. +__global__ void mxfp4_dequant_kernel( + const uint8_t* __restrict__ blocks, // [OUT, nblk, 16] + const uint8_t* __restrict__ scales, // [OUT, nblk] + __nv_bfloat16* __restrict__ out, // [IN, OUT], IN = nblk*32 + int out_dim, int nblk +) { + long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x; + long long total = (long long)out_dim * nblk * 16; + if (idx >= total) return; + + int b = idx & 15; // byte within block (0..15) + long long t = idx >> 4; // (out_row, block) + int block = t % nblk; + int out_row = t / nblk; + + uint8_t byte = blocks[idx]; + int code_lo = byte & 0x0F; + int code_hi = (byte >> 4) & 0x0F; + + // E8M0 scale: 2^(e - 127). e==0 is the canonical zero scale. + int e = scales[(long long)out_row * nblk + block]; + float scale = exp2f((float)e - 127.0f); + + int in_lo = block * 32 + 2 * b; // even element index in the block + int in_hi = in_lo + 1; // odd element + + // transposed write: out[in_idx, out_row] = out[in_idx*out_dim + out_row] + out[(long long)in_lo * out_dim + out_row] = __float2bfloat16(kFp4[code_lo] * scale); + out[(long long)in_hi * out_dim + out_row] = __float2bfloat16(kFp4[code_hi] * scale); +} + +extern "C" { + +// Dequantize one expert. `blocks`/`scales` point at this expert's slice. +// Output `out` must hold IN*OUT bf16 (IN = nblk*32). Runs on `stream`. +void launch_mxfp4_dequant_bf16( + const void* blocks, const void* scales, void* out, + int out_dim, int nblk, void* stream +) { + long long total = (long long)out_dim * nblk * 16; + int threads = 256; + int grid = (int)((total + threads - 1) / threads); + mxfp4_dequant_kernel<<>>( + (const uint8_t*)blocks, (const uint8_t*)scales, + (__nv_bfloat16*)out, out_dim, nblk + ); +} + +} // extern "C"