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 <noreply@anthropic.com>

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;

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

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 <mxfp4-model-dir>
+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::<half::bf16>();
+    // out[in_idx, 0] = s[in_idx*out_dim + 0]
+    let vals: Vec<f32> = (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]");
+}

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 <cuda_bf16.h>
+#include <cstdint>
+
+// 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<<<grid, threads, 0, (cudaStream_t)stream>>>(
+        (const uint8_t*)blocks, (const uint8_t*)scales,
+        (__nv_bfloat16*)out, out_dim, nblk
+    );
+}
+
+}  // extern "C"