Repository navigation
perf(mistral4): the attention scale promotes the query to f32 #1711
Description
mistral4.rs builds the Llama-4 position-dependent attention scale from a host &[f32], so from_slice_f32 returns f32 and the multiply into q promotes on the wider operand. q leaves that line as f32 and carries through the attention and the output projection into the residual stream, after which every matmul in every later layer promotes its own weight to match. It is #1709 reached through a model file rather than a bridge helper, and the invariant is the same: hand back the dtype you were given.
Fixed in commit 3eb79be on branch bench/0.7.0-refresh (PR #1617, not for merge). ministral3.rs carries the identical line and gets the identical fix; llama4.rs is the original of the pattern and already captures q_dtype before building the scale, so it needed no change.
M1 Ultra, mistral-small-4-119b-2603-4bit, pp512/tg128: decode 18.9 to 55.51 tok/s, which is 105% of mlx-vlm on the same host against 37% before, and effective bandwidth 62 to 184 GB/s. That lands it beside gpt-oss-120b-4bit at 174 GB/s. Prefill is unchanged at 381 tok/s, as expected: prefill measured 112% of mlx-vlm while decode measured 37%, which is the signature of a per-decode-step cost rather than a whole-forward one. Output is unchanged.
The elimination order is worth recording, because all of it came from same-host benchmark data before any code was read. Model size is not the cause: gpt-oss-120b-4bit is the same 120B class and sits at 105%. Expert count and top-k are not the cause: the same 128-expert top-4 configuration appears in models at 103-111%. The fused decode-MoE kernel is not the cause: MLXCEL_FUSED_MOE=0 moves the number 18.97 to 19.22, inside noise. MLA is not the cause: deepseek-v2-lite-4bit 100%, glm4-flash-4bit 105%, minicpm3-4b-4bit 109%. Host-side graph building is not the cause: MLXCEL_PROFILE_PIPELINE_DETAIL puts 1.59 ms in forward against 79.6 ms in async_eval. Only then does reading mistral4.rs pay off.
Two follow-ups. The ministral3.rs change is unmeasured, because no checkpoint on this host reaches that module: ministral-3b-4bit, mistral-small-3.1-24b-4bit and mistral-small-4-119b-2603-4bit all detect as mistral3_vlm, and only the last routes its text decoder through mistral4. And a static audit of from_slice_f32 and hardcoded-FLOAT32 constructors across src/models and src/vision leaves nine candidates inside decode-path functions after excluding tests and sites that restore the dtype; most look like routing gates that are already f32, epsilons and mask construction, but none has been checked against a measurement.
Update the affected rows in benchmarks/metal_m1ultra_2026-09-08.csv and benchmarks/metal_m1ultra_vlm_2026-09-08.csv when this lands, and regenerate the per-hardware docs.
Activity
inureyes commentedon Sep 8, 2026
The fix is implemented in commit 3eb79be on branch
bench/0.7.0-refresh(PR #1617, not for merge). Same root cause as #1709. Awaiting the benchmark refresh before it lands.bench: re-measure mistral-small after the attention-scale fix
inureyes commentedon Sep 8, 2026
The nine remaining audit candidates are resolved: eight are false positives and one is real.
The four
routed_scaling_factorsites (dots1.rs:398,glm4_moe.rs:815,glm4_moe_lite.rs:669,solar_open.rs:638) are character-for-character copies of one block, and all four are safe:moe_weighted_sumcasts the scores to the expert output's dtype immediately before the multiply and casts the sum back tooutput_dtypebefore returning, so a promotedtopk_scoresnever escapes. That is the same boundary restore the compiled bridge variants do, and it is the exact contrast with this issue: inmistral4.rsthe promoted product went on through the attention ando_projinto the residual with no cast at all.llama4_helpers.rs:43widens integer indices to f32 to build a chunked mask, which is not an activation.deepseek_v4_hyper.rs:143adds an epsilon to a sigmoid gate on the routing path.The one real leak is
src/vision/encoders/gemma3n.rs:58: the RMSNorm epsilon is built as f32, sorsqrt(v + eps)is f32 and multiplying it intoxpromotes the activation with no restore. It sits in a vision tower, so it runs once per request rather than once per decoded token, and it is being handled alongside the gemma3n investigation on the M5 Max host.Reproduced on M5 Max:
mistral-small-4-119b-2603-4bittext decode 19.04 to 100.12 tok/s.mistral-small-3.1-24b-4bitandministral-3b-4bitare unchanged there, which independently confirms that onlymistral-small-4routes its text decoder throughmistral4.rs. No checkpoint withmodel_type: ministral3exists on either host, so theministral3.rshalf of the fix stays unmeasured on both.inureyes commentedon Sep 9, 2026
Confirmed on the second host and closing. The fix is on
main:3eb79bee"perf(mistral4,ministral3): cast the attention scale to the query dtype" is an ancestor offea72611, and both halves are in place,mistral4.rs:317andministral3.rs:191.M5 Max,
mistral-small-4-119b-2603-4bit, same host and day,.venv-mlxlmatmlx-vlm-0.6.17:a50ff440)fea72611)That is 5.19x on the text harness and 5.3x on the VLM harness, against 2.91x on M1 Ultra, and it lands at parity with the reference runtime rather than merely improving. Prefill moves too on this host, 902.50 to 1307.07, which the M1 Ultra measurement did not show.
Worth recording because it cost a wrong conclusion first: the M5 Max row had not been re-measured after this landed, so the cross-host table paired M1 Ultra's post-fix 55.03 against M5 Max's pre-fix 19.04 and computed 0.35x, the largest apparent hardware deficit in the set. Three controls measured in the same session (
qwen2.5-7b-instruct-4bit,meta-llama-3.1-8b-instruct-4bit,qwen3-30b-a3b-4bit) reproduce their 2026-09-06 values at 1.00x, 1.00x and 0.98x, which is what separates one stale row from a stale snapshot and avoided re-sweeping the host.The follow-up in the description is not closed by this. The static audit of
from_slice_f32and hardcoded-f32 constructors is being filed separately, because re-running it showed the first pass undercounted: a constructor that takes a dtype captured into a variable, as infull_f32(&[1], alpha, input_dtype), is dtype-following and has to be excluded, and doing that properly leaves 52 sites rather than nine. Most are vision or audio encoders, which run per image rather than per token, and MoE routing gates that are f32 by design, so the list still needs ranking before any of it is worth measuring.Both CSVs named in the description are updated:
benchmarks/metal_m5max_2026-09-09_single_mistral-small-4-119b-2603-4bit.csv,benchmarks/metal_m5max_2026-09-09_recheck_single_*.csvandbenchmarks/metal_m5max_vlm_2026-09-09_single_mistral-small-4-119b-2603-4bit.csv, with the baseline inbenchmarks/pylm_m5max_vlm_2026-09-09_mistral4check_single_*.csv.bench: close the M5 Max commit gap and fix two false doc claims
bench: close the M5 Max commit gap and correct two false doc claims (#…
fix(solar_open): pair the GPTQ biases dtype with the scales dtype (#1731