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]");
+}