quantization: MXFP4 W4A16 expert weights (memory-optimization foundation)

Weight-only 4-bit for the gpt-oss MoE experts: weights stored MXFP4 (E2M1 +
per-32-element UE8M0 block scale, tools/quantize_mxfp4.py), a fused kernel reads
the 4-bit weights and dequantizes on-chip to BF16. Decode (M=1) uses a fused
dequant-GEMV (batched_gemv_mxfp4) with shared-memory activation tiling; prefill
(M>1) dequantizes to BF16 then reuses the BF16 batched GEMM. MXFP4 is detected
by the scale tensor's rank (3-D [E,N,K/32]) vs FP8's 1-D [E].

Verified on dash5 (gpt-oss-20b, TP=2, 5090): byte-identical greedy tokens to
FP8/BF16, smallest footprint (13 GB vs 22 GB FP8, 39 GB BF16) — fits one 32 GB
5090 with room for KV cache.

NOT a decode speedup: the hand-written W4A16 GEMV (no tensor cores) is less
efficient than cuBLASLt's FP8 tensor-core GEMM, so even at half the weight bytes
decode is 17.0 ms vs FP8 13.5 ms (faster than BF16 18.8 ms); prefill regresses
(350 vs 134 ms, dequant fallback). Committed as a correct memory-optimization
foundation. Beating FP8 on speed needs FP4 tensor cores (W4A4, cuBLASLt
block-scaled MXFP4) or a Marlin-class kernel; see
docs/benchmarks/mxfp4-and-llama-decode.md.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

crates/xserv-kernels/src/quantization.rs

@@ -30,6 +30,14 @@ unsafe extern "C" {
         num_rows: i32, cols: i32, tokens: i32,
         stream: *mut c_void,
     );
+    fn launch_batched_gemv_mxfp4_bf16(
+        x: *const c_void, w_packed: *const c_void, w_scales: *const c_void, y: *mut c_void,
+        e: i32, n: i32, k: i32, stream: *mut c_void,
+    );
+    fn launch_dequant_mxfp4_to_bf16_t(
+        w_packed: *const c_void, w_scales: *const c_void, out: *mut c_void,
+        e: i32, n: i32, k: i32, stream: *mut c_void,
+    );
 }
 
 // ============================================================
@@ -428,3 +436,50 @@ pub fn batched_gemm_fp8(
 
     c
 }
+
+// ============================================================
+// MXFP4 W4A16 (weight-only 4-bit) for MoE experts
+// ============================================================
+
+/// MXFP4 W4A16 batched GEMV for decode (M=1).
+///
+/// x:        [E, K] BF16 (per-expert activation; replicated across experts)
+/// w_packed: [E, N, K/2] byte tensor — two E2M1 nibbles per byte (lo = even k)
+/// w_scales: [E, N, K/32] byte tensor — UE8M0 scale per 32-element block
+///
+/// Returns:  [E, N] BF16, where y[e,n] = sum_k x[e,k] * dequant(W[e,n,k]).
+pub fn batched_gemv_mxfp4(x: &Tensor, w_packed: &Tensor, w_scales: &Tensor, n: usize, k: usize) -> Tensor {
+    assert_eq!(x.dtype(), DType::BF16);
+    assert!(x.is_contiguous());
+    let e = x.shape()[0];
+    assert_eq!(x.shape()[x.ndim() - 1], k, "GEMV K mismatch");
+
+    let y = Tensor::empty(&[e, n], DType::BF16, x.device());
+    unsafe {
+        launch_batched_gemv_mxfp4_bf16(
+            x.data_ptr() as *const c_void,
+            w_packed.data_ptr() as *const c_void,
+            w_scales.data_ptr() as *const c_void,
+            y.data_ptr() as *mut c_void,
+            e as i32, n as i32, k as i32,
+            std::ptr::null_mut(),
+        );
+    }
+    y
+}
+
+/// Dequantize MXFP4 weights [E, N, K] → BF16 [E, K, N] for the prefill GEMM path
+/// (the BF16 batched GEMM expects weights as [E, K, N]).
+pub fn dequant_mxfp4_to_bf16_t(w_packed: &Tensor, w_scales: &Tensor, e: usize, n: usize, k: usize) -> Tensor {
+    let out = Tensor::empty(&[e, k, n], DType::BF16, w_packed.device());
+    unsafe {
+        launch_dequant_mxfp4_to_bf16_t(
+            w_packed.data_ptr() as *const c_void,
+            w_scales.data_ptr() as *const c_void,
+            out.data_ptr() as *mut c_void,
+            e as i32, n as i32, k as i32,
+            std::ptr::null_mut(),
+        );
+    }
+    out
+}