Repository navigation
fix(phi,stablelm): port the upstream f32 attention guard #1710
Description
Upstream promotes the attention inputs to f32 in phi, phixtral and stablelm; mlxcel had ported that guard only into phixtral.rs. phi.rs and stablelm.rs ran their scores at the input dtype.
Without it phi-2-4bit goes NaN at layer 29 of 32 in f16, and every token after that is the argmax of a NaN row, which decodes as !. It is reproducible in bare MLX with no mlxcel in the loop: the same forward with a f32 residual stays finite through all 32 layers and picks ' a' for "Paris is", and with a f16 residual it goes NaN. phixtral.rs's own comment already records the same failure ("NaN from layer 30 of 32"), so the family had been diagnosed once and only one of its files was fixed.
Nothing caught it because a second defect was masking it. gelu_approx in the cxx bridge returned f32 for a half-precision input, which widened the whole residual stream from the first MLP onward and kept the scores out of f16 range by accident. Fixing that exposed this.
Upstream reference: https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/models/phi.py and https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/models/stablelm.py, queries.astype(mx.float32) into the attention call and .astype(values.dtype) out of it; stablelm widens the keys as well.
The guard costs throughput, and the honest accounting is that phi-2-4bit's faster number was the speed of a broken model. On M1 Ultra it reads 164.10 tok/s with the widened attention removed and 129.93 with it, against 128.55 before any of this work. stablelm-1.6b-4bit goes 245.40 to 209.32, which is 98% of mlx-lm on the same host, the reference paying the same cost; its previous 115% was the guard not being paid. Measured headroom on stablelm is 11.8x against f16's ceiling on a 16-token prompt, worst |q·k^T| 5534.5, so its guard is precautionary. Phi's is not.
Two follow-ups. The other families upstream guards this way should be checked against their mlxcel ports: mlx-vlm applies it in minimax_m3_vl, paligemma, phi, phixtral and nemotron_labs_diffusion. And a f16 range assertion would catch the next one, since the symptom is a silent NaN and every parity check passes right up to the layer that overflows.
Update the affected rows in benchmarks/metal_m1ultra_2026-09-06.csv and the M5 Max equivalent when this lands, and regenerate the per-hardware docs. phi-2-4bit and stablelm-1.6b-4bit both move.
Activity
inureyes commentedon Sep 8, 2026
Implemented in commit f4ecc92 on branch
bench/0.7.0-refresh(PR #1617, which is not for merge), awaiting the benchmark refresh on both hosts before it lands. Bothphi-2-4bitandstablelm-1.6b-4bitmove, so their rows are re-measured as part of that refresh.bench: refresh the rows the activation dtype fix moved
bench: re-measure the models the dtype fix touched
inureyes commentedon Sep 12, 2026
Status, so this does not read as untouched work: the guard itself landed in
3576d734(src/models/phi.rsandsrc/models/stablelm.rsnow widen exactly what upstream widens), and the M1 Ultra rows were re-measured after it inbenchmarks/metal_m1ultra_2026-09-08.csv:phi-2-4bit128.76 andstablelm-1.6b-4bit211.05, against 128.55 and 245.39 before. Both move in the direction this issue predicted, and the oldstablelmnumber was the speed of a model running its scores in f16.What is left is the M5 Max half.
benchmarks/metal_m5max_2026-09-03.csvstill holds the pre-fix rows (phi-2-4bit106.08 with the early-EOS flag,stablelm-1.6b-4bit428.36), anddocs/benchmark_results/model_tests_m5max.mdstill quotes them. That measurement needs the M5 Max host and is now with the session that owns it; it will also record whether the early-EOS flag onphi-2-4bitsurvives the guard.The two follow-ups named in the description, checking the other families upstream guards this way and adding an f16 range assertion, have no issue of their own yet.
inureyes commentedon Sep 12, 2026
Correcting my previous comment: the M5 Max measurement already existed, and I missed it.
benchmarks/metal_m5max_2026-09-06.csvcarriesphi-2-hf-4bit-mlxat 173.07 andstablelm-2-1_6b-chat-4bitat 283.30, both dated 2026-09-08 and both stampedmlxcel_commit=3576d734, which is the commit that added the guard. So the M5 Max numbers are post-guard, and against the M1 Ultra values inmetal_m1ultra_2026-09-08.csvthey read 1.34x on both models.I looked for the M1 Ultra directory names (
phi-2-4bit,stablelm-1.6b-4bit), and the M5 Max host stores those checkpoints under different directory names, so my grep returned nothing and I read that as "not measured". A directory name is a host-local artifact, whichdocs/benchmarks.mdsays in as many words. Joining across hosts has to go through the checkpoint identity indocs/model-catalog.tsv, not the directory name.What was actually stale is
docs/benchmark_results/model_tests_m5max.md, which still quoted the pre-guard 181.80 and 394.89. #1827 aligns it with the committed rows and also records that the old early-EOS flag on phi-2 does not reproduce: measured with--no-ignore-eosat39c597b0, it generates 14 tokens on the old 7-token chat prompt and the full 128 at pp512, so that flag was the pre-guard NaN behaviour rather than an EOS defect.docs(bench): carry the #1710 guard into the M5 Max tables (#1827)
inureyes commentedon Sep 12, 2026
Done on both hosts. Closing.
The guard landed in
3576d734:src/models/phi.rsandsrc/models/stablelm.rsnow widen exactly what upstream widens, queries into the attention call and the result back out, with stablelm also widening the keys.Measurements, all taken at a commit that contains the guard. M1 Ultra, in
benchmarks/metal_m1ultra_2026-09-08.csv:phi-2-4bit128.76 against 128.55 before,stablelm-1.6b-4bit211.05 against 245.39. M5 Max, inbenchmarks/metal_m5max_2026-09-06.csvunder that host's directory names:phi-2-hf-4bit-mlx173.07 andstablelm-2-1_6b-chat-4bit283.30, both stampedmlxcel_commit=3576d734, which is 1.34x the M1 Ultra figure on both models. #1827 carried those numbers intodocs/benchmark_results/model_tests_m5max.md, which had still been quoting the pre-guard 181.80 and 394.89.The early-EOS flag this issue's numbers carried for
phi-2-4bitdoes not survive the guard: re-measured with--no-ignore-eos, the checkpoint generates 14 tokens on the old 7-token chat prompt and the full 128 at pp512, so that flag was the pre-guard NaN behaviour rather than an EOS defect. The M5 Max table now records the model as passing.Two things named in the description are not done and have no issue of their own yet: checking the other families upstream guards this way (
minimax_m3_vl,paligemma,phi,phixtral,nemotron_labs_diffusionin mlx-vlm), and adding an f16 range assertion so the next silent NaN is caught by a test rather than by a decode that prints!. Separately, a re-measurement ofphi-2on currentmainis reading lower than the committed row; that is being repeated on a quiet host before anyone calls it a regression, and it will get its own issue if it holds up, since the MLX pin also moved in between.inureyes commentedon Sep 12, 2026
Closing the loose end the previous comment left open: the lower
phi-2reading on currentmaindoes not hold up, so no regression issue is needed.Repeated on the M5 Max host under an idle gate,
phi-2-hf-4bit-mlxmeasures 173.23 tok/s at pp512 with prefill 6843.44, against the committed 173.07 and 6802.24. The immediately preceding repetition of the same script read 117.56, and thestablelmpair in the same session read 284.65 and 198.02. The spread was contention: mlxcel is developed from several checkouts on that host, andcargo clippy --features metal,accelerateandcargo test --workspacewere running in another one during the low samples. The MLX pin hypothesis closes with it, since 173.23 was measured at pin81ba1c6awhile the committed row is at9a795735.One thing worth recording for the next benchmark campaign. The contention hit
phi-2and missedstablelminside the same window, so a control model reproducing its baseline does not certify the measurement taken beside it. The suspect model has to be repeated itself.inureyes commentedon Sep 12, 2026
One correction to my closing comment, because the wrong version of it would send the next reader at the wrong commit.
I wrote that the guard landed in
3576d734. It landed inf4ecc9269, the #1709 commit, in a wider form than upstream:phi.rswidenedq,cache_kandcache_vthrough anf32_dtypebinding.3576d734then narrowed it to what upstream actually widens, queries only for phi and queries plus keys for stablelm.That distinction matters for anyone writing the regression test in #1830: reverting
3576d734restores the wider guard and leaves the model finite, so a test would still pass. The change that removes the protection is deleting the singleastype(&q, mlxcel_core::dtype::FLOAT32)atsrc/models/phi.rs:199.docs(models): audit the f32 attention widen, narrow phixtral
test(core): assert attention scores stay inside the f16 range
test(core): assert attention scores stay inside the f16 range (#1882)
docs(models): audit the f32 attention widen, narrow phixtral
docs(models): audit the f32 attention widen, narrow phixtral (#1890)