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

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>

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.

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};

crates/xserv-model/src/bin/gptoss-gen.rs

@@ -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:?}");
+}

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(
+    let mut cache = xserv_model::GpuKVCache::new(&model.config, 512, xserv_tensor::DType::BF16, 0);
-        model.config.num_layers(), model.config.num_kv_heads(),
-        model.config.head_dim(), xserv_tensor::DType::BF16, Device::Cuda(0),
-    );
     let mut dlog = model.decode_step(tokens[0], &mut cache);
     for &tok in &tokens[1..] {
         dlog = model.decode_step(tok, &mut cache);

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<u32>) -> Vec<u32> {
         assert!(!prompt.is_empty());
-        let mut cache = GpuKVCache::new(
+        // gpt-oss max_position_embeddings is 131072; a full-length KV pool would
-            self.config.num_layers(), self.config.num_kv_heads(),
+        // be ~12GB. Cap to a practical context (AIME/GSM8K fit easily).
-            self.config.head_dim(), DType::BF16, Device::Cuda(0),
+        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].