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

9
crates/xserv-model/src/bin/gptoss-logits.rs
View File

@@ -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 <bf16-model-dir> <tok0> <tok1> ...
//! differences). Usage: gptoss-logits <mxfp4-model-dir> <tok0> <tok1> ...
use std::path::PathBuf;
use half::bf16;
use xserv_model::loader;
@@ -13,10 +13,11 @@ fn main() {
let tokens: Vec<u32> = 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]