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

moe: correct + deterministic KV-cached gpt-oss decode (single card)

Browse Source 
Root-caused the decode non-determinism: matmul's m==1 custom-GEMV fast path
reduces over K with a grid-split atomicAdd, whose float accumulation order
is non-deterministic. Negligible for attention's stable pre-transposed
weights, but for gpt-oss's wide expert GEMMs (K=2880, N up to 5760) over
freshly-dequantized MXFP4 weights it produced visibly different results
run-to-run (and a wrong argmax). Added gemm::matmul_dense (plain cublasGemmEx,
no GEMV shortcut) and route the expert GEMMs through it.

Now decode_step (KV cache + GPU sink-attention + MXFP4 experts) is:
- deterministic: 3/3 identical runs
- correct: top-1 token 12650 = " Paris" for "The capital of France is",
  MATCH_TOP1 with the host-attention reference forward
- end-to-end: gptoss-gen generates 32 tokens at ~6.85 tok/s on one 5090.

Removed the temporary A/B debug dumps. gptoss-logits runs both paths and
asserts the top-1 match; gptoss-gen times greedy generation.

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 603b6eb270
5 changed files with 85 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
crates/xserv-kernels/src/gemm.rs
View File

@@ -201,6 +201,48 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
c
}
/// Dense cuBLAS GEMM that never takes the m==1 custom-GEMV fast path.
///
/// The custom GEMV kernel reduces over K with a grid-split `atomicAdd`, whose
/// float accumulation order is non-deterministic. For most decode matmuls
/// (stable pre-transposed weights) the effect is negligible, but for gpt-oss's
/// wide expert GEMMs (K=2880, N up to 5760) over freshly-dequantized MXFP4
/// weights it produces visibly different results run-to-run. Routing those
/// matmuls here (plain `cublasGemmEx`) makes the MoE forward deterministic and
/// matches the batched (m>1) reference path.
pub fn matmul_dense(a: &Tensor, b: &Tensor) -> Tensor {
assert_eq!(a.ndim(), 2);
assert_eq!(b.ndim(), 2);
assert_eq!(a.shape()[1], b.shape()[0], "inner dimension mismatch");
assert_eq!(a.dtype(), b.dtype(), "dtype mismatch");
assert!(a.is_contiguous() && b.is_contiguous(), "matmul_dense requires contiguous");
assert!(matches!(a.device(), Device::Cuda(_)));
let (m, k, n) = (a.shape()[0], a.shape()[1], b.shape()[1]);
let dtype = a.dtype();
let c = Tensor::empty(&[m, n], dtype, a.device());
let (a_ptr, b_ptr, c_ptr) = (a.data_ptr() as *const c_void, b.data_ptr() as *const c_void, c.data_ptr() as *mut c_void);
let (alpha, beta) = (1.0f32, 0.0f32);
let (a_type, b_type, c_type) = match dtype {
DType::F32 => (CUDA_R_32F, CUDA_R_32F, CUDA_R_32F),
DType::BF16 => (CUDA_R_16BF, CUDA_R_16BF, CUDA_R_16BF),
_ => panic!("unsupported dtype for matmul_dense"),
};
with_cublas(|handle| unsafe {
cublasSetStream_v2(handle, std::ptr::null_mut());
error::check(cublasGemmEx(
handle, CUBLAS_OP_N, CUBLAS_OP_N,
n as i32, m as i32, k as i32,
&alpha as *const f32 as *const c_void,
b_ptr, b_type, n as i32,
a_ptr, a_type, k as i32,
&beta as *const f32 as *const c_void,
c_ptr, c_type, n as i32,
CUBLAS_COMPUTE_32F, -1,
)).expect("cuBLAS GEMM failed");
});
c
}
/// Batched matrix multiplication via cuBLAS: C[b] = A[b] @ B[b]
/// a: [..., M, K], b: [..., K, N] → [..., M, N]
/// Leading dimensions must match and tensors must be contiguous.

4
crates/xserv-kernels/src/lib.rs
View File

@@ -12,9 +12,9 @@ pub mod transpose;
pub use activation::{add, gelu, mul, scale, silu, silu_mul};
pub use transpose::{merge_heads_gpu, repeat_kv_gpu, reshape_heads_gpu, strided_to_contiguous_gpu, transpose_for_rope_gpu, transpose_from_rope_gpu};
pub use attention::{attention, decode_attention, flash_attention, paged_decode_attention};
pub use attention::{attention, decode_attention, decode_attention_sink, flash_attention, paged_decode_attention};
pub use embedding::embedding;
pub use gemm::{batched_matmul, matmul, GemmBackend};
pub use gemm::{batched_matmul, matmul, matmul_dense, GemmBackend};
pub use layernorm::layernorm;
pub use quant::dequant_mxfp4;
pub use rmsnorm::{add_rmsnorm, rmsnorm};

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].