gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

moe: MXFP4-resident experts on GPU (single-card gpt-oss)

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-29 21:50:38 +08:00
parent f59ba73938
commit 94957c5727
3 changed files with 125 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

123
crates/xserv-model/src/gptoss.rs
View File

@@ -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<Tensor>, // per expert: [hidden, 2*inter]
gate_up_bias: Vec<Tensor>, // [2*inter]
down_wt: Vec<Tensor>, // per expert: [inter, hidden]
down_bias: Vec<Tensor>, // [hidden]
gate_up_blocks: Vec<GpuBuffer>, // per expert: U8 [2*inter, nblk, 16]
gate_up_scales: Vec<GpuBuffer>, // per expert: U8 [2*inter, nblk]
gate_up_bias: Vec<Tensor>, // [2*inter] BF16
down_blocks: Vec<GpuBuffer>, // per expert: U8 [hidden, nblk, 16]
down_scales: Vec<GpuBuffer>, // per expert: U8 [hidden, nblk]
down_bias: Vec<Tensor>, // [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<String, Tensor>) -> 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<String, Tensor>,
mut u8s: HashMap<String, (Vec<u8>, Vec<usize>)>,
) -> Self {
crate::init_kernels();
let dev = Device::Cuda(0);
let take = |w: &mut HashMap<String, Tensor>, n: &str| -> Tensor {
w.remove(n).unwrap_or_else(|| panic!("missing weight: {n}"))
};
let take_u8 = |u: &mut HashMap<String, (Vec<u8>, Vec<usize>)>, n: &str| -> (Vec<u8>, Vec<usize>) {
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<Tensor> = 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::<bf16>();
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].