test(core): assert attention scores stay inside the f16 range - #1882
Conversation
|
The errors are all This branch adds one |
`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
e553478 to
cea0f64
Compare
tests/attention_f16_range.rsdecodes 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_scoreandF16_MAXinlayers.rsmake 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 undersrc/models/or ingenerate.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 atsrc/models/phi.rs:199: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 statusclean 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_alldoes 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.rsreverted toorigin/main, the two release binaries have byte-identical__textsections; 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 warningsandcargo fmt --all -- --checkclean.Closes #1830