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