Skip to content

test(core): assert attention scores stay inside the f16 range - #1882

Merged
inureyes merged 2 commits into
mainfrom
test/issue-1830-attention-f16-range
Sep 14, 2026
Merged

inureyes merged 2 commits into
mainfrom
test/issue-1830-attention-f16-range

Conversation

@inureyes

Copy link
Copy Markdown
Member

tests/attention_f16_range.rs decodes 32 steps on each f16 family carrying the f32 attention widen (phi, stablelm, phixtral) and checks every step, so a regression in that guard fails the build instead of emitting garbage. max_abs_attention_score and F16_MAX in layers.rs make the headroom measurable: mlxcel hands q, k and v to MLX's fused SDPA and never forms the scores, so anything inspecting them has to recompute them. Both are diagnostics with no caller under src/models/ or in generate.rs.

The discrimination experiment

Required by #1830: a test that cannot fail proves nothing.

phi, widen deleted: RED. Removing the single astype(&q, FLOAT32) line at src/models/phi.rs:199:

phi-2-hf-4bit-mlx: steps 1..=8 all sampled token 0, with max logits
[4.52e-41, 4.52e-41, 4.52e-41, 4.52e-41, 4.52e-41, 4.52e-41, 4.52e-41, 4.52e-41].
The row is finite, so this is a collapsed logit distribution rather than a NaN one;
the f32 widen on this family's attention scores is the guard for it.

stablelm, both widen lines deleted: GREEN. 25 distinct tokens and ordinary text, matching the 11.8x headroom #1710 measured for it. #1830 asks for it either way, and it bounds what this file covers.

Guards restored by path; git status clean after each run.

The failure shape is not the one #1710 described

This is why the file checks two things. Deleting phi's widen produces no NaN row on this prompt: every step's max logit stayed finite at 4.52e-41, a denormal, and argmax returned token 0 at all 33 steps. A finiteness assertion alone would have passed a model emitting nothing but !. max_all does propagate NaN, checked directly against a row containing one, so that assertion is sound for the NaN shape; this prompt simply collapses before reaching it. The collapse check therefore runs inside the decode loop, so its failure names the step where the run began rather than only reporting that the whole generation was degenerate.

Decode throughput

Timing could not settle this and the binaries can. Built at this commit and with layers.rs reverted to origin/main, the two release binaries have byte-identical __text sections; 99 of 53,488,176 bytes differ, all outside the code (load commands and __LINKEDIT), and the new symbol survives into neither binary because nothing calls it. Identical code cannot decode at a different speed.

Timing is reported as a caution. Position-balanced over 32 runs of -n 128 --temp 0, after/before is 0.9799 by median. That 2.0% gap is smaller than the 2.6% effect of which binary ran first within a rep, which was measured by reversing the order and watching the sign flip, and both sit inside a 17.3% spread on a box at load ~3.3. Before the order was reversed the first pass looked like a consistent 4.9% regression in 6 of 7 reps.

Validation

M1 Ultra, macOS 27.0, all three checkpoints resolved under MLXCEL_REQUIRE_MODELS=1. Workspace gate 11247 passed / 0 failed / 124 suites; cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check clean.

Closes #1830

@inureyes inureyes added type:test Test related changes status:review Under review priority:medium Medium priority area:models Model architectures, weights, loading, metadata area:core mlxcel-core: MLX FFI, primitives, KV cache, layers platform:macos macOS (Apple Silicon) specific labels Sep 14, 2026
@inureyes

Copy link
Copy Markdown
Member Author

OpenXLA feature compile is red here and it is not this branch. It fails on main too: red at 7e4577b1 and at b8d10fb1, green at d67b64d3, so it entered with #1871 (7e4577b1, "enable model-free WebUI startup"), which touches 30 Rust files including every file the errors name.

The errors are all unused import and never used under the XLA feature set, in src/server/app.rs, src/server/router_server.rs, src/server/auth.rs and src/models/mod.rs: ui_model_actions, ui_operation_cancel, invalid_webui_field, is_webui_token, request_id and others. Nothing in them names layers.rs, tests/attention_f16_range.rs, max_abs_attention_score or F16_MAX.

This branch adds one pub const, one uncalled pub fn and one test file. The other 11 checks pass and the default-feature cargo clippy --workspace --all-targets -- -D warnings is clean locally.

`phi-2-4bit` ran its `q @ k^T` in f16 for months and nothing in the tree failed. The products pass 65504 in the deep layers, and layers 0 through 29 of 32 track the reference to three decimal places first, so a parity check that samples an early layer passes right up to the layer that breaks. #1709 additionally masked it by widening the residual stream through `gelu_approx`; fixing that exposed it.

`tests/attention_f16_range.rs` decodes 32 steps on each f16 family that carries the widen today (phi, stablelm, phixtral) and checks every step, so a regression in the guard fails the build instead of emitting garbage text. `max_abs_attention_score` and `F16_MAX` in `layers.rs` make the headroom itself measurable: mlxcel hands q, k and v to MLX's fused SDPA and never forms the scores, so anything inspecting them has to recompute them. Both are diagnostics with no caller under `src/models/` or in `generate.rs`, because reading a score costs a host readback per call and that boundary was removed once already for costing 3.31x of decode.

The failure shape is not the one #1710 described, which is why the test checks two things. Deleting phi's widen and running this file: every step's max logit stayed finite at 4.52e-41, a denormal, and argmax returned token 0 at all 33 steps. A finiteness assertion alone would have passed a model emitting nothing but `!`. `max_all` does propagate NaN, checked directly, so that assertion is sound for the NaN shape; this prompt simply collapses before reaching it. The collapse check therefore runs inside the decode loop, so its failure names the step where the run began.

Validated on M1 Ultra with all three checkpoints resolved under `MLXCEL_REQUIRE_MODELS=1`: three tests green with the guards present; phi red with its widen deleted, reporting `steps 1..=8 all sampled token 0`; stablelm still green with both of its widen lines deleted, matching the 11.8x headroom #1710 measured for it. Guards restored by path, tree confirmed clean. Full workspace gate, clippy and fmt pass.

Closes #1830
The signature reads `[B, H, L, D]` and `[B, H, S, D]`, which is accurate and easy to read as permissive. It is not: on a GQA family the head counts differ and the matmul aborts rather than broadcasting. Measuring Gemma 2 for issue #1829, at 8 query heads over 4 kv heads, walked into it.

Refs #1830, #1829
@inureyes
inureyes force-pushed the test/issue-1830-attention-f16-range branch from e553478 to cea0f64 Compare September 14, 2026 12:01
@inureyes inureyes self-assigned this Sep 14, 2026
@inureyes
inureyes merged commit 301618f into main Sep 14, 2026
16 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:core mlxcel-core: MLX FFI, primitives, KV cache, layers area:models Model architectures, weights, loading, metadata platform:macos macOS (Apple Silicon) specific priority:medium Medium priority status:review Under review type:test Test related changes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

test(core): assert attention scores stay inside the f16 range

1 participant