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

[inureyes]

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

1. 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.
2. 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.
3. 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).
4. 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

1. 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.
2. 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.
3. 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.
4. 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.

[inureyes]

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.

[inureyes]

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.

[inureyes]

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.