From 94957c57270192abd4620b044d86a24004a79819 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Fri, 29 May 2026 21:50:38 +0800 Subject: [PATCH] moe: MXFP4-resident experts on GPU (single-card gpt-oss) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Experts now stay MXFP4-packed on GPU (~10GB whole model, fits one 32GB card) instead of dequantized to ~38GB BF16. loader::load_model_dir_split returns BF16 tensors + raw U8 (_blocks/_scales) in one pass; GptOss slices each expert's MXFP4 bytes to a GpuBuffer at load, and expert_forward dequantizes the selected expert to a BF16 scratch (dequant_mxfp4) right before its GEMM — no per-token CPU->GPU upload, no 38GB BF16 dir. Verified: gptoss-logits on the original MXFP4 dir (/opt/wjh/models/gpt-oss-20b) gives logits byte-identical to the BF16 path — top-1 token 12650 = " Paris" @ 15.3125, full top-10 unchanged — running on a single GPU. Build green on dash5 (release). Co-Authored-By: Claude Opus 4.8 --- crates/xserv-model/src/bin/gptoss-logits.rs | 9 +- crates/xserv-model/src/gptoss.rs | 123 ++++++++++++-------- crates/xserv-model/src/loader.rs | 44 +++++++ 3 files changed, 125 insertions(+), 51 deletions(-) diff --git a/crates/xserv-model/src/bin/gptoss-logits.rs b/crates/xserv-model/src/bin/gptoss-logits.rs index f026694..7790b22 100644 --- a/crates/xserv-model/src/bin/gptoss-logits.rs +++ b/crates/xserv-model/src/bin/gptoss-logits.rs @@ -1,6 +1,6 @@ //! Dump gpt-oss next-token logits for a fixed token-id sequence, to compare //! against the llama.cpp oracle (isolates the model forward from tokenizer -//! differences). Usage: gptoss-logits ... +//! differences). Usage: gptoss-logits ... use std::path::PathBuf; use half::bf16; use xserv_model::loader; @@ -13,10 +13,11 @@ fn main() { let tokens: Vec = args[2..].iter().map(|s| s.parse().expect("token id")).collect(); assert!(!tokens.is_empty(), "need at least one token id"); + xserv_cuda::device::set_device(0).unwrap(); let config = ModelConfig::from_file(&model_dir.join("config.json")); - eprintln!("[gptoss-logits] loading {} ...", model_dir.display()); - let weights = loader::load_model_dir(&model_dir, Device::Cpu); - let model = GptOss::from_weights(config, weights); + eprintln!("[gptoss-logits] loading {} (MXFP4) ...", model_dir.display()); + let (floats, u8s) = loader::load_model_dir_split(&model_dir, Device::Cpu); + let model = GptOss::from_weights(config, floats, u8s); eprintln!("[gptoss-logits] forward over {} tokens", tokens.len()); let logits = model.forward(&tokens); // [T, vocab] diff --git a/crates/xserv-model/src/gptoss.rs b/crates/xserv-model/src/gptoss.rs index 2cd03c7..e71c097 100644 --- a/crates/xserv-model/src/gptoss.rs +++ b/crates/xserv-model/src/gptoss.rs @@ -8,9 +8,10 @@ //! - alternating sliding-window attention (window from config on flagged layers) //! - q/k/v/o projection biases; head_dim 64; no q/k norm; rotate_half RoPE (θ=150000) //! -//! Weights are loaded from a plain BF16 safetensors dir (MXFP4 experts are -//! dequantized to BF16 offline by tools/gptoss_dequant.py), so the standard -//! loader feeds us BF16 tensors and this file needs no quantization code. +//! Expert weights stay MXFP4-resident on GPU (~10GB for the whole model, fits +//! one 32GB card; BF16 would be ~38GB). Each selected expert is dequantized to a +//! BF16 scratch with our `dequant_mxfp4` kernel right before its GEMM. Dense +//! weights (attn/router/norms/embed/lm_head + expert biases) are BF16. //! //! v1 is a self-contained non-paged forward (contiguous KV built per call) used to //! validate next-token agreement with llama.cpp. Paged-cache + PP + server wiring @@ -18,6 +19,7 @@ use std::collections::HashMap; use half::bf16; +use xserv_cuda::GpuBuffer; use xserv_kernels::*; use xserv_tensor::{Device, Tensor}; @@ -46,23 +48,39 @@ struct Block { o_bias: Tensor, sinks: Tensor, // [n_heads] (f32 on host) sliding: bool, - // MoE. + // MoE. Experts stay MXFP4 on GPU; dequantized per use (see expert_forward). router_wt: Tensor, // [hidden, n_experts] router_bias: Tensor, // [n_experts] - gate_up_wt: Vec, // per expert: [hidden, 2*inter] - gate_up_bias: Vec, // [2*inter] - down_wt: Vec, // per expert: [inter, hidden] - down_bias: Vec, // [hidden] + gate_up_blocks: Vec, // per expert: U8 [2*inter, nblk, 16] + gate_up_scales: Vec, // per expert: U8 [2*inter, nblk] + gate_up_bias: Vec, // [2*inter] BF16 + down_blocks: Vec, // per expert: U8 [hidden, nblk, 16] + down_scales: Vec, // per expert: U8 [hidden, nblk] + down_bias: Vec, // [hidden] BF16 + gate_up_out: usize, // 2*inter (dequant out_dim) + gate_up_nblk: usize, // hidden/32 + down_out: usize, // hidden + down_nblk: usize, // inter/32 } impl GptOss { - /// Load gpt-oss from a BF16 (dequantized) HF-format weight map. - pub fn from_weights(config: ModelConfig, mut w: HashMap) -> Self { + /// Load gpt-oss from the original MXFP4 HF dir. `floats` holds the BF16 + /// tensors, `u8s` the MXFP4 expert `_blocks`/`_scales` (both from + /// `loader::load_model_dir_split`). Experts are sliced per-expert and kept on + /// GPU as MXFP4; dense weights are uploaded as BF16. + pub fn from_weights( + config: ModelConfig, + mut w: HashMap, + mut u8s: HashMap, Vec)>, + ) -> Self { crate::init_kernels(); let dev = Device::Cuda(0); let take = |w: &mut HashMap, n: &str| -> Tensor { w.remove(n).unwrap_or_else(|| panic!("missing weight: {n}")) }; + let take_u8 = |u: &mut HashMap, Vec)>, n: &str| -> (Vec, Vec) { + u.remove(n).unwrap_or_else(|| panic!("missing MXFP4 tensor: {n}")) + }; let repl = |t: Tensor| t.to_device(dev); // pre-transpose a [out, in] linear weight to [in, out] for x@wt. let wt = |t: Tensor| t.to_device(dev).transpose(0, 1).contiguous(); @@ -70,6 +88,10 @@ impl GptOss { let hidden = config.hidden(); let n_experts = config.num_experts(); let inter = config.intermediate_size.expect("intermediate_size"); + // dequant dims: gate_up out=2*inter, in=hidden (nblk=hidden/32); + // down out=hidden, in=inter (nblk=inter/32). + let (gate_up_out, gate_up_nblk) = (2 * inter, hidden / 32); + let (down_out, down_nblk) = (hidden, inter / 32); let embed_tokens = repl(take(&mut w, "model.embed_tokens.weight")); let norm = repl(take(&mut w, "model.norm.weight")); @@ -81,24 +103,33 @@ impl GptOss { let mut layers = Vec::with_capacity(n_layers); for i in 0..n_layers { let p = format!("model.layers.{i}"); - // Experts are stored fused as [E, in, out]; slice per expert into [in, out]. - let gate_up = take(&mut w, &format!("{p}.mlp.experts.gate_up_proj")); // [E, hidden, 2*inter] + // MXFP4 expert tensors: blocks [E, OUT, nblk, 16], scales [E, OUT, nblk]. + let (gu_blk, _) = take_u8(&mut u8s, &format!("{p}.mlp.experts.gate_up_proj_blocks")); + let (gu_scl, _) = take_u8(&mut u8s, &format!("{p}.mlp.experts.gate_up_proj_scales")); + let (dn_blk, _) = take_u8(&mut u8s, &format!("{p}.mlp.experts.down_proj_blocks")); + let (dn_scl, _) = take_u8(&mut u8s, &format!("{p}.mlp.experts.down_proj_scales")); let gate_up_b = take(&mut w, &format!("{p}.mlp.experts.gate_up_proj_bias")); // [E, 2*inter] - let down = take(&mut w, &format!("{p}.mlp.experts.down_proj")); // [E, inter, hidden] - let down_b = take(&mut w, &format!("{p}.mlp.experts.down_proj_bias")); // [E, hidden] - let mut gate_up_wt = Vec::with_capacity(n_experts); + let down_b = take(&mut w, &format!("{p}.mlp.experts.down_proj_bias")); // [E, hidden] + + // per-expert byte spans + let gu_blk_pe = gate_up_out * gate_up_nblk * 16; + let gu_scl_pe = gate_up_out * gate_up_nblk; + let dn_blk_pe = down_out * down_nblk * 16; + let dn_scl_pe = down_out * down_nblk; + + let mut gate_up_blocks = Vec::with_capacity(n_experts); + let mut gate_up_scales = Vec::with_capacity(n_experts); + let mut down_blocks = Vec::with_capacity(n_experts); + let mut down_scales = Vec::with_capacity(n_experts); let mut gate_up_bias = Vec::with_capacity(n_experts); - let mut down_wt = Vec::with_capacity(n_experts); let mut down_bias = Vec::with_capacity(n_experts); - // Experts are kept on CPU (the 32 experts per layer total ~36GB for - // the whole model, which won't fit one GPU). Each selected expert's - // weights (~50MB) are uploaded on demand in expert_forward; only - // top-k experts per token are touched, so the H2D traffic is small. for e in 0..n_experts { - gate_up_wt.push(slice_expert(&gate_up, e, hidden, 2 * inter)); // CPU - gate_up_bias.push(slice_row(&gate_up_b, e, 2 * inter)); // CPU - down_wt.push(slice_expert(&down, e, inter, hidden)); // CPU - down_bias.push(slice_row(&down_b, e, hidden)); // CPU + gate_up_blocks.push(upload_u8(&gu_blk[e * gu_blk_pe..(e + 1) * gu_blk_pe])); + gate_up_scales.push(upload_u8(&gu_scl[e * gu_scl_pe..(e + 1) * gu_scl_pe])); + down_blocks.push(upload_u8(&dn_blk[e * dn_blk_pe..(e + 1) * dn_blk_pe])); + down_scales.push(upload_u8(&dn_scl[e * dn_scl_pe..(e + 1) * dn_scl_pe])); + gate_up_bias.push(slice_row(&gate_up_b, e, gate_up_out).to_device(dev)); + down_bias.push(slice_row(&down_b, e, down_out).to_device(dev)); } layers.push(Block { @@ -116,7 +147,9 @@ impl GptOss { sliding: config.layer_uses_sliding_window(i), router_wt: wt(take(&mut w, &format!("{p}.mlp.router.weight"))), router_bias: repl(take(&mut w, &format!("{p}.mlp.router.bias"))), - gate_up_wt, gate_up_bias, down_wt, down_bias, + gate_up_blocks, gate_up_scales, gate_up_bias, + down_blocks, down_scales, down_bias, + gate_up_out, gate_up_nblk, down_out, down_nblk, }); } @@ -204,8 +237,7 @@ impl GptOss { let xr = row_view(x, ti); let mut acc: Option = None; for (j, &e) in top.iter().enumerate() { - let y = expert_forward(&xr, &layer.gate_up_wt[e], &layer.gate_up_bias[e], - &layer.down_wt[e], &layer.down_bias[e], limit); // [1, hidden] + let y = expert_forward(&xr, layer, e, limit); // [1, hidden] let yw = scale_tensor(&y, weights[j]); acc = Some(match acc { Some(a) => add(&a, &yw), None => yw }); } @@ -275,20 +307,18 @@ fn matmul2(a: &Tensor, b: &Tensor) -> Tensor { matmul(a, b, GemmBackend::CuBlas) } -/// One expert: clamped SwiGLU. x:[*,hidden] -> [*,hidden]. -/// gate_up = x@gate_up_wt + bias; gate=even cols, up=odd cols (interleaved); -/// gate.clamp(max=limit); up.clamp(-limit,limit); h=(up+1)*gate*sigmoid(gate*1.702); h@down_wt+bias. -fn expert_forward(x: &Tensor, gate_up_wt: &Tensor, gate_up_bias: &Tensor, - down_wt: &Tensor, down_bias: &Tensor, limit: f32) -> Tensor { - // Upload this expert's CPU-resident weights to x's device just for this call. - let dev = x.device(); - let gate_up_wt = gate_up_wt.to_device(dev); - let gate_up_bias = gate_up_bias.to_device(dev); - let down_wt = down_wt.to_device(dev); - let down_bias = down_bias.to_device(dev); - let gate_up = add_bias(&matmul2(x, &gate_up_wt), &gate_up_bias); // [*, 2*inter] - let h = clamped_swiglu(&gate_up, limit); // [*, inter] - add_bias(&matmul2(&h, &down_wt), &down_bias) // [*, hidden] +/// One expert `e` of `layer`: clamped SwiGLU. x:[*,hidden] -> [*,hidden]. +/// Dequantizes the expert's MXFP4 weights to BF16 scratch on GPU, then: +/// gate_up = x@gate_up + bias; gate=even cols, up=odd cols (interleaved); +/// gate.clamp(max=limit); up.clamp(-limit,limit); h=(up+1)*gate*sigmoid(gate*1.702); h@down+bias. +fn expert_forward(x: &Tensor, layer: &Block, e: usize, limit: f32) -> Tensor { + let gate_up_w = dequant_mxfp4(&layer.gate_up_blocks[e], &layer.gate_up_scales[e], + layer.gate_up_out, layer.gate_up_nblk, 0); // [hidden, 2*inter] + let gate_up = add_bias(&matmul2(x, &gate_up_w), &layer.gate_up_bias[e]); // [*, 2*inter] + let h = clamped_swiglu(&gate_up, limit); // [*, inter] + let down_w = dequant_mxfp4(&layer.down_blocks[e], &layer.down_scales[e], + layer.down_out, layer.down_nblk, 0); // [inter, hidden] + add_bias(&matmul2(&h, &down_w), &layer.down_bias[e]) // [*, hidden] } /// Clamped interleaved SwiGLU on host (correctness-first). [*, 2I] -> [*, I]. @@ -372,12 +402,11 @@ fn add_bias(x: &Tensor, bias: &Tensor) -> Tensor { add(x, &broadcast) } -/// Slice expert `e` out of a fused [E, rows, cols] tensor -> [rows, cols]. -fn slice_expert(t: &Tensor, e: usize, rows: usize, cols: usize) -> Tensor { - let host = t.to_device(Device::Cpu); - let s = host.as_slice::(); - let stride = rows * cols; - Tensor::from_slice(&s[e * stride..(e + 1) * stride], &[rows, cols]) +/// Upload raw U8 bytes (an MXFP4 expert slice) to a GPU buffer. +fn upload_u8(bytes: &[u8]) -> GpuBuffer { + let mut buf = GpuBuffer::alloc(bytes.len()).expect("alloc expert U8"); + buf.copy_from_host(bytes).expect("upload expert U8"); + buf } /// Slice row `e` out of [E, n] -> [n]. diff --git a/crates/xserv-model/src/loader.rs b/crates/xserv-model/src/loader.rs index 00c6815..a0539f5 100644 --- a/crates/xserv-model/src/loader.rs +++ b/crates/xserv-model/src/loader.rs @@ -63,6 +63,50 @@ pub fn load_model_dir(dir: &Path, device: Device) -> HashMap { all_tensors } +/// Load a model dir splitting tensors by dtype: float tensors (F32/F16/BF16) +/// become `Tensor`s on `device`; U8 tensors (gpt-oss MXFP4 `_blocks`/`_scales`, +/// which are not an xserv Tensor dtype) are returned as raw `(bytes, shape)`. +/// One pass over the shards (the 13GB MXFP4 file is read once). +pub fn load_model_dir_split( + dir: &Path, device: Device, +) -> (HashMap, HashMap, Vec)>) { + let mut files: Vec = Vec::new(); + let single = dir.join("model.safetensors"); + if single.exists() { + files.push(single); + } else { + let mut entries: Vec<_> = std::fs::read_dir(dir).unwrap() + .filter_map(|e| e.ok()) + .filter(|e| e.path().file_name() + .map(|f| f.to_string_lossy().ends_with(".safetensors")).unwrap_or(false)) + .collect(); + entries.sort_by_key(|e| e.file_name()); + files.extend(entries.into_iter().map(|e| e.path())); + } + assert!(!files.is_empty(), "no safetensors files in {}", dir.display()); + + let mut floats = HashMap::new(); + let mut u8s = HashMap::new(); + for path in &files { + let data = std::fs::read(path) + .unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display())); + let st = SafeTensors::deserialize(&data) + .unwrap_or_else(|e| panic!("failed to parse {}: {e}", path.display())); + for (name, view) in st.tensors() { + let shape: Vec = view.shape().to_vec(); + let raw = view.data(); + match view.dtype() { + safetensors::Dtype::F32 => { floats.insert(name, make_tensor(raw, &shape, DType::F32).to_device(device)); } + safetensors::Dtype::F16 => { floats.insert(name, make_tensor(raw, &shape, DType::F16).to_device(device)); } + safetensors::Dtype::BF16 => { floats.insert(name, make_tensor(raw, &shape, DType::BF16).to_device(device)); } + safetensors::Dtype::U8 => { u8s.insert(name, (raw.to_vec(), shape)); } + other => eprintln!("load_model_dir_split: skipping {name}: dtype {other:?}"), + } + } + } + (floats, u8s) +} + fn make_tensor(raw_bytes: &[u8], shape: &[usize], dtype: DType) -> Tensor { match dtype { DType::F32 => {