Repository navigation
Description
Summary
After #910 (the qmm_sm80 async race fix) made 4-bit CUDA decode deterministic, gemma-4-26b-a4b-it-4bit still deterministically emits a semantically wrong token at one position in a specific long multi-turn context: the wine paragraph of the 5-paragraph repro reads "During the Middle War, monasteries became vital centers for preserving winemaking" where it should read "Middle Ages". The rest of the paragraph is coherent. This is NOT the #910 race (see evidence) and NOT random garbage (no foreign-script, fully deterministic). It is a standard-4-bit (RTN/affine) quantization-quality artifact: the quantization-aware-trained variant does not exhibit it.
Reproduction
printf '%s\n' \
"Remember this secret code: PARETO-7391. Now write a detailed 160-word paragraph about the history of coffee." \
"Now write a detailed 160-word paragraph about the history of tea." \
"Write a detailed 160-word paragraph about the history of chocolate." \
"Write a detailed 160-word paragraph about the history of wine." \
"Write a detailed 160-word paragraph about the history of bread." \
"/bye" | ./target/release/mlxcel run <path>/gemma-4-26b-a4b-it-4bit -t 0 -n 256
Wine paragraph contains "During the Middle War, monasteries...". Reproduces 4/4 runs, byte-identical.
Evidence that this is quantization quality, not a runtime bug
- Deterministic and matches the serialized oracle: 4/4 repro runs are byte-identical, and the default (non-blocking) output hash equals the
CUDA_LAUNCH_BLOCKING=1output hash exactly. So the qmm race (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 #910) is fully closed here; "War" is the true deterministic 4-bit greedy result, not a race remnant. - No foreign-script / no garbage: unlike the pre-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 #910 corruption (Japanese/Hangul/Ethiopic injections), the output is clean prose with a single wrong word.
- QAT variant is correct:
gemma-4-26b-a4b-it-qat-4biton the identical repro does NOT emit "Middle War" (only correct "Middle Ages"/"Middle East"). Same architecture, same runtime, different quantization -> points squarely at standard RTN/affine 4-bit weight quantization of this checkpoint. - Context-dependent: a short single-turn wine prompt yields "Middle Ages"; the error only surfaces in the long accumulated multi-turn context (the decision margin at that position is evidently degraded by 4-bit only in that regime).
Root-cause hypothesis (to confirm)
Standard 4-bit (affine, group_size 64) quantization degrades one near-decision in a long-context regime enough to flip the greedy argmax. The #910 investigation measured (pre-fix, on cleaner contexts) that "Ages" normally wins by 14-17 logits, so the flip is context-specific; the likely sensitive component is the 262k-token embedding/unembedding (logit head) and/or specific decoder layers under 4-bit. The QAT variant's immunity supports a weight-precision cause.
Investigation and fix directions
- Quantify the margin post-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 #910: dump top-k logits at the exact "During the Middle ___" decode step in the multi-turn context, for standard-4bit vs qat-4bit, to measure how narrowly "War" beats "Ages" and whether the logit head is the dominant contributor.
- Test targeted higher-precision: does keeping the unembedding / logit head (or the embedding) at higher precision (e.g. 8-bit or bf16) eliminate the flip? If a small, bounded precision bump on one sensitive tensor fixes it at acceptable memory/perf cost, that is the fix (and a general 4-bit quality win, not gemma4-specific).
- Confirm QAT robustness across a few more prompts (is it a reliable mitigation to recommend).
- If no bounded code fix exists (i.e. it is inherent RTN 4-bit degradation), conclude with evidence and document mitigations (use the QAT checkpoint, or an 8-bit / higher-precision-head quantization) rather than shipping a speculative change.
Acceptance criteria
- Either: a bounded, verified code/quantization change (e.g. higher-precision logit head by default, or a documented mixed-precision option) that makes the standard-4bit repro emit "Middle Ages" (matching QAT / higher precision) without unacceptable memory/perf regression and with no regression on other models; the 5-paragraph repro is clean and deterministic.
- Or: an evidence-backed determination that this is inherent standard-4-bit RTN degradation, with a documented mitigation (QAT / 8-bit / higher-precision head) and the logit-margin data captured on the issue.
Priority / impact
Low. Deterministic single wrong word (not garbage, not non-deterministic) on the standard-4-bit checkpoint only; a working mitigation (the QAT variant) already exists. Much milder than the pre-#910 corruption.
Environment / related
- GB10 (CUDA sm_121), mlxcel HEAD (post-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 #910
4fd625b),gemma-4-26b-a4b-it-4bitvsgemma-4-26b-a4b-it-qat-4bit. - Related: 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 #910 (qmm race, fixed, made this deterministic and isolable), bug: Gemma 4 26B A4B multi-turn corruption from numerically unfaithful fused decode-MoE kernel #886 (fused MoE f32). This is the residual quality item split out from the 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 #910 verification.
Activity
inureyes commentedon Jul 27, 2026
Investigation result: no bounded compute-precision fix; inherent standard-4-bit degradation (can-not-fix)
Confirmed on current main (4fd625b, post-#910) on GB10. The 5-paragraph repro deterministically emits "During the Middle War" on
gemma-4-26b-a4b-it-4bit, while the bread paragraph in the same run correctly emits "During the Middle Ages". This is standard-4-bit (RTN/affine) quantization quality, not a runtime bug, and the prime hypothesis (that the logit-head accumulation precision is the lever, analogous to #886) is disproven by direct measurement.Where the argmax is decided and what the margin is
The failing token is the projection
hidden(2816) -> vocab(262144)at the wine "During the Middle ___" decode step. Instrumenting the tied-embedding logit head to compute BOTH the current bf16 head and an f32-accumulated head from the IDENTICAL bf16 hidden state at that exact step (token ids: " War"=5289, " Ages"=52901; values are recovered pre-softcap logits, cap=30):The bf16-vs-f32 accumulation difference on these logits is about 0.045 (War) and 0.002 (Ages), roughly 100x smaller than the ~4.0-logit margin, and it does not change the argmax. For contrast, at the bread "Middle Ages" step the model is confident (Ages raw ~+21.6, next candidate " Age" ~+1.35, " War" ~-5.5); the wine long-context regime is where every candidate collapses into a low clustered band (top-1 raw only about -18.5) and " War" edges out " Ages".
Enabling the f32-accumulated logit head for the entire run (which perturbs the whole greedy trajectory via the butterfly effect) still emits "During the Middle War".
The accumulation site (for the record)
The CUDA qmv kernel (
src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmv.cu) does accumulate the per-thread K-partials in the activation dtype forbits < 8(conditional_t<(bits >= 8), float, T>), i.e. bf16 for this 4-bit head, whilebits >= 8accumulates in f32. So there is a genuine bf16-accumulation site, and forcing it to f32 is a real change (note: MLX's CUDAquantized_matmuloutput dtype tracks the scales/biases dtype, so an honest f32 test must upcast scales+biases, not just the activation). But the measured effect of that change at the failing decision is negligible relative to the margin, so it is not a fix. The ~4-logit gap is baked into the bf16 hidden state and the 4-bit weights before the projection; f32-accumulating the final projection moves it by <0.05.Conclusion and mitigations
This is inherent standard-4-bit (RTN/affine, group_size 64) weight degradation surfacing only in the long accumulated multi-turn context, not a bounded compute/accumulation-precision defect. No speculative code change is being shipped (same posture as #817 / #824).
Confirmed mitigations:
gemma-4-26b-a4b-it-qat-4bit(quantization-aware-trained): does not exhibit the flip on the identical repro (only correct "Middle Ages" / "Middle East").Priority stays low: deterministic single wrong word (not garbage, not non-deterministic) on the standard-4-bit checkpoint only, with a working mitigation already available. Reopen if a product decision is made to change defaults (e.g. a higher-precision logit head by default, or f32 accumulation across the whole 4-bit decode path), but neither is a bounded logit-head change and the QAT mitigation already covers the case.
inureyes commentedon Jul 27, 2026
Additional data point from a release-regression sweep (post-#910, binary 4fd625b): the stray-token artifact is not limited to the wine/"Middle War" context. On a simple prompt ("List three interesting facts about Mars..."), standard-4bit gemma-4-26b-a4b deterministically emits "...from New York City to Los जाते Los Angeles" (a stray Devanagari token) while the rest of the output is fully correct. It is deterministic (byte-identical across runs, matches the CUDA_LAUNCH_BLOCKING=1 oracle), and the QAT variant (gemma-4-26b-a4b-it-qat-4bit) is clean on the same prompt. This confirms the wontfix diagnosis: inherent standard-4bit RTN degradation of this checkpoint (occasional deterministic stray tokens across contexts), mitigation is the QAT checkpoint. In the same sweep, 13 other models (dense/MoE/VLM/bf16/mxfp4, 4-bit incl. mixtral, qwen3-30b-a3b, gpt-oss-20b) were clean and deterministic, so this is isolated to the standard-4bit gemma-4-26b-a4b weights.