moe: FIX decode — route all gpt-oss matmuls through dense cuBLAS

Root cause of the broken/non-deterministic KV-cache decode: matmul's m==1
fast path is a custom GEMV that reduces over K with a grid-split atomicAdd,
whose float accumulation order is non-deterministic. At decode (m==1) the
router logits jittered enough to flip the top-4 EXPERT SELECTION run-to-run,
so each step picked different experts -> wrong, non-deterministic output and
garbage generation. (Earlier partial fixes failed because only the expert
GEMMs were switched; router/qkv/o_proj still used the GEMV path.)

Fix: gptoss matmul2 -> gemm::matmul_dense (plain cublasGemmEx, no GEMV) for
every matmul. Forward (m>1) already used dense cuBLAS, unchanged + correct.

Verified THIS run (read full output file):
  gptoss-logits "The capital of France is":
    forward top5 == decode top5, byte for byte
      12650 " Paris" 15.50, 625, 25, 354, 261 ;  MATCH_TOP1: YES
    forward 3/3 identical, decode 3/3 identical (deterministic).
  gptoss-gen: first token 12650 " Paris", ~10.4 tok/s on one 5090.

KV-cache GPU decode (sink-attention + MXFP4 experts + YaRN RoPE) is now
correct and deterministic. Next: harmony chat template + server + AIME/GSM8K.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-05-29 22:34:55 +08:00
parent fcd7fa62b7
commit b8868d59ac

View File

@@ -389,8 +389,14 @@ fn yarn_rope_cache(config: &ModelConfig) -> RopeCache {
RopeCache { cos: cos_buf, sin: sin_buf, max_seq_len: max_seq, half_dim: half }
}
// All gpt-oss matmuls use the dense cuBLAS path (never the m==1 custom GEMV).
// The GEMV kernel reduces over K with a grid-split atomicAdd whose float order
// is non-deterministic; at decode (m==1) that made the router logits jitter
// enough to flip the top-4 expert selection run-to-run, so decode output was
// wrong and non-deterministic. Dense cuBLAS is deterministic for fixed inputs.
// Forward (m>1) already used dense cuBLAS, so it is unaffected.
fn matmul2(a: &Tensor, b: &Tensor) -> Tensor {
matmul(a, b, GemmBackend::CuBlas)
matmul_dense(a, b)
}