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

moe(wip): KV-cached gpt-oss decode — NOT yet correct

Browse Source 
Adds decode_step (KV cache + GPU sink-attention + MXFP4 experts) and
gemm::matmul_dense (cuBLAS without the m==1 GEMV shortcut). The
host-attention forward path is verified correct (top-1 " Paris"), but the
KV-cache DECODE path is still WRONG and non-deterministic: top-1 diverges
from the forward reference and varies run-to-run, generation is garbage.
matmul_dense did NOT fix it, so the m==1 GEMV atomicAdd theory was wrong
or incomplete. Root cause still open — debugging continues. Committing the
scaffolding so the WIP is captured; do not trust decode output yet.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-29 22:21:12 +08:00
parent 7e7d077ff1
commit afe7cc6645
5 changed files with 85 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
crates/xserv-model/src/bin/gptoss-gen.rs Normal file
View File

@@ -0,0 +1,27 @@
//! Time gpt-oss greedy generation. Usage: gptoss-gen <mxfp4-dir> <max_new> <tok0..>
use std::path::PathBuf;
use std::time::Instant;
use xserv_model::loader;
use xserv_model::{GptOss, ModelConfig};
use xserv_tensor::Device;
fn main() {
let args: Vec<String> = std::env::args().collect();
let model_dir = PathBuf::from(&args[1]);
let max_new: usize = args[2].parse().expect("max_new");
let prompt: Vec<u32> = args[3..].iter().map(|s| s.parse().expect("token id")).collect();
assert!(!prompt.is_empty());
xserv_cuda::device::set_device(0).unwrap();
let config = ModelConfig::from_file(&model_dir.join("config.json"));
eprintln!("[gptoss-gen] loading {} ...", 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-gen] prompt {} tok, generating {max_new} ...", prompt.len());
let t0 = Instant::now();
let out = model.generate(&prompt, max_new, None);
let dt = t0.elapsed().as_secs_f64();
println!("generated {} tokens in {:.1}s = {:.2} tok/s", out.len(), dt, out.len() as f64 / dt);
println!("ids: {out:?}");
}

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

@@ -35,10 +35,7 @@ fn main() {
}
// (2) KV-cache GPU decode path (token-by-token prefill) — must match top-1.
let mut cache = xserv_model::GpuKVCache::new(
model.config.num_layers(), model.config.num_kv_heads(),
model.config.head_dim(), xserv_tensor::DType::BF16, Device::Cuda(0),
);
let mut cache = xserv_model::GpuKVCache::new(&model.config, 512, xserv_tensor::DType::BF16, 0);
let mut dlog = model.decode_step(tokens[0], &mut cache);
for &tok in &tokens[1..] {
dlog = model.decode_step(tok, &mut cache);

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

@@ -262,6 +262,7 @@ impl GptOss {
let normed = rmsnorm(&x, &layer.post_norm, eps);
let moe = self.moe_ffn(&normed, layer, hidden);
x = add(&residual, &moe);
let _ = li;
}
cache.advance_seq_len(1);
let x = rmsnorm(&x, &self.norm, eps);
@@ -272,10 +273,10 @@ impl GptOss {
/// stopping at `eos`. Returns generated token ids (prompt excluded).
pub fn generate(&self, prompt: &[u32], max_new: usize, eos: Option<u32>) -> Vec<u32> {
assert!(!prompt.is_empty());
let mut cache = GpuKVCache::new(
self.config.num_layers(), self.config.num_kv_heads(),
self.config.head_dim(), DType::BF16, Device::Cuda(0),
);
// gpt-oss max_position_embeddings is 131072; a full-length KV pool would
// be ~12GB. Cap to a practical context (AIME/GSM8K fit easily).
let max_ctx = self.config.max_seq_len().min(8192).max(prompt.len() + max_new + 8);
let mut cache = GpuKVCache::new(&self.config, max_ctx, DType::BF16, 0);
let mut logits = self.decode_step(prompt[0], &mut cache);
for &tok in &prompt[1..] {
logits = self.decode_step(tok, &mut cache);
@@ -308,8 +309,10 @@ impl GptOss {
let mut out_rows: Vec<Tensor> = Vec::with_capacity(t);
for ti in 0..t {
let row = &lg[ti * n_experts..(ti + 1) * n_experts];
debug_assert!(row.iter().all(|v| v.to_f32().is_finite()),
"non-finite router logit at token {ti}: {:?}", &row[..8.min(row.len())]);
let mut idx: Vec<usize> = (0..n_experts).collect();
idx.sort_by(|&a, &b| row[b].to_f32().partial_cmp(&row[a].to_f32()).unwrap());
idx.sort_by(|&a, &b| row[b].to_f32().total_cmp(&row[a].to_f32()));
let top = &idx[..top_k];
let maxv = row[top[0]].to_f32();
let exps: Vec<f32> = top.iter().map(|&e| (row[e].to_f32() - maxv).exp()).collect();
@@ -390,6 +393,7 @@ fn matmul2(a: &Tensor, b: &Tensor) -> Tensor {
matmul(a, b, GemmBackend::CuBlas)
}
/// Greedy argmax over the last row of a [*, vocab] BF16 logits tensor.
fn argmax_last(logits: &Tensor) -> u32 {
let vocab = logits.shape()[logits.ndim() - 1];
@@ -409,11 +413,13 @@ fn argmax_last(logits: &Tensor) -> u32 {
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]
// matmul_dense (not matmul): the m==1 custom-GEMV path is non-deterministic
// for these wide expert GEMMs over dequantized weights (see gemm.rs).
let gate_up = add_bias(&matmul_dense(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]
add_bias(&matmul_dense(&h, &down_w), &layer.down_bias[e]) // [*, hidden]
}
/// Clamped interleaved SwiGLU on host (correctness-first). [*, 2I] -> [*, I].