metal: fused Sdpa via ported MLX attention kernels (native GQA, decode kernel)

[czoli1976]

Follow-up to #2320, same fused-Sdpa slot: MLX's attention kernels ported as
owned .metal source — the decode-specialised sdpa_vector family (single-pass
and split-KV two-pass) and the tiled steel_attention prefill kernel — behind a
single chooser translator that prefers them, keeps the vendored MFA metallib for
shapes they don't cover, and explodes the rest.

Two things the metallib can't do. It predates GQA, so llama-class attention
falls out of fusion entirely; and it is a prefill kernel, so at decode shapes
(qL=1) it is slower than not fusing at all — which is what main does today.

The mask input is wired through as an additive float mask. That turned out to be
the difference between a synthetic win and a real one: an exported causal LLM
passes its mask as a fourth Sdpa input rather than setting is_causal, so
without it nothing in a real model fuses. It arrives as f32 next to f16
activations, and the kernels template the mask on the activation type, so a
rewrite casts it to the query dtype — the mask is a constant, so it folds.

Same vendoring idiom as MlxGemm/MlxGemv: mlx_sdpa.metal is a mechanical
flattening of the MLX include closure at a pinned commit, each section verbatim,
f32/f16 entry points instantiated at the end; resync by re-flattening.

End to end

Qwen2.5-7B-Instruct q40ef16 (--metal) on the M1 Pro. All 28 attention layers
fuse to MetalMlxSdpa at both shapes; on main the same graphs show 28
MetalScaledMaskedSoftmax, i.e. nothing fused.

	main	this PR
decode, S=1 P=1024	135.5 ms/token	105.8	1.28×
prefill, S=512 P=0	1022 ms	816	1.25×

Best of five interleaved runs each. This box is noisy under a 4 GB model, so the
per-pair spread is wide — decode 1.00–1.49× (ahead in 4/5), prefill 0.90–1.40×
(ahead in 4/5, one pair against).

Same model on the M4, which is a quieter box — three interleaved runs, best of:

	main	this PR
decode, S=1 P=1024	108.2 ms/token	88.7	1.22×
prefill, S=512 P=0	247.3 ms	184.1	1.34×

Per-pair there was 1.22–1.27× decode and 1.34–2.91× prefill, ahead in every pair.

Model output stays finite — logits and all 28 KV cache pairs — on a model that
historically produced NaNs on this backend.

Second device (M4, 8-core GPU, macOS 26.6)

8-layer stacks, ms/layer. The metallib gap widens on this GPU, and prefill
improves over the M1 Pro's 1.34–1.46×:

	MLX port	MFA metallib	explode
prefill f32 S=512	0.390	0.697 (1.79×)	—
decode f16 kvL=4096	0.060	1.758 (29.1×)	0.085 (1.41×)
decode f32 kvL=4096	0.104	2.037 (19.6×)	0.108 (1.04×)
decode f32 kvL=512	0.058	0.272 (4.68×)	0.039 (0.67×)

One case needs a caveat rather than a claim. bench_sdpa_decode first showed
f32 with a short cache at 0.67× against exploding on this GPU, so I went looking
for a gate rule and swept 32 decode shapes (f16/f32 × D 64/128 × MHA and GQA 32/8
× kvL 256..4096) on both machines — 64 measurements, bench_sdpa_gate_sweep.
No rule emerged: the fused path wins 28/32 on the M1 Pro and 26/32 on the M4, and
the handful of losses are 0.82–0.99×, scattered, and do not replicate across the
two devices.

Re-running the original case several times explains it — the spread is warm-up,
not shape. The MLX kernels are compiled at pipeline-load time, and the fused
model is benched first, so the first measurement in a process pays that:
0.57× then 0.94× then 1.09× across three consecutive runs on the M4. So: no
measured regression I can reproduce, and no gate narrowing proposed. Flagging it
because the first number was real and someone else will hit it.

bench_vector_passes also confirms the split-KV gate on a second device: the M4
reports applegpu_g16g, mlx does not treat it as large, and single-pass is indeed
faster there at 4096 (0.167 vs 0.186) and 16384 (0.634 vs 0.655).

Kernel and op numbers

M1 Pro (14-core GPU), macOS 26.5.2, 8-layer Sdpa stacks, ms/layer, best of
5×N, KV resident on the device. Range over three runs.

Prefill, f32, H=8 S=512 D=64:

	MLX port	MFA metallib	explode
ms/layer	0.44–0.46	0.61–0.66	1.08–1.11

→ 1.34–1.46× over the metallib, ~2.4× over explode.

Decode, qL=1, H=8 D=64:

	MLX port	MFA metallib	explode
f32, kvL=512	0.069–0.083	0.238–0.290	0.076–0.098
f32, kvL=4096	0.124–0.173	1.258–1.290	0.208–0.417
f16, kvL=4096	0.144	1.083	0.725

→ 3.2–10.2× over the metallib, 1.0–5.0× over explode. Note the metallib
column is 2.6–3.1× slower than explode here: today a decode-shaped Sdpa
fuses to it, so this also removes that regression.

bench_sdpa_multilayer_mlx_vs_mfa swaps the op in the transformed graph, so
both kernels are measured in one process against the same graph, allocator and
sync pattern.

One deviation from the parked branch: the two-pass split-KV gate mirrors mlx's
own rule (large GPUs at ≥1024 keys, otherwise GQA at ≥4096) instead of splitting
unconditionally at 1024, so devices mlx does not treat as large keep the
single-pass kernel. This M1 Pro reports applegpu_g13s, which mlx does treat
as large, and the kernel-level numbers say that is the right call here
(bench_vector_passes, ms/dispatch):

kvL	1-pass	2-pass
1024	0.0475	0.0546
4096	0.1395	0.0831
16384	0.7356	0.4033

So the split pays from ~4096 and is marginally wrong at exactly 1024; mlx's
threshold is kept rather than retuned on one device.

Validation

cargo test -p tract-metal --release: 93 passed, 8 ignored (benches), plus 16
new mlx_sdpa correctness cases (f32/f16 × 1-pass/2-pass/steel, GQA, causal,
unaligned, batched) checked against a CPU reference. cargo build --workspace
clean, cargo fmt --all clean, cargo clippy -p tract-metal clean.

One failure, mfa::tests::test_mfa_attention_causal_const_is_noop, is
pre-existing: it fails identically on main at 405dd8d on macOS 26.5.2 (the
triangular function constant does mask now, so the probe's assertion no longer
holds). Untouched here; happy to send a separate fix.

Not included: the M5 tensor-op (NAX) variant of the steel kernel. It needs M5
silicon to validate and I have none — parked until someone can run it.

MLX (ml-explore/mlx) is MIT, Copyright (c) 2023-2025 Apple Inc.; attributed in
the source header. Dispatch decisions follow MLX's own
scaled_dot_product_attention.cpp.

🍍

[czoli1976]

The pre-existing failure mentioned above is now split out as #2546 — that one is independent of this PR and applies to main on its own.

[czoli1976]

The M5 tensor-op variant mentioned above is now filed separately as #2548, with the branch rebased on top of this PR — it needs one run on M5 hardware, which I don't have.

[kali]

/ci llm

[github-actions]

⚠️ Bench vs main — no speed regressions · 1 secondary regression(s)

Reference: 2026-08-02 morning nightly run (1d old) · full report → run

Speed — evaltime · prefill · decode

no inference-speed regressions

Improvements

Δ	metric	device	main → PR
🟢 -21.7%	parakeet_tdt_600m_v3_f32f32_decoder_pass evaltime · cpu	i9-11900kb_rtx-4060	0.838 ms → 0.656 ms

⚠️ 1 secondary regression(s)

Δ	metric	device	main → PR
⚠️ +6.4%	llama_3_2_3B_instruct_q40ef16_541 load+optimize · cuda	jetson-orin-nx	4.14 s → 4.41 s

[github-actions]

⚠️ Bench vs main — no speed regressions · 6 secondary regression(s)

Reference: 2026-08-03 morning nightly run (0d old) · full report → run

Speed — evaltime · prefill · decode

no inference-speed regressions

Improvements

Δ	metric	device	main → PR
🟢 -21.7%	parakeet_tdt_600m_v3_f32f32_decoder_pass evaltime · cpu	i9-11900kb_rtx-4060	0.838 ms → 0.656 ms

⚠️ 6 secondary regression(s)

Δ	metric	device	main → PR
⚠️ +22.6%	arm_ml_kws_cnn_m load · pass	cortex-a9	84 ms → 103 ms
⚠️ +18.3%	arm_ml_kws_cnn_m load+optimize · pass	cortex-a9	131 ms → 155 ms
⚠️ +11.9%	hey_snips_v1 load · 400ms	cortex-a7	67 ms → 75 ms
⚠️ +11.5%	arm_ml_kws_cnn_m load · pass	cortex-a7	78 ms → 87 ms
⚠️ +10.7%	hey_snips_v1 load+optimize · 400ms	cortex-a7	112 ms → 124 ms
⚠️ +6.4%	llama_3_2_3B_instruct_q40ef16_541 load+optimize · cuda	jetson-orin-nx	4.14 s → 4.41 s

[github-actions]

✅ CI / large-models: success

• ✅ cli: success
• ✅ foundation-llms: success
• ✅ foundation-llm: success
• ✅ parakeet-tdt-600m-v3: success
• ✅ nemotron-speech-streaming-en-06b: success

View workflow run