fix(core): f32 matmul and single-query SDPA lose ~fp16 precision on M5, breaking four numeric tests

Problem / Background

Four deterministic test failures occur on Apple M5 Max that do not occur on the M1 Ultra nightly runner. All four are numeric-agreement assertions on f32 computations. The observed error magnitudes cluster at fp16 epsilon scale (2^-10 = 9.77e-4), roughly three orders of magnitude larger than f32 epsilon, which indicates the f32 inputs are being computed at reduced internal precision on M5-class hardware.

The four failing tests

mlxcel-core layers::tests::chunked_causal_attention_matches_fast_causal at src/lib/mlxcel-core/src/layers.rs:6344. Message: q_len=4 k_len=4 chunk=1 diverged from do_causal SDPA. Measured max_abs_diff = 0.0009678602 against a tolerance of 1e-5. Fails only at chunk=1; chunk 2 and 3 pass.
mlxcel-core layers::tests::chunked_query_attention_matches_unchunked_with_mask at src/lib/mlxcel-core/src/layers.rs:6288. Message: chunk=1 diverged from unchunked SDPA. Measured max_abs_diff = 0.0009678602 against a tolerance of 1e-5. Fails only at chunk=1; chunks 2, 4 and 5 pass.
mlxcel-core mla::decode::decode_tests::absorbed_attention_respects_an_additive_causal_mask at src/lib/mlxcel-core/src/mla/decode_tests.rs:93. Message: masked absorbed attention drifted by 0.0014691837, against F32_TOL = 2e-5 (relative error).
mlxcel-core mla::decode::decode_tests::expand_latent_reproduces_the_up_projection at src/lib/mlxcel-core/src/mla/decode_tests.rs:132. Message: k_nope[0][0][1] 0.30357933044433594 != 0.3037950135765861, absolute tolerance 1e-4, observed absolute error 2.157e-4 (relative error about 7.1e-4).

Why this is a precision problem and not a test problem

Test 4 compares the GPU result against a reference dot product accumulated in f64 in plain Rust, over kv_lora_rank = 32 elements (the TINY fixture is num_heads: 4, kv_lora_rank: 32, qk_nope_head_dim: 16, qk_rope_head_dim: 8, v_head_dim: 12). A correctly executed f32 dot product of length 32 on values of order 1 should differ from the f64 reference by about 1e-6 or less. The observed 2.157e-4 is roughly two orders of magnitude larger than that, and sits at fp16 epsilon scale. The reference side is exact, so the GPU side is the imprecise one.
Tests 1 and 2 both report the identical divergence 0.0009678602, which is just below 2^-10 = 0.0009765625, that is fp16 epsilon. Both fail only in the chunk=1 configuration, which is the single-query path.
Test 4 is a plain up-projection matmul, not attention, so the reduced precision is not confined to SDPA kernels.

Hardware dependence, established

Failing host: Apple M5 Max, macOS 26.6, toolchain 1.93.1 (also reproduced under 1.97.1 with byte-identical values).
Passing host: the self-hosted-macos-26-arm64 nightly runner (M1 Ultra). In the last successful nightly-verify.yml run (run id 30939291504) all four tests report ... ok.
The failures are fully deterministic on M5: three consecutive reruns produced byte-identical values.
The repository already carries M5-specific routing at src/lib/mlxcel-core/src/lib.rs:3178, which dispatches to layers::metal4_causal_attention when hw.has_neural_accelerator && hw.macos_supports_na. There is currently no environment variable to disable that path, so the hypothesis could not be isolated by toggling it off. Adding such an override would make this class of problem diagnosable.

Why it matters beyond the test suite

chunk=1 is the single-query shape, which is what every decode step uses in production. If the single-query path on M5 computes at fp16-level precision while the multi-token prefill path computes at f32, then prefill and decode disagree numerically on M5 by about 1e-3 per step. Test 3's failure is on the MLA decode path specifically. This should be assessed for output-quality impact on real checkpoints, not only fixed at the tolerance level.

Suggested investigation

Determine whether MLX selects a neural-accelerator (NAX) kernel for f32 matmul and single-query SDPA on M5, and whether that kernel accumulates at reduced internal precision.
Decide whether the reduced precision is intended behavior on M5. If it is, the affected tolerances need per-hardware calibration rather than a blanket loosening, and the divergence must be documented. If it is not, the dispatch needs to select a full-precision kernel for these shapes.
Measure decode output quality on at least two real checkpoints on M5 versus M1 to establish whether the numeric difference is observable in generated tokens.
Add an environment override to disable the M5 neural-accelerator routing so this class of divergence can be bisected locally.

Related test-quality gap (fold into the fix)

The assertions in tests 1 and 2 do not print the measured difference; the message says only "diverged". The divergence magnitude had to be obtained by patching the assertion locally. Those two assertions should include the measured value, the way tests 3 and 4 already do.

Not part of this issue

A fifth test, tests/mllama_parity.rs::sub_max_real_tiles_keep_the_legacy_real_rows_byte_identical, also fails on M5 but has a different and benign root cause: it asserts exact f32 bit equality and the observed difference is 5.9604645e-8 (2^-24), which is 0.5 ULP of the output's largest element (1.1521907). That is the same class as issue #939 and is being fixed separately by applying a tolerance, following the precedent set by PR #953 for its sibling test in src/models/mllama/text.rs.

Acceptance Criteria

Root cause established: confirmed whether MLX dispatches a reduced-precision (neural-accelerator) kernel for f32 matmul and single-query SDPA on M5, with evidence.
A decision is recorded on whether the reduced precision is intended M5 behavior, along with the rationale.
If unintended: the dispatch selects a full-precision kernel for these shapes, and all four tests pass on M5 with their original tolerances.
If intended: tolerances are calibrated per hardware (not blanket-loosened), the M5 divergence is documented in the repository docs, and all four tests pass on both M5 and the M1 Ultra nightly runner.
Decode output quality measured on at least two real checkpoints on M5 versus M1, with results reported in this issue.
An environment override exists to disable the M5 neural-accelerator routing at src/lib/mlxcel-core/src/lib.rs:3178, and it is documented.
The assertions in chunked_causal_attention_matches_fast_causal and chunked_query_attention_matches_unchunked_with_mask print the measured divergence in their failure messages.
The fix is integrated into the actual inference code path, not left as a standalone helper, and cargo test --release passes on both M5 and the nightly runner.

Technical Considerations

Affected files: src/lib/mlxcel-core/src/layers.rs, src/lib/mlxcel-core/src/mla/decode.rs and its tests, src/lib/mlxcel-core/src/lib.rs (M5 routing at line 3178).
Per the project's Apple Silicon precision rules, M5-specific code paths gate on hw.has_neural_accelerator && hw.macos_supports_na. Any new override must respect that gate rather than bypassing hardware detection.
Loosening tolerances alone would hide a real production-path divergence on the decode hot path, so it is not an acceptable standalone fix.

fix(tests): tolerate f32 reassociation in the mllama tile-selection p…

Narrowing from the #1186 investigation, plus the override this issue asked for

Two contributions from the block-vs-chain work on #1186 and #1188. The first narrows the suspect list; the second removes the obstacle that made testing it awkward.

The override now exists: PR #1196

Acceptance criterion 6 here asks for an environment override that disables the M5 neural-accelerator routing so this class of divergence can be bisected locally. There was none, and the M5 investigation on #1182 had to patch should_use_metal4_attention by hand and revert it afterwards.

MLXCEL_METAL4_ATTENTION=0 now forces the route off. Off-switch only, so it is inert on hardware that lacks the accelerator. It also fixes a drift the change exposed: lib.rs's causal-attention dispatch carried its own inline copy of the hardware test rather than calling the shared predicate, so a switch applied to one site only would have produced runs that look like the route is off and are not.

Suggested use here: run the four tests this issue names with and without the variable and report whether the failures move. That is the evidence acceptance criterion 1 wants.

What the new diagnostic does and does not say about the SDPA half

tests/metal_block_vs_chain_op_parity.rs (in main since #1192) compares one T = K call against K single-position calls, byte for byte. On both an M1 Ultra and an M5 Max, at Qwen3.8-27B and Gemma 4 12B shapes:

op	M5 Max	M1 Ultra
`fast_rms_norm`	equal	equal
`fast_rope`	equal	equal
causal SDPA, block vs per-position	equal	equal
dense matmul f16	DIVERGES	DIVERGES
dense matmul f32	DIVERGES	DIVERGES
`quantized_matmul` (affine)	DIVERGES	equal below the batch limit

Read this carefully, because it is easy to overclaim. The SDPA row says the single-query form and the T = K block form agree with each other on f16 inputs. This issue's failures are f32 computations compared against a higher-precision reference. Those are different axes: agreement between two paths does not establish that either is accurate, and an f16 result says nothing directly about the f32 kernel.

What it does contribute: on the block-vs-chain axis, the single-query SDPA path is not distinguishable from the multi-position one on either machine, so whatever separates M5 from M1 in tests 1 and 2 is not a T = 1 versus T > 1 dispatch split in SDPA. Test 4 is a plain up-projection matmul rather than attention, and the matmul rows above are where every measured divergence lives, which is consistent with this issue's own observation that the reduced precision is not confined to SDPA kernels.

A second lever worth trying, from the same investigation

MLX gates its NAX gemm on metal::is_nax_available() && !complex && (env::enable_tf32() || a.dtype() != float32) (https://github.com/ml-explore/mlx/blob/main/mlx/backend/metal/matmul.cpp). For float32 specifically that last clause makes MLX_ENABLE_TF32=0 a way to avoid the NAX gemm without touching our code at all, which is directly relevant to this issue's f32 matmul half. Measured on M5 Max during the #1182 investigation: with MLX_ENABLE_TF32=0 the block arm's digest moved while the chain arm's did not, and the divergence shrank from 72% of bytes to 37%, so the toggle demonstrably reaches the kernel.

Combining the two switches gives four arms on one machine: NA attention on/off crossed with TF32 on/off. Tests 1 to 3 should respond to the first, test 4 to the second, if the hypothesis in this issue is right.

Not addressed

Nothing here measures decode output quality on real checkpoints (acceptance criterion 5), and nothing here was run on M5, which is where the four failures are.

feat(core): add MLXCEL_METAL4_ATTENTION as a kill switch for the M5 N…

Root cause: `MLX_ENABLE_TF32`, not the M5 neural-accelerator attention route

Ran the cross this issue's suggested investigation asks for, now that #1196 has shipped the kill switch. Apple M5 Max, 40 GPU cores, 128 GB, macOS 26.6.1, Xcode 26.5, main at 33951285, cargo test --release --features metal,accelerate -p mlxcel-core.

`MLXCEL_METAL4_ATTENTION`	`MLX_ENABLE_TF32`	result
(unset)	(unset)	4 FAILED
1	1	4 FAILED
0	1	4 FAILED
1	0	4 passed
0	0	4 passed

All four tests respond to the second variable only. MLXCEL_METAL4_ATTENTION changes nothing, in either direction, for any of them.

The failure values are byte-identical across all three failing arms and match this issue's body exactly:

q_len=4 k_len=4 chunk=1 diverged from do_causal SDPA by 0.0009678602
chunk=1 diverged from unchunked SDPA by 0.0009678602
masked absorbed attention drifted by 0.0014691837
k_nope[0][0][1] 0.30357933044433594 != 0.3037950135765861

The default-environment arm matters as much as the cross: with neither variable set, all four fail. So TF32 is on by default on this host and the default state is exactly what this issue reports. With MLX_ENABLE_TF32=0 all four pass at their original tolerances, unchanged.

On the first acceptance criterion

This is evidence that MLX dispatches a reduced-precision kernel for f32 matmul on M5 and that the selection is governed by TF32. It is consistent with every observation in the issue body: the fp16-epsilon-scale error, the identical 0.0009678602 in tests 1 and 2, and test 4 being a plain up-projection matmul rather than attention. A single switch explaining an attention test and a non-attention test together is what a matmul-level cause looks like.

What this does not establish, and please do not read it as such

MLXCEL_METAL4_ATTENTION having no effect here is not a verification that the switch works, and it is not evidence that the NA attention route is well behaved. These four tests never reach it.

The route is consulted at exactly three sites: lib.rs:3207, layers::metal4_attention and layers::metal4_causal_attention. The failing tests call attention_dispatch, chunked_causal_attention, chunked_query_attention, ffi_fast_scaled_dot_product_attention_causal and the MLA decode paths directly, none of which pass through those three. Confirmed empirically as well: re-running test 1 with MLXCEL_LOG_NA_ATTENTION=all produces zero dispatch log lines with the switch on and zero with it off.

So this test set is out of the route's call path, the no-effect result is uninformative about the switch, and #1196's "effect not verified" stays open. Following up separately with a real-checkpoint decode A/B that does engage the route.

Suggested next step for the remaining criteria

The second criterion (is the reduced precision intended M5 behavior) is now a narrower question than the issue posed: it is about MLX's TF32 default on generation 17, not about our dispatch. The fourth (decode output quality on real checkpoints, M5 against M1) is unchanged and still needs measuring, and MLX_ENABLE_TF32 gives it a clean A/B on a single host.

Follow-up: TF32 is invisible to decode on quantized checkpoints

Same host and build as the four-arm comment above. This is a partial answer to the fourth acceptance criterion, and it is partial in a way worth being explicit about.

96 tokens, --temp 0, arms alternated, two reps, default environment against MLX_ENABLE_TF32=0:

checkpoint	default rep1	default rep2	TF32 off rep1	TF32 off rep2	generated text
`llama-3.2-1b-4bit`	517.23	519.81	519.96	523.41	identical
`gemma-3-4b-it-4bit`	178.23	178.51	178.01	178.16	identical
`qwen2.5-7b-instruct-4bit`	119.49	119.76	119.45	119.48	identical

Byte-identical output on all three, throughput within noise.

Why this is not yet the criterion's answer

All three are 4-bit checkpoints, so their hot matmuls are quantized_matmul, not f32 dense matmul, and TF32 governs the f32 arm. The M5 probe on #1182 saw this directly: MLX_ENABLE_TF32=0 changed the dense f32 block digest and left the f16 arm alone. So this measurement mostly demonstrates that the production decode path does not go through the ops TF32 controls, rather than that reduced precision is harmless where it does apply.

That is still useful, and it narrows the issue: the four failures are f32 fixtures, and a 4-bit checkpoint's decode does not reach that path. It is an argument that the blast radius is smaller than the issue's "prefill and decode disagree numerically on M5 by about 1e-3 per step" framing suggested, at least for quantized checkpoints.

Not measured, and needed to actually close the criterion: a checkpoint whose forward pass does f32 dense matmul, and the M1 comparison. This is a single-host result.

Also worth flagging for whoever picks this up: --temp 0 byte-identity over 96 tokens is a test for argmax flips, not for per-op precision. A quality assessment that matters would want perplexity or a longer generation, not a diff.

This is not M5-specific. It reproduces on CUDA, from the same cause.

make verify-test-cuda at fa260a5d on NVIDIA hardware fails three tests in -p mlxcel --lib:

models::bailing_moe_linear::tests::chunked_gla_matches_the_sequential_recurrence
models::deepseek_v2::tests::absorbed_mla_attention_matches_the_decompressed_block_step_for_step
models::florence2::florence2_tests::incremental_decode_matches_full_sequence

All three are members of the M5 Max set, and all three pass under MLX_ENABLE_TF32=0 on CUDA, exactly as all seventeen do on M5:

MLX_ENABLE_TF32=0 cargo test --profile test-fast --features cuda -p mlxcel --lib -- --test-threads=1 <the three>
test result: ok. 3 passed; 0 failed

So the variable is MLX_ENABLE_TF32 on both backends. TF32 is on by default for f32 matmul on Ampere and later just as it is on an M5's neural accelerator, and this issue's tests are numeric-agreement assertions on f32.

What that changes about this issue

The title and the "Hardware dependence, established" section are wrong, and in a way that matters. They read as "M5 computes at reduced precision and M1 Ultra does not", which frames it as an Apple generation 17 defect. The M1 Ultra is not the control it looks like: it passes because it has no TF32-equivalent path, not because the other hosts are broken. Two unrelated architectures failing the same assertions under the same switch is a property of the assertions and of MLX's default, not of a chip.

It also changes what a fix would cost. "Calibrate tolerances per hardware" now means per hardware class, at least Metal-NA and CUDA-Ampere-plus, and the set is open. The alternative, deciding whether TF32 should be on by default for these ops, is one decision rather than a growing table.

The subset is itself information

CUDA fails 3 where M5 fails 17. Whatever the other 14 touch on M5, CUDA either does not route through TF32 there or its kernels close the gap. Worth knowing before anyone loosens a tolerance: a tolerance wide enough for M5 would be far wider than CUDA needs on the same test.

Not a regression from the MLX pin bump

Established from both directions rather than assumed. On M5 the failure set is identical at pin 2c46b953 and 9a795735, deterministic over three isolated runs each (#1198). On CUDA the three failures vanish under MLX_ENABLE_TF32=0, which is a property of the ops and not of the pin.

Recommend retitling this issue away from M5 and restating the acceptance criteria around the TF32 default, since "all four tests pass on M5 with their original tolerances" is now a subset of the actual problem in two directions: seventeen tests, two backends.

Close-out in #2132, and a correction to this issue's framing

The title and the "Hardware dependence, established" section are superseded by the 2026-08-17 comments: the cause is MLX's default MLX_ENABLE_TF32=1 (10-bit mantissa f32 GEMM inputs) on CUDA Ampere+ as well as Apple generation 17. It is not M5-specific, and the MLXCEL_METAL4_ATTENTION route is not involved. The body is left as filed; this comment is the current reading.

criterion	status
Root cause	`MLX_ENABLE_TF32`, four-arm cross on M5 (2026-08-17) plus CUDA reproduction
Intended or not	MLX's default, not our dispatch. Tests pin it to 0 (#1260); shipped numerics stay on the default
Tests pass at original tolerances	#1260, no tolerance changed
Env override for the NA route	#1196 (`MLXCEL_METAL4_ATTENTION`)
Assertions print the divergence	already true for the two named; #2132 adds the last two chunked-SDPA ones
Decode quality on real checkpoints	measured on GB10 in #2132, below; M5 vs M1 not run

New from GB10 at 5211f58d: the four mlxcel-core tests named here pass on CUDA even under MLX_ENABLE_TF32=1 (3 runs each), because their shapes use MLX's gemv and SDPA kernels rather than cuBLAS. So nothing in mlxcel-core caught a lost pin on CUDA; #2132 adds a sentinel GEMM test that fails at 5.84e-3 under =1.

Greedy decode, 256 tokens, default vs =0: the f32 checkpoint plamo-2-1b diverges on 1 of 5 prompts, deterministically in both arms (5/5 runs), same throughput. f16, bf16 and 4-bit checkpoints show nothing attributable to the variable. The default stays; docs/environment-variables.md now documents MLX_ENABLE_TF32=0 for f32-exact output.

test(core): guard the TF32 test pin and record its decode impact

Navigation Menu

Repository navigation

fix(core): f32 matmul and single-query SDPA lose ~fp16 precision on M5, breaking four numeric tests #1065

Problem / Background

The four failing tests

Why this is a precision problem and not a test problem

Hardware dependence, established

Why it matters beyond the test suite

Suggested investigation

Related test-quality gap (fold into the fix)

Not part of this issue

Acceptance Criteria

Technical Considerations

Narrowing from the #1186 investigation, plus the override this issue asked for

The override now exists: PR #1196

What the new diagnostic does and does not say about the SDPA half

A second lever worth trying, from the same investigation

Not addressed

Root cause: `MLX_ENABLE_TF32`, not the M5 neural-accelerator attention route

On the first acceptance criterion

What this does not establish, and please do not read it as such

Suggested next step for the remaining criteria

Follow-up: TF32 is invisible to decode on quantized checkpoints

Why this is not yet the criterion's answer

This is not M5-specific. It reproduces on CUDA, from the same cause.

What that changes about this issue

The subset is itself information

Not a regression from the MLX pin bump

Close-out in #2132, and a correction to this issue's framing

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

fix(core): f32 matmul and single-query SDPA lose ~fp16 precision on M5, breaking four numeric tests #1065

Description

Problem / Background

The four failing tests

Why this is a precision problem and not a test problem

Hardware dependence, established

Why it matters beyond the test suite

Suggested investigation

Related test-quality gap (fold into the fix)

Not part of this issue

Acceptance Criteria

Technical Considerations

Activity

inureyes commented on Aug 16, 2026

Narrowing from the #1186 investigation, plus the override this issue asked for

The override now exists: PR #1196

What the new diagnostic does and does not say about the SDPA half

A second lever worth trying, from the same investigation

Not addressed

inureyes commented on Aug 17, 2026

Root cause: MLX_ENABLE_TF32, not the M5 neural-accelerator attention route

On the first acceptance criterion

What this does not establish, and please do not read it as such

Suggested next step for the remaining criteria

inureyes commented on Aug 17, 2026

Follow-up: TF32 is invisible to decode on quantized checkpoints

Why this is not yet the criterion's answer

inureyes commented on Aug 17, 2026

This is not M5-specific. It reproduces on CUDA, from the same cause.

What that changes about this issue

The subset is itself information

Not a regression from the MLX pin bump

inureyes commented on Oct 6, 2026

Close-out in #2132, and a correction to this issue's framing

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

Root cause: `MLX_ENABLE_TF32`, not the M5 neural-accelerator attention route