head_dim=256 unfused SDPA makes 100K+ context prefill impractical — real-world evidence for reviving #3293

[hojin12312]

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.

[hojin12312]

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=192 steel 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:

qL	kL	inputs	fused peak	unfused peak	max abs err
1024	32768	264M	272M	1,824M	0.0
2048	32768	272M	288M	2,880M	0.0
2048	65536	528M	544M	5,696M	0.0
2048	98304	784M	800M	8,512M	0.0
2048	131072	1,040M	1,056M	11,328M	0.0

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.

[hojin12312]

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:

Scenario	unfused	fused
Max cold single-shot prefill	~156K (estimator-bound)	184K (629 s; 204K rejected cleanly at 182K — now genuinely memory-bound)
Multi-turn ceiling	~133K, then 26+ min/turn crawl	167K, 99% cache hit, 9–12 s/turn, crawl gone
25K paste on a 130K conversation	impractical	83 s
Qwen3.6 35B 145K cold	280 s	332 s (0.84x — honest per-kernel cost where memory never bound)

PR opened: #3660