From 603b6eb27026b9077c45af35abf73e52da057c1d Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Fri, 29 May 2026 22:21:12 +0800 Subject: [PATCH] moe: correct + deterministic KV-cached gpt-oss decode (single card) 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 --- crates/xserv-kernels/src/gemm.rs | 42 +++++++++++++++++++++ crates/xserv-kernels/src/lib.rs | 4 +- crates/xserv-model/src/bin/gptoss-gen.rs | 27 +++++++++++++ crates/xserv-model/src/bin/gptoss-logits.rs | 5 +-- crates/xserv-model/src/gptoss.rs | 20 ++++++---- 5 files changed, 85 insertions(+), 13 deletions(-) create mode 100644 crates/xserv-model/src/bin/gptoss-gen.rs diff --git a/crates/xserv-kernels/src/gemm.rs b/crates/xserv-kernels/src/gemm.rs index 916ea6a..ef4263d 100644 --- a/crates/xserv-kernels/src/gemm.rs +++ b/crates/xserv-kernels/src/gemm.rs @@ -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. diff --git a/crates/xserv-kernels/src/lib.rs b/crates/xserv-kernels/src/lib.rs index 75ca199..923b66b 100644 --- a/crates/xserv-kernels/src/lib.rs +++ b/crates/xserv-kernels/src/lib.rs @@ -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}; diff --git a/crates/xserv-model/src/bin/gptoss-gen.rs b/crates/xserv-model/src/bin/gptoss-gen.rs new file mode 100644 index 0000000..00d1313 --- /dev/null +++ b/crates/xserv-model/src/bin/gptoss-gen.rs @@ -0,0 +1,27 @@ +//! Time gpt-oss greedy generation. Usage: gptoss-gen +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 = std::env::args().collect(); + let model_dir = PathBuf::from(&args[1]); + let max_new: usize = args[2].parse().expect("max_new"); + let prompt: Vec = 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:?}"); +} diff --git a/crates/xserv-model/src/bin/gptoss-logits.rs b/crates/xserv-model/src/bin/gptoss-logits.rs index 7f19648..8442bc8 100644 --- a/crates/xserv-model/src/bin/gptoss-logits.rs +++ b/crates/xserv-model/src/bin/gptoss-logits.rs @@ -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); diff --git a/crates/xserv-model/src/gptoss.rs b/crates/xserv-model/src/gptoss.rs index 046c833..64a23d3 100644 --- a/crates/xserv-model/src/gptoss.rs +++ b/crates/xserv-model/src/gptoss.rs @@ -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) -> Vec { 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 = 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 = (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 = 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].