Repository navigation
test(harness): pin f32 GEMMs to full precision, and accept minijinja's Python-style bools - #1260
Merged
Merged
Conversation
…s Python-style bools Eighteen tests were red on Apple GPU generation 17 hosts and nowhere else, for two unrelated reasons (#1259). Seventeen assert full-f32 algorithm equivalence (chunked vs sequential, prefill vs the single-token chain, absorbed vs decompressed), and MLX routes f32 GEMMs to its reduced-precision NAX kernel whenever MLX_ENABLE_TF32 is unset: the dispatch is 'is_nax_available() && (enable_tf32() || dtype != float32)' with the env defaulting to 1, so on generation 17 the two computation shapes stop agreeing to the pinned tolerances (measured 2.2e-4 against a 1e-4 bound on the MLA up projection). A ctor in each lib test binary now sets MLX_ENABLE_TF32=0 before MLX latches the value, with an explicit operator setting winning; verified that all seventeen pass under the pin and that the override still reproduces the red. Shipped numerics stay on MLX defaults and are covered by the runtime exactness probes. The same mechanism is the open CUDA report #1088; this pin applies to those test binaries too but has only been validated on Metal. The eighteenth expected 'ARG:recursive=true;' from a rendered chat template. minijinja 2.24 (bumped in #1162) renders a Jinja2 bool Python-style, as HF transformers does, so the expectation moves to 'ARG:recursive=True;'; the test's actual subject, the bool surviving argument-mapping normalization as a bool, is unchanged. Both suites are green on this M5 Max with no env vars set: mlxcel-core 1513 passed, root lib 5711 passed. Closes #1259 Refs #1088
inureyes
added
status:review
This was referenced Aug 20, 2026
inureyes
added a commit
that referenced
this pull request
Aug 20, 2026
…on (#1262) ## Summary #1088's second acceptance criterion asked that whichever fix was chosen be documented at the gate's own definition, so the next person does not rediscover the mechanism from a confusing failure set. #1260 fixed the mechanism but documented it only in the ctor doc comments inside `src/lib.rs` and `src/lib/mlxcel-core/src/lib.rs`. Somebody starting from the `Makefile` and asking "why is this gate green now, and what am I allowed to change" finds nothing. This adds the missing note as a third bullet in the `verify-test-cuda` comment block. Comment-only, no recipe change. ## Related issues Refs #1088. Refs #1259, #1260 (the pin this documents). ## Type of change - [x] `docs` ## What the note records Four things a future reader needs, in the place they will actually look: - Neither gate setting `MLX_ENABLE_TF32` is deliberate, not an oversight. The pin lives in the test binaries so it covers every gate running them without the macOS and CUDA command lines diverging from each other. - MLX defaults the variable to 1 in `mlx/utils.h`, which routes f32 GEMMs through a reduced-precision kernel on both backends: TF32 on CUDA, NAX on Apple GPU generation 17. That is why #1088 and #1259 are the same bug on different hardware. - The encoded policy: algorithm-equivalence unit tests run at full f32, while shipped numerics stay on MLX defaults and are covered by the runtime exactness probes (#1188, #1189). - The trap to avoid: do not answer a future red here by widening a parity tolerance before checking whether the run had TF32 on. ## Validation The evidence cited in the note is from the CUDA-runner validation of #1088, run on GB10 (sm_121, Linux aarch64, driver 580.173.02) at `670512c2` with MLX pin `9a795735`: - `make verify-test-cuda`, unmodified: **8167 passed, 0 failed, 310 ignored** across 101 test binaries. No `FAILED`, no `panicked`, no `SIGABRT`/`SIGSEGV`, no compile errors. - `MLX_ENABLE_TF32=1` at the same commit still reds three of the four tests #1088 listed (bailing GLA, deepseek_v2 MLA, florence2), confirming the pin is load-bearing rather than decorative. florence2 reproduces at `max abs diff 0.00014263391`, identical to the value recorded in #1088 at `0f727572`. - klear is the one that no longer reproduces under TF32; its margin was the thinnest of the four (1.1e-3 measured against a 1e-3 tolerance). It was checked instead by deliberately perturbing `src/models/klear.rs:617` to skip the causal kernel, which fails the test by 0.28, roughly 280x its tolerance. The perturbation was reverted. ## Test plan - [x] `make -n verify-test-cuda` emits the same recipe as before the change - [x] Comment-only diff, no Rust sources touched, so no `fmt`/`clippy` surface
inureyes
added a commit
that referenced
this pull request
Aug 20, 2026
## What was actually wrong The run-to-run variance is a **test fixture defect**, not a backend property. `filled_weights` walks a `WeightMap` and advances one LCG seed per key. `WeightMap` is `HashMap<String, UniquePtr<MlxArray>>`, and Rust's `RandomState` randomizes iteration order per process, so every run handed a different noise block to a different tensor: **each process built a different random model.** The klear causality test then measured a different model's numerics each time, which is why its max logit delta wandered over 5.9e-4 to 2.0e-3 and straddled its own 1e-3 bound. The fix is to sort the key walk. The same fixture bug was present in three sibling files, so all four are fixed here. ## Two hypotheses in the issue are refuted, with measurements **"The prime suspect is accumulation order in the quantized MoE path (GatherQMM / grouped GEMM)."** The fixture is not quantized. `synthetic_weights` emits only `.weight` for the expert stacks and no `.scales`, and `SwitchLinear::from_stacked_parts` selects the non-quantized `Regular` path when scales are absent. Probed directly through the public API on GB10: `SwitchGLU::forward_fused_kernel(..).is_some()` is **false**, so neither the fused decode-MoE kernel nor the quantized `gather_qmm` path runs in this test at all. **"The variance exists at both settings, so TF32 is an amplifier and not the cause."** Both halves needed correcting. Once the fixture is deterministic there is no variance at either setting, and the pin's effect on this test is exactly what it is on the other three in #1088. Measured on GB10 (CUDA sm_121, MLX pin `9a795735`), 10 processes per mode after the fix, max abs logit delta between the single-token forward and row 0 of the 4-token prefill: | | `MLX_ENABLE_TF32=0` (the #1260 pin) | `MLX_ENABLE_TF32=1` (MLX default) | | --- | --- | --- | | delta | 5.364418e-7 | 1.0485351e-3 | | distinct values over 10 runs | **1** | **1** | | verdict over 10 runs | 10 passed, 0 failed | 0 passed, 10 failed | Before the fix the same test was 8 failures in 10 with a moving value. It is now bit-identical run to run in both modes, so klear stops being the exception among the four tests #1088 listed and behaves like the other three: a deterministic magnitude that clears the bound under the pin and misses it without. ## The backend is not implicated on this path Repeated identical forwards inside one process are bit-identical at both precision settings, measured separately for each arm (`0e0` for both the single-token and the 4-token prefill arm, at `MLX_ENABLE_TF32=0` and `=1`). Scoped claim on purpose: this says nothing about the quantized MoE accumulation-order nondeterminism recorded in #629 / PR #726, which is a different path that this fixture never reaches. ## What the tolerance is The 1e-3 bound is a precision bound, not a causality bound, and the test now says so. Measured in the same binary by handing the prefill an all-zero additive mask, a genuinely bidirectional prefill moves row 0 by about **1.6**, roughly 1.5e3x the default-precision delta and 3e6x the pinned one. The assertion keeps its full power over the property it names, so the tolerance is left alone rather than widened. ## Acceptance criteria - [x] The source of the run-to-run variance is identified: per-process `HashMap` iteration order in `filled_weights`. - [x] A decision is recorded: it was neither a backend defect nor an accepted backend property, but a fixture defect, and it is fixed. - [x] Nothing to document for users: the premise (nondeterministic greedy decode for quantized MoE on CUDA, inferred from this test) is refuted, so no such note is added. The corrected reasoning is recorded at the `verify-test-cuda` gate definition in the `Makefile`, replacing the claims #1264 left there. - [x] The runtime exactness probes (#1188, #1189) are checked. See below. - [x] Re-running the klear test 10 times on GB10 gives a stated, reproducible verdict. ## On the exactness probes The probes compare a verify block against a single-token chain byte for byte, within one process, so run-to-run variance would make their verdict noise-determined. The in-process repeat measurement above shows forwards on this path are bit-reproducible, so the probes are robust on that axis. One asymmetry surfaced while checking and is recorded rather than changed here, since it is not a nondeterminism problem: the Qwen 3.5 gate short-circuits on `!mlxcel_core::metal_is_available()` in both the server path (`models/qwen3_5.rs`) and the offline CLI path (`commands/generate.rs`), while the three Gemma 4 arms have no such guard and therefore do run the probe on CUDA. Not exercised against a real Gemma 4 checkpoint on CUDA here. ## Verification - `models::klear` 22 passed, `models::afmoe` 23 passed, `models::phixtral` 18 passed, `models::bailing_moe_linear` 35 passed; 98 tests, 0 failed. - `cargo fmt --all -- --check` clean; `cargo clippy --lib --tests --features cuda -- -D warnings` clean. - The klear causality test run 10x per precision mode as tabulated above. Closes #1265
This was referenced Oct 6, 2026
inureyes
added a commit
that referenced
this pull request
Oct 6, 2026
Issue #1065 was filed as an M5 precision defect. The cause is MLX's default MLX_ENABLE_TF32=1, which runs f32 GEMMs with 10-bit mantissa inputs on CUDA Ampere and later (cuBLAS FAST_TF32) as well as on Apple GPU generation 17 (NAX). #1260 already pins the variable to 0 in both lib test binaries. This change closes what was left. - Add tf32_pin_tests, an f32 GEMM (16x512x16, off MLX's gemv path) checked against an f64 host reference. On GB10 it passes with the pin and fails at 5.8e-3 against a 1e-4 bound under MLX_ENABLE_TF32=1. Before this, nothing in mlxcel-core went red on CUDA when the pin was lost: the four tests #1065 named stay green there under TF32 because their shapes use gemv and SDPA, not cuBLAS. - Correct both ctor doc comments, which said the breakage was Apple-only. - Print the measured divergence in the two remaining chunked-SDPA assertions. - Document MLX_ENABLE_TF32 in docs/environment-variables.md with the GB10 decode A/B, and stop naming the M5 attention kill switch as #1065's suspect. Decode A/B on GB10, greedy, 256 tokens: the f32 checkpoint plamo-2-1b changed output on 1 of 5 prompts, deterministically in both arms over 5 runs, at unchanged throughput. bf16, f16 and 4-bit checkpoints showed no difference attributable to the variable. Shipped numerics stay on the MLX default. Closes #1065
inureyes
added a commit
that referenced
this pull request
Oct 6, 2026
## Summary #1065's title and "Hardware dependence" section are wrong: the failures are not M5-specific. The cause is MLX's default `MLX_ENABLE_TF32=1`, which runs f32 GEMMs with 10-bit mantissa inputs on CUDA Ampere+ (cuBLAS `CUBLAS_COMPUTE_32F_FAST_TF32`) and Apple GPU generation 17 (NAX). The M1 Ultra passed only because it has no TF32-class path. The main fix already landed: #1260 pins `MLX_ENABLE_TF32=0` in both lib test binaries, #1262/#1263 document it at the `verify-test-cuda` gate, #1196 added the attention kill switch (which does not affect these tests). This PR closes what remained: - `mlxcel-core/src/tf32_pin_tests.rs`: an f32 GEMM (16x512x16, off MLX's gemv path) against an f64 host reference. Nothing in `mlxcel-core` went red on CUDA when the pin was lost, because the four tests #1065 named pass on CUDA even under TF32 (their shapes use gemv/SDPA, not cuBLAS). - Ctor doc comments in `src/lib.rs` and `mlxcel-core/src/lib.rs` no longer say "Apple only". - The two remaining chunked-SDPA assertions print the measured divergence (the two #1065 named already did). - `docs/environment-variables.md`: new `MLX_ENABLE_TF32` row with the measured decode impact; the `MLXCEL_METAL4_ATTENTION` row no longer names the route as #1065's suspect. Tolerances are unchanged. Shipped numerics stay on the MLX default, per the existing policy (unit tests verify full-f32 algorithm equivalence; TF32 output is not f32 evidence). ## Decision on production decode GB10, greedy, 256 tokens, default vs `MLX_ENABLE_TF32=0`, arms alternated: | checkpoint | weights | result | |---|---|---| | plamo-2-1b | f32 | 1 of 5 prompts diverges (after ~40 words), deterministic in both arms, 5/5 runs each; 38 tok/s both arms | | llama-3.1-8b-bf16 | f16 | identical, 3 prompts x 2 runs | | qwen2.5-0.5b-bf16 | bf16 | run-to-run variation in both arms alike, not attributable | | qwen3-0.6b | 4-bit | run-to-run variation in both arms alike, not attributable | TF32 changes greedy output only where dense f32 GEMMs exist, which in practice means f32 checkpoints. On CUDA, MLX sends unbiased single-row matmuls to a gemv kernel the setting does not reach, so prefill-sized GEMMs are the likely site; that split is read from MLX source, not measured. The default stays; the docs tell operators to set `MLX_ENABLE_TF32=0` for f32-exact output. ## Test plan - [x] At origin/main `5211f58d`, CUDA, `--test-threads=1`: the 4 `mlxcel-core` tests from #1065 pass by default and under `MLX_ENABLE_TF32=1` (3 runs each). The 3 root-lib tests from the 2026-08-17 comment pass by default and fail 3/3 under `MLX_ENABLE_TF32=1` (deepseek_v2 2.70e-4, florence2 1.43e-4) - [x] New sentinel: pass by default and with `=0`, fail 3/3 under `=1` (5.84e-3 vs 1e-4 bound) - [x] After the change: `layers::tests::chunked_`, `mla::decode::decode_tests`, `tf32_pin_tests` all pass on CUDA - [x] `cargo fmt --check`, `cargo clippy -p mlxcel-core --lib --tests --features cuda -- -D warnings` ## Not verified on this host - The four tests on M5 Max, and the M1 Ultra nightly runner: no Metal hardware here. #1065's 2026-08-17 comment already showed all four pass on M5 under `MLX_ENABLE_TF32=0`, and #1260 reported the suite green on M5 with the pin. - Sentinel failing under `MLX_ENABLE_TF32=1` on M5 (expected from the NAX dispatch, not run). - M5 vs M1 decode quality on real checkpoints. Closes #1065
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Eighteen tests were red on Apple GPU generation 17 hosts and nowhere else, and every recent PR test plan on this machine has been diffing its failure set against
mainto prove it added nothing. Both causes are identified in #1259; this PR fixes them, and both suites run green on an M5 Max with no env vars set (mlxcel-core 1513 passed, root lib 5711 passed).Related issues
Closes #1259. Refs #1088 (same mechanism on the CUDA gate; see below).
Type of change
testCause 1: MLX runs f32 GEMMs at TF32-class precision on NAX hardware (17 tests)
MLX's matmul dispatch is
is_nax_available() && (enable_tf32() || dtype != float32), andMLX_ENABLE_TF32defaults to 1 inmlx/utils.h. On generation 17 that routes f32 GEMMs through the reduced-precision NAX kernel, so the suite's algorithm-equivalence tests (chunked vs sequential GLA, prefill vs the single-token chain, absorbed vs decompressed MLA, incremental vs full decode) stop agreeing to their pinned tolerances: measured 2.2e-4 against a 1e-4 bound on the MLA up projection, 1.4e-3 against 1e-3 on a klear position-0 logit.MLXCEL_METAL4_ATTENTION=0changes nothing, because the mechanism is the matmul dispatch, not the attention route.The fix is a
ctorin each lib test binary that setsMLX_ENABLE_TF32=0before MLX latches the value into its process-wide static, with an explicit operator setting winning over the pin. The policy this encodes: unit tests verify algorithm equivalence at full f32; shipped numerics stay on MLX defaults and are covered by the runtime exactness probes (#1189, #1188).Verified three ways: all 17 pass individually under
MLX_ENABLE_TF32=0before the change; both full suites pass after it; andMLX_ENABLE_TF32=1still reproduces the red, so the override is real.Cause 2: minijinja 2.24 renders Jinja2 bools Python-style (1 test)
prepare_chat_request_renders_tool_call_arguments_as_mappingpinnedARG:recursive=true;and the #1162 bump to minijinja 2.24 rendersARG:recursive=True;, which is what HFtransformers(Python Jinja2) produces for the same template. The expectation moves to the Python-style form with the rationale recorded; the test's actual subject, the bool surviving argument-mapping normalization as a bool rather than a string, is unchanged.The CUDA gate (#1088)
The pin lives in the test binaries, so a CUDA build's test run gets it too, and #1088's "the gate leaves MLX TF32 enabled" is the same default this pin overrides. Validation needs the CUDA runner; this PR does not claim it.
Test plan
cargo fmt --checkcargo clippy --workspace --all-targets --features metal,accelerate -- -D warningscargo test --release -p mlxcel-core --features metal,accelerate: 1513 passed, 0 failedcargo test --release -p mlxcel --lib --features metal,accelerate: 5711 passed, 0 failedMLX_ENABLE_TF32=1reproduces the pre-fix failure (operator override wins)