Repository navigation
Allow force using fused kernel in scaled_dot_product_attention #3658
Description
Summary
PR #3293 added a fused full-attention path for head_dim=192/256 (steel_attention instantiation + use_fallback routing at kL>16384). The core change was reviewed favorably and validated by third parties on M2 Ultra and M3 Ultra, but the PR was closed when the author stepped away from this line of work. The last open question from @angeloskath was:
Simply put, a real world scenario would be even better. Something like
mlx_lm.generate ....results in OOM ...
We have been running head_dim=256 models in production on Apple Silicon for months, and have exactly that evidence.
Setup
- Mac Studio (Apple Silicon, 36 GB unified memory), MLX 0.31.x
- Runtime: oMLX (OpenAI-compatible server on MLX) with chunked prefill, prefix KV cache, and a Metal memory guard
- Model: Gemma 4 26B-A4B MoE (3-bit), head_dim=256 — 30 layers, of which only 5 attend over the full context (25 are sliding-window 1024)
Chunked prefill does not solve it
Chunked prefill avoids the single qL x kL self-attention, as noted in the PR thread. But on the unfused path each chunk still materializes fp32 scores against the whole cache, per full-attention layer:
transient ≈ n_q_heads × chunk_len × kv_len × 4 bytes
This grows linearly with kv_len, so a chunk size that is safe at 16K is not safe at 130K. In practice on our setup:
| Scenario (Gemma 4 26B, 36 GB machine) | Result |
|---|---|
| Cold prefill of a ~130K-token prompt, 1024-token chunks | 226 s — completes, transient-bound |
| Default 2048-token chunks at long context | Metal OOM abort (kIOGPUCommandBufferCallbackErrorOutOfMemory) — we had to patch the chunk size down and add memory guards |
| Multi-turn conversation growing past ~133K, memory ceiling active | guard shrinks chunks to 32–512 tokens to stay alive → a single turn's prefill takes 26+ minutes |
At kv_len=130K, even a 32-token chunk allocates ~665 MB for scores alone. Decode is unaffected (sdpa_vector already covers head_dim=256) — the gap is exactly the full-attention path #3293 addressed.
Note this is with only 5/30 layers attending over the full context; dense head_dim=256 models hit the same wall proportionally harder.
Control group (same hardware, same runtime)
head_dim=128 models — Granite 4.1 30B (dense), Nemotron-3-Nano 30B-A3B (hybrid MoE), and a cohere2-MoE 30B — show flat-to-linear per-turn prefill latency up to 164K context on the same setup. The wall is specific to the unfused head_dim=256 full-attention path, not the runtime or the hardware.
Ask
Could the final form of #3293 (bd=192/256 steel_attention instantiation + kL>16384 routing — unfused stays the default where it is faster) be reconsidered, either by reviving the branch or reimplementing it? head_dim=256 is common and growing across model families (Gemma 4, Qwen3.5/3.6 dense and MoE). Happy to benchmark on our workloads or test a branch.
Activity
DT-HJ commentedon Jun 11, 2026
Follow-up with concrete numbers: I rebased #3293's three commits onto current main (plus two small completeness fixes), built from source, and benchmarked on M4 Max (36 GB).
Branch: https://github.com/hojin12312/mlx/tree/sdpa-256-revival (original commits preserved with authorship; fixes: missing
bd=192steel instantiation that the routing already targeted, and excluding head_dim>=192 from the NAX dispatch branch since the NAX family has no 192/256 instantiations).Peak memory — the actual point of the change
n_q=16, n_kv=8 (GQA), head_dim=256, fp16(Gemma 4 26B attention shape), inputs included:Fused transient is ~16 MB regardless of kL; unfused grows linearly to +10.3 GB at 131K. Outputs agree exactly.
Speed — honest numbers
Same shapes, min/mean of 10 iters: above the kL>16384 routing threshold the fused kernel is 0.74–0.79x of unfused on M4 Max (e.g. qL=2048, kL=65536: 275 ms fused vs 214 ms unfused). Consistent with @angeloskath's earlier measurements — per-chunk, unfused is faster when the transient fits.
Why the memory column wins in production
A serving runtime holds 15–25 GB of weights+KV on a 36 GB machine. It cannot absorb a 3–10 GB fp32 transient per prefill chunk, so under memory pressure it shrinks chunks (2048 → 512 → 32 tokens) to keep the transient bounded — and GPU utilization collapses: we measure 26+ minute single-turn prefills at ~133K context (Gemma 4 26B, 5/30 full-attention layers). With the fused path the same step is a few hundred ms per chunk at full chunk size.
The kL>16384 routing already keeps unfused wherever it is faster and safe; the threshold is tunable if you'd prefer it gated on estimated transient size (
n_q*qL*kL*4) instead of kL alone.Happy to open this as a PR if there's interest.
DT-HJ commentedon Jun 11, 2026
Update — we put the branch into production serving and measured
end-to-end (Mac Studio 36 GB, oMLX, chunked prefill 1024, 27 GB prefill
safety cap, Gemma 4 26B 3-bit). The serving admission estimator was the
remaining blocker (it assumed the unfused transient and rejected long
prompts before the kernel ever ran); with the estimator taught the
fused O(1) transient:
PR opened: #3660
zcbenz commentedon Jun 13, 2026
@hojin12312 Would it be possible to measure with the mlx-lm instead of oMLX? We are not familiar with oMLX's code so it is hard for us to check if this is something that we could fix without sacrificing the performance.
Thump604 commentedon Jun 13, 2026
I ran a clean official
mlx-lmA/B locally against current MLX main and this PR branch.Scope: official
mlx-lmbenchmark path only. Excluded: oMLX, Runtime serving,Jobs, draft/MTP, custom scheduler.
Environment:
a025496c8f6da9d599e2010019ea5ee328c2d1190aa50825acca3be22c8f02c5581d17aa442ba742mlx-lm:2c008fd0252b2c569227d12568356ab88ab0560a/0.31.3python -m mlx_lm.benchmarkResults:
Interpretation: official
mlx-lmreproduces the memory/speed tradeoff on Qwen35B: lower peak memory, materially slower prompt throughput. In this benchmark
shape, Gemma 4 26B did not show a peak-memory improvement at 32K or 64K.
I also tried an optional Qwen 35B 131K row. Main completed with
prompt_tps=516.049,peak_memory=49.971 GB; the PR branch aborted beforecompleting the trial with:
I would treat the 131K row as an unresolved follow-up, not a valid A/B
comparison.
Full local artifact:
/opt/ai-runtime/run/mlx-sdpa-3660-validation/20260613T041902Z/summary.jsonDT-HJ commentedon Jun 13, 2026
Thanks. I re-ran this through the stock
mlx-lmbenchmark path — no oMLX, no serving layer: vanillamlx 0.31.2 vs the PR wheel at
0aa50825,python -m mlx_lm.benchmark, Gemma 4 26B UD-MLX 3-bit and4-bit, 32K→131K, prefill-step 2048. (The two arms are release builds 71 commits apart, not a
single-commit A/B; the static dispatch analysis below is what isolates the PR's behavior.)
No difference. Peak memory is identical between the two builds in every row (to the benchmark's
0.001 GB precision; e.g. 4-bit 131K = 28.088 GB both), and prompt_tps matches within 0.1%. This agrees
with @Thump604's earlier Gemma result.
Root cause — and a correction to my own earlier micro-benchmark: Gemma 4 26B/31B run their
full-attention (global) layers at
head_dim=512(global_head_dim), not 256. Only thesliding-window layers are head_dim 256, and those are bounded (kL ≈ window+chunk, well under the 16384
threshold), so they never hit a large-kL full-attention SDPA. Since 512 is not in the fused dispatch
set, both vanilla and the PR build take the same unfused
matmul→softmax→matmulfallback for thoselayers — so the new 192/256 kernel can't change anything for these models. My earlier table labeled
(n_q=16, n_kv=8, head_dim=256)as "the Gemma 4 26B shape", but that's actually the sliding layer;the real large-kL shape is
(n_q=16, n_kv=2, head_dim=512). Apologies for the mislabel.One more correction: the fallback here materializes the score tensor in bf16 (the model's
activation dtype;
precise=Truesoftmax only raises the accumulator to fp32), ~8 GiB at 131K withqL=2048 — not an fp32 transient. The peak grows about 2.8–3.0 GB per 32K of context in my runs; a
bf16 score plus global KV predicts ~2.82 GB (with a small residual from other buffers), whereas an
fp32 score would predict ~4.96 GB and clearly does not match the observed slope. So this is a real,
kL-scaling materialization that vanilla does not stream — the 156K→184K admitted-context change I
reported in production came from relaxing our serving-side admission estimator, not the kernel, and I
now consider that estimator change an under-estimate we need to fix on our side, not evidence about MLX.
Net for PR #3660: its new 192/256 full-attention kernel is effectively a no-op for Gemma 4 26B/31B
(their long-context layers are 512). It can still help models whose actual full-attention head dim is
genuinely 192 or 256; my result only shows Gemma 4 26B/31B aren't such models, so I wouldn't gate any
Gemma-4 decision on it. Full 16-run table below.
Full 16-run mlx-lm A/B (peak GB / prompt_tps)
Models:
unsloth/gemma-4-26b-a4b-it-UD-MLX-3bitandunsloth/gemma-4-26b-a4b-it-UD-MLX-4bit(download revision not recorded). mlx-lm 0.31.3; vanilla mlx 0.31.2 vs fused wheel
0.32.0.dev20260611+0aa50825;
mlx_lm.benchmark --model <dir> -p <ctx> -g 8 --num-trials 1 --prefill-step-size 2048; Mac Studio 36 GB.peakat the benchmark's 0.001 GB precision.(The no-difference conclusion is the prefill peak + prompt_tps above, plus the static dispatch
analysis; decode
-g 8/1-trial is too short to compare and isn't part of the claim.)zcbenz commentedon Jun 14, 2026
Thanks for running the benchmark. The performance hit is not small so I think we should not blindly route to fused sdpa? It might be nice if we have heuristics to estimate the needed RAM and then check memory limitation but it is going to be quite hard to get right for all cases on framework's side.
I wonder if it makes sense adding a
force_fusedoption to fast sdpa op and let the inference engines or users to decide. @angeloskathDT-HJ commentedon Jun 14, 2026
That sounds reasonable to me. The Qwen results show that a fixed
kLthreshold is too blunt, since whether the memory reduction is worth the throughput cost depends on the runtime's resident weights, KV cache, and memory budget.Making the 192/256 full-attention fused path opt-in via
force_fusedseems like a cleaner boundary: MLX can expose the supported kernel, while inference engines or users with the relevant memory context decide when to select it. I would expect the default to preserve the existing MLX dispatch, with the option only requesting the fused path when the shape and other capability checks support it.I'm happy to update #3660 accordingly once the intended API semantics are agreed on.
ddalcu commentedon Jul 11, 2026
Independent confirmation from a second production runtime (mlx-serve), plus data on how far the composed path can be pushed, and +1 on the
force_fusedopt-in.We ship mlx-serve (a Zig serving runtime on mlx-c) and hit exactly this class: effectively every checkpoint we serve is head_dim=256 (Gemma-4 E4B/26B/31B, Qwen 3.5/3.6 dense and MoE), so the unfused full-attention path is our default long-context reality, not an edge case.
1. Isolated measurement (MLX 0.32.0, M4 Max 128 GB). One
mx.fast.scaled_dot_product_attentioncall, bf16,mask="causal", GQA shapes, peak viareset_peak_memory:On 0.32.0 with bf16 inputs the fallback's materialized scores measure ≈
heads × qL × kL × 2bytes (+~5%).2. Lazy eval multiplies it. Chaining N attention layers with an
mx.evalevery 4 layers, ~3.2 layers' score transients coexist at peak (3.44 GB measured for a 1.07 GB/layer shape); eval-per-layer bounds it to ~1× at sub-1% wall-clock cost. So a runtime's effective transient is 1–3× one layer's scores depending on its eval cadence.3. Real-world A/B — Gemma-4 26B-A4B (4-bit, quantized KV), cold prefill of a 102,448-token prompt, M4 Max 128 GB, byte-identical outputs:
Two takeaways that I think support the direction this thread is heading:
kLthreshold inside MLX can't pick correctly, andforce_fusedas an opt-in is the right boundary: the caller is the only party that knows resident weights, KV size, per-request cache quantization, and the machine's budget. A serving engine already computing an admission estimate can flipforce_fusedprecisely when the floor matters and eat the throughput cost knowingly.Happy to benchmark a revived #3293 wheel on M4 Max across Gemma-4 (E4B / 26B-A4B) and Qwen 3.6 (27B dense / 35B MoE) at 32K–262K, with dense and quantized KV, if that helps get it over the line.
Thump604 commentedon Jul 11, 2026
@zcbenz @angeloskath Before rebasing #3660, could you confirm the intended API and semantics for a caller-selected fused full-attention path?
Exact local reproduction used the stock
mlx-lmbenchmark path, not a serving runtime:Against the June comparison bases (MLX main
a025496and #36600aa50825), Qwen 3.6 35B measured 41.316 -> 39.882 GB peak at 32Kand 44.201 -> 40.068 GB at 64K, while prompt throughput changed
1005.743 -> 871.489 and 802.929 -> 624.513 tokens/s respectively.
The #3660 arm did not complete the optional 131K row, so I am not
claiming a result for that context.
The relevant upstream paths are
use_fallback()andsdpa_full_self_attention_metal()inmlx/backend/metal/scaled_dot_product_attention.cpp, plus the compiledSteel full-attention shapes in
mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal.The minimal ask is not a new default threshold. It is confirmation of the
preferred opt-in surface: for example, a caller-selected
force_fusedmode, or another API name you prefer, which requests the compatible fused
full-attention path while preserving current automatic dispatch by default.
That would let an inference engine make the memory/throughput decision from
its own resident-memory and request context, while unsupported shapes retain
the normal path.
I will prepare a fresh, narrow PR on current main only after that API
direction is confirmed. This is not a request to merge the current dirty
#3660 branch, not a claim that its historical
kL > 16384threshold is theright global policy, and not a Gemma 4 26B/31B long-context claim: those
models' global attention uses head_dim=512 and does not exercise this
192/256 full-attention route.
apocryphx commentedon Jul 21, 2026
Another real-world datapoint for reviving #3293, from a different corner of the same model family: we run Gemma-4-31B dense (from-scratch MLX C++ engine, M4 Max 128 GB, bf16 compute, Q4 weights) — 50 sliding-window layers at d=256 (32 q / 16 kv heads) plus 10 global layers at d=512. With the current routing (
sdpa_full_supported_head_dim ∈ {64, 80, 128}, re-verified on 0c537a4), no layer of this model fuses at prefill at all.Where the OP's evidence is memory (OOM at 100K+), ours is a throughput/scaling datapoint at moderate context, isolated to single
mx.fast.scaled_dot_product_attentioncalls (array mask, bf16):Per-FLOP the composite path is respectable at this L — the dominant costs are the materialized
[n_kv, n_rep, qL, kL]score transients (≈1 GB/layer at 4K, growing quadratically, i.e. exactly the OP'schunk × kv_lenwall at long context) and the memory-bandwidth those cost. So our data supports the same conclusion from the throughput side: the fused kernel's value at d=256 is primarily the O(L) memory behavior, and it compounds with context length.One nuance worth having in the thread, since both reported models are Gemma-4 variants where most d=256 layers are sliding-window (25/30 for the OP, 50/60 for us): a fused kernel with an arbitrary array mask still processes the full qL×kL area unless it skips fully-masked tiles. For the sliding layers there's an algorithmic mitigation available today that composes with the requested kernel rather than substituting for it: chunk the prefill and trim K/V to the last (window + chunk) positions before the SDPA call — every dropped key is outside every query's window, so outputs are unchanged (we gate this greedy-token-exact). That took our d=256 composite cost from 46.7 → 18.1 ms/layer at L=4096 and scales as O(L·window) instead of O(L²). With #3293 revived, the fused kernel would then run on the bounded
[chunk × (window+chunk)]problem — memory and compute both bounded. It does nothing for the global (full-attention) layers, which is where the fused kernel is irreplaceable — that plus the d=512 case is filed separately as #3885.Happy to benchmark a revived #3293 on M4 Max (both the dense 31B shapes above and isolated kernels).
beyondsimulations commentedon Aug 5, 2026
Another data point for this issue, from a shipping macOS app (on-device
handwriting OCR with Qwen3-VL via mlx-swift): the unfused fallback is not just
slow — at vision-model prefill sizes it is a hard, unrecoverable crash, and it
is silent until then.
Qwen3-VL-8B's vision tower has
hidden_size 1152 / num_heads 16→ head_dim72, which misses
sdpa_full_supported_head_dim(64/80/96/128 on currentmain), so every ViT prefill materializes the full score matrix. Byte-exact
example from a real 1404×4079 px document page:
N = 88 × 254 = 22,352patches (16 px patches)16 heads × 22,352² × 2 B= 15,987,580,928 B (14.89 GiB)maxBufferLengthon a 16 GB M-series machine = 9,534,832,640 B (8.88 GiB)The result is
[metal::malloc] Attempting to allocate 15987580928 bytes which is greater than the maximum allowed buffer size of 9534832640 bytesraised asan uncatchable
fatalError, with the process left spinning afterwards — a hostapp cannot catch it, fall back, or even die cleanly. Qwen3-VL-4B (head_dim 64)
is immune at any image size, which isolates the gate as the cause.
The model-level fix exists (mlx-swift-lm #455 and mlx-vlm both zero-pad 72→80
onto the fused kernel — exact, since padded dims contribute nothing to the dot
products), but every downstream port has to rediscover and re-implement the
same trick after debugging a crash. Two asks for mlx core, either of which
would have prevented this:
by padding internally (or at minimum add 72, which two model families ship);
maxBufferLengthas a recoverable error instead offatalError— callersthat could chunk their input currently get no chance to.
Related: the silent-fallback observability gap discussed in #3885 is exactly
what made this expensive to debug — there is no way to know you fell off the
fused path until the allocator aborts.
zcbenz commentedon Aug 11, 2026
Let's add a
force_fused=Trueoption tomx.fast.scaled_dot_product_attentionAPI that ignores the heuristics and uses fused path, and throws when a fused kernel is not available for the shape. I would be happy to review a PR implementing it.[-]head_dim=256 unfused SDPA makes 100K+ context prefill impractical — real-world evidence for reviving #3293[/-][+]Allow force using fused kernel in scaled_dot_product_attention[/+]apocryphx commentedon Aug 11, 2026
@zcbenz We'll take that PR, if nobody has started — with one safety pairing from the #3885 investigation that
force_fusedmakes load-bearing: the vector dispatches got theircheck_kernel_threadgroup_sizein #4018 and #3099, but the full-attention dispatch (sdpa_full_self_attention_metal) still launches unchecked. Under heuristic routing it only runs at shapes whose pipelines are known to fit;force_fusedmakes register-limited configurations reachable — and on M1-class GPUs an oversized launch is silently dropped and returns zeros (measured; details in that thread). Wide head dims are exactly where this bites: in our head_dim=256 experiments the specialized full-kernel pipelines on M4 Max already reportmaxTotalThreadsPerThreadgroup = 128— precisely the 4-warp dispatch size, zero headroom — so devices with smaller register budgets would cap below the launch. The check should extend to the full dispatch together with the option.Proposed semantics, objections welcome:
force_fused=Truebypassesuse_fallbackentirely; vector-vs-full selection still follows query length as today.false; zero behavior change for existing callers.Vector-kernel d=512 instantiations (with the BN=16 threadgroup shape M1-class needs — three-device sweep data in #3885) would follow as a separate small PR. Since
force_fusedis explicit opt-in, they would ship with no routing-heuristic changes and zero behavior change for default callers — which removes the regression risk that stalled the earlier head-dim expansion.zcbenz commentedon Aug 19, 2026
Closed by #4185.
ddalcu commentedon Aug 19, 2026
Thank you @zcbenz, @hojin12312, @Thump604 !