bug: Gemma 4 26B A4B temp-0 decode is non-deterministic run-to-run and corrupts near-tied tokens independently of the fused MoE kernel

[inureyes]

Summary

Greedy (temp 0) decode of gemma-4-26b-a4b-it-4bit on GB10 (CUDA sm_121) is not deterministic run-to-run: identical piped multi-turn repro runs of the same binary diverge bitwise mid-generation, and at a recurring near-tie hotspot ("During the Middle ___" in the wine paragraph) the flip sometimes lands on a garbage token (Hangul injection "During the Middle 이루어진 Middle Ages", wrong word "During the Middle War"). Critically, this reproduces with the fused decode-MoE kernel DISABLED, so it is a pipeline-level defect independent of the #886 kernel, and it invalidates #886's original discriminating experiment (MLXCEL_FUSED_MOE=0 was reported 5/5 clean during triage but corrupts on current main).

Evidence matrix

All runs: same binary (branch fix/issue-886-fused-moe-decode-parity, commit 5451e34, which includes the #886 f32-partials fix), same piped 5-turn repro from #886, /home/inureyes/models/gemma-4-26b-a4b-it-4bit, -t 0 -n 256, GB10.

• Fused MoE default (post-bug: Gemma 4 26B A4B multi-turn corruption from numerically unfaithful fused decode-MoE kernel #886 fix): 2 identical runs differ bitwise (first divergence in the tea paragraph, both continuations individually sane); 1 of 3 total runs produced "During the Middle War" in the wine paragraph.
• MLXCEL_FUSED_MOE=0 (reference gather_qmm path): 2 identical runs differ bitwise; one of the two contains "During the Middle 이루어진 Middle Ages" (4 Hangul characters injected mid-phrase).
• MLXCEL_FUSED_MOE=0 MLX_USE_CUDA_GRAPHS=0: 2 identical runs still differ bitwise (both happened to be clean of foreign script), so CUDA graphs are not the source.
• Single-turn control (ONE prompt, 256 tokens, MLXCEL_FUSED_MOE=0): 2 identical runs differ bitwise, so multi-turn KV rotation and chunked prefill are not required to trigger it; the source is in the base decode path (attention, matmul, norms, or logits).

Fused kernel exoneration (from #886)

The #886 in-situ probe (MLXCEL_FUSED_MOE_PARITY_CHECK=1) re-ran the fused kernel pair on identical inputs at every MoE call across a full repro run: 0 bitwise mismatches in 25600+ calls, and the isolated tests show byte-identical fused output across the entire MLXCEL_FUSED_MOE_SGY range. Post-fix the fused kernel matches an all-f32 dense reference at nrms 1.9e-6 (one seed bitwise identical). The non-determinism and residual corruption documented here are therefore not attributable to that kernel.

Candidate causes (to investigate)

• A non-deterministic base op in the MLX CUDA decode path: SDPA implementation (split-K or atomics), cuBLAS/cuBLASLt or cuDNN plan selection that varies per process (timing-based heuristics), or a reduction with run-order-dependent accumulation.
• fix(gemma4): size chunked-prefill masks to returned keys per family #891 changed gemma4 chunked-prefill mask sizing between the bug: Gemma 4 26B A4B multi-turn corruption from numerically unfaithful fused decode-MoE kernel #886 triage and now; whether the pre-fix(gemma4): size chunked-prefill masks to returned keys per family #891 reference path was bitwise deterministic on this repro is untested (the single-turn control above suggests the trigger does not require chunked prefill, but the triage baseline shift is unexplained).
• The model has a genuine near-tie at the recurring "During the Middle ___" position in structurally repetitive multi-turn output; run-to-run jitter picks different branches, and some branches emit implausible tokens, which suggests the logit gap at that step is abnormally small for a 26B model and may deserve its own look (4-bit quantization of the 262k-token embedding/unembedding).

Impact

• User-visible corrupted tokens at temp 0 on Gemma 4 26B A4B, regardless of the MoE path.
• Blocks the bug: Gemma 4 26B A4B multi-turn corruption from numerically unfaithful fused decode-MoE kernel #886 acceptance gate of 10 consecutive clean default-config repro runs, since the residual failure mode is outside that kernel.

Suggested next steps

1. At the first divergent token between two runs, dump per-layer activations (or logits) for both runs on identical context and bisect which op first differs.
2. Rebuild at 199fa40 (pre-fix(gemma4): size chunked-prefill masks to returned keys per family #891) and A/B the MLXCEL_FUSED_MOE=0 repro for bitwise determinism to settle whether the triage-era baseline was deterministic.
3. Audit the MLX CUDA SDPA and matmul paths used by gemma4 decode for atomics and per-process plan selection; try forcing deterministic algorithms.
4. Quantify the "During the Middle ___" logit gap with a clean forward pass at that context to establish how much jitter is needed to flip it.

Environment

GB10 (DGX Spark, CUDA sm_121), Linux 6.17, mlxcel HEAD of main plus #886 fix, MLX pin b7c3dd6d. Repro transcripts captured during the #886 investigation; regenerate with the #886 repro command.

Related: #886 (fused-kernel numeric faithfulness fix, includes the in-situ probe used here), #891 (gemma4 chunked-prefill mask change merged between the #886 triage and this finding).

[inureyes]

Root cause found: async race in the MLX CUDA quantized-matmul path (not gemma4-specific, not a near-tie)

Investigation result (GB10, temp 0, single-turn, bisected):

It is general to CUDA 4-bit quantized-matmul decode, independent of MoE and model architecture. The bf16/4bit fork is decisive:

Model	Quant	Result (2+ temp-0 runs, byte-diff)
llama-3.1-8b-bf16	none	deterministic (4/4 byte-identical)
llama-3.1-8b-4bit	4-bit	non-deterministic
llama-3.2-1b-4bit	4-bit	non-deterministic
qwen3-4b-4bit	4-bit	non-deterministic
gemma-4-26b-a4b-4bit	4-bit MoE	non-deterministic (also with MLXCEL_FUSED_MOE=0)
qwen3-30b-a3b-4bit	4-bit MoE	non-deterministic

So the non-deterministic op is the quantized-matmul dispatch (qmv/qmm_sm80/qmm_naive in mlx/backend/cuda/quantized/), not a shared base op (SDPA/RMSNorm/RoPE are atomics-free, fixed-order) and not the fused MoE kernel (exonerated by #886).

Two stacked sources:

• Layer A (all quantized models): CUDA_LAUNCH_BLOCKING=1 alone makes 4-bit dense llama fully deterministic; MLX_USE_CUDA_GRAPHS=0 does not. The qmm kernels have no atomics/split-K (deterministic by construction), so non-determinism that vanishes under host/device serialization is a missing stream-dependency edge / async-allocator race, not float-reduction noise. Prime suspect: the per-call cu::malloc_async(out.nbytes(), encoder) + kernel-node submission in the quantized path (the machinery that differs from the deterministic cuBLASLt dense-matmul path).
• Layer B (MoE-only residual, survives CUDA_LAUNCH_BLOCKING=1): router top-8 expert set flips run-to-run (e.g. {31,86,...} vs {57,86,...} by the 4th token) because the router input (128 expert-score logits = qmv outputs) differs by a hair and tips a ranking boundary. ArgPartition itself is deterministic. The MoE-only branch is the batched (M=1,B=8) gather_qmm -> qmm_sm80/qmm_naive dispatch; exact extra race not pinned within budget.

Not a near-tie / quantization issue. At "During the Middle ___", clean runs decide rank-1 "Ages" by a 14-17 logit margin (very confident). In a corrupted run "Ages" is not even in the top-5; the top-5 is all garbage ("War", hiragana, Hangul, "Middlehood", Ethiopic) clustered within ~0.5-0.8 logits at low confidence. The jitter displaces the correct token entirely, it does not nudge a tie. So fixing determinism is sufficient; embedding quantization needs no separate remediation.

Fix plan: (1) root-fix Layer A: add the missing dependency edge in the CUDA async-allocator/command-encoder around quantized-matmul output buffers; (2) re-run the determinism matrix; (3) if the MoE residual survives, bisect the batched gather_qmm path's temporary/JIT-module lifecycle; interim mitigation is an env-gated deterministic-qmm stream-sync (cheaper than global CUDA_LAUNCH_BLOCKING, which measured ~2-2.5x slower). This is an upstream MLX CUDA bug affecting all 4-bit CUDA models and should also be reported upstream.