Add a fused full-attention path for head_dim 256 on NAX devices

[wyanzhao]

Proposed changes

Add a fused NAX full-attention path for head_dim == 256 without
materializing the [n_heads, qL, kL] score tensor.

Area	Change
Kernel configuration	bq=64, bk=32, wm=4, wn=2
Parallel decomposition	Separate attention_nax_dsplit kernel: split D and Dv across the WN = 2 simdgroups of the second warp dimension
Routing	Causal, no array mask, qL >= 1024, NAX-compatible dtype/TF32 mode
Ragged inputs	No padding copies. Ragged qL runs the kernel's unaligned-Q path (Q is loaded once, outside the KV loop, so it is free). Ragged kL from a KV-cache slice is read past its end up to the next bk multiple (the extra columns are causally masked); ragged kL from a plain array runs the unaligned-K path (~13% slower over the KV loop, still ~2x the unfused graph)
Fallback safety	Preserve fallback for unsupported shapes and float32 with TF32 disabled

Motivation

use_fallback currently allows fused full attention only for
head_dim in {64, 80, 128}. Full-attention layers in the Qwen3.5/3.6
family use head_dim == 256, so they fall back to an unfused graph that:

• materializes the full score tensor;
• computes the full rectangle before applying the causal mask; and
• becomes increasingly expensive at long context lengths.

A direct bd=256 instantiation of steel_attention_nax reaches only about
half the efficiency of bd=128 on M5 Max (22 vs 48 TF). The limiting
resource is the per-simdgroup accumulator byte budget: the bd=256 output
and score accumulator set is twice as large. A half-precision-accumulator
probe recovered +51%, while register pressure, scheduling fences,
accumulation-chain order, instruction footprint, and Q reload traffic were
ruled out by measurement.

Design

Head-dimension split

attention_nax_dsplit, a separate kernel from attention_nax dispatched
for bd = 256 with wn = 2, assigns each simdgroup one half of D and
Dv:

1. Each half computes its partial Q @ K.T scores.
2. The two halves exchange score fragments pair-by-pair through a 16 KB
   threadgroup buffer.
3. Both halves run softmax on the complete score tile.
4. Each half accumulates P @ V for its owned half of Dv.

Both halves use the same fragment-to-lane mapping, so the exchange can use a
per-lane-linear layout without coordinate conversion. Q fragments remain
register-resident across the KV loop.

This reduces the per-simdgroup accumulator set to 2 S + 8 O fragments
(320 B/lane), matching the working set of the existing bd=128 kernel.

Ragged causal inputs

There is no padding. Ragged query lengths take the kernel's unaligned-Q
path (align_Q = false), which measures within noise of the aligned one:
Q is loaded into registers once, before the KV loop, so raggedness never
touches the hot loop.

Ragged key/value lengths matter more: the unaligned-K pipeline
(align_K = false) carries bounds-checked K/V loads and column masking
inside the KV loop and measures ~13% slower over the whole loop on M5 Max.
The dispatch avoids it for the common case without copying anything: K/V
that come out of a preallocated KV cache (mlx-lm's KVCache grows its
buffer 256 positions at a time) are slices of a longer buffer,
keys[..., :kL, :]. When such a slice can be read past kL up to the next
bk multiple — the head's row pitch (strides[1] / strides[2]) reaches
that far and the buffer bound allows it — the kernel is simply told the
widened kL. This is safe under causal masking with qL_off computed from
the true lengths: every extra key column lies strictly beyond the
diagonal of every real query row, so it receives exactly zero probability
and its (finite) V row never reaches the output.

The check deliberately differs from the CUDA use_cudnn_for_decoding
heuristic in two ways: Metal buffers are page-rounded, so it tests
offset + extent <= buffer_size instead of equality, and it does not
require the cache capacity to be a multiple of 256 — after a chunked
prefill mlx-lm's cache is e.g. 2049 + 1280 = 3329 positions long, and
only the row pitch reaching the widened length matters. K/V that are not
such slices (a plain array with a ragged kL) run the unaligned-K path.

Routing and fallback behavior

Condition	Behavior
NAX, causal, no array mask, head_dim == value_head_dim == 256, qL >= 1024	Fused wn=2 path
qL < 1024	Existing fallback
Array mask present	Existing fallback
float32 with MLX_ENABLE_TF32=0	Existing fallback
Non-NAX device	Existing fallback

The qL >= 1024 threshold is based on measured M5 Max behavior. Below that
point, long-cache rectangles do not provide enough query blocks to fill the
machine (512 x 8192 measured approximately 1.08x the unfused time on an
earlier variant).

Performance

Unless noted otherwise: M5 Max, 16 Q / 2 KV heads, bf16, causal.

Kernel benchmarks

Shape	Before	After	Improvement
8192 x 8192	24.92 ms / 22.1 TF	12.84 ms / 42.8 TF	1.94x throughput
4096 x 4096	8.87 ms	3.45 ms	2.57x faster
2048 x 8192	8.92 ms	5.77 ms	1.55x faster
1024 x 8192	4.65 ms	3.58 ms	1.30x faster
2048 x 2048	2.39 ms	1.19 ms	2.01x faster

At 8192 x 8192, the new kernel reaches approximately 90% of the
bd=128 kernel's throughput.

Ragged lengths

M5 Max, 16 Q / 2 KV heads, bf16, causal, min of alternating A/B runs:

Shape (qL x kL)	Path	Time
4096 x 4096	aligned reference	3.46 ms
4064 x 4096	ragged Q, unaligned-Q path	3.61 ms
4096 x 4112	ragged K sliced from a longer cache (widened, aligned pipeline)	3.53 ms
4096 x 4112	ragged K from a plain array (unaligned-K path)	4.02 ms
8192 x 8200	ragged K, aligned vs unaligned-K path	15.2 ms vs 17.2 ms

Padding Q instead would have cost more than it saved (3.84 ms for
4064 x 4096), which is why the earlier padding copies were dropped.

End-to-end prefill

Qwen3.6-35B-A3B 4-bit, ragged prompts:

Prompt / chunking	Before	After	Change	Peak memory
8k, default 2048-token chunks	4039 tok/s	4193 tok/s	+3.8%	—
8k, single chunk	4209 tok/s	4822 tok/s	+14.6%	—
16k, 8k-token chunks	3457 tok/s	4180 tok/s	+20.9%	29.9 -> 26.1 GB
32k, 8k-token chunks	2833 tok/s	3598 tok/s	+27.0%	34.6 -> 26.5 GB
32k, single chunk	1771 tok/s	2873 tok/s	+62.2%	72.9 -> 41.8 GB

Transient memory for the 4096 x 4096 kernel case drops from 663 MB to
75 MB because the fused path no longer materializes the score tensor.

The end-to-end numbers were taken with an earlier revision that padded
ragged inputs with copies; mlx-lm hands over cache slices, which now run
the same aligned kernel without the pad/copy work.

Numerical behavior

Against an f32 reference, the fused bf16 output is more accurate than the
previous unfused graph:

Implementation	Relative L2 error
Unfused bf16 graph	2.3e-3
Fused NAX kernel	1.6e-3

Validation

• Fallback and fused shapes checked against the primitives reference in
  float32 and bfloat16.
• Committed coverage includes aligned/ragged squares and rectangles, plus
  ragged K/V sliced out of a longer preallocated cache with large, finite
  stale rows behind the slice (any leak into the output would be loud).
• Verified directly against mlx_lm.models.cache.KVCache slices from a
  chunked prefill (2049 then 1031 tokens) in bf16 and fp32/TF32.
• Additional local boundary checks cover odd lengths 1031 and 2049, plus
  the exact qL == 1024 routing boundary.
• A subprocess test with MLX_ENABLE_TF32=0 verifies that
  float32 + head_dim 256 remains on the fallback path.
• Full local test_fast_sdpa.py result: 20 tests, 1 platform-dependent
  skip, otherwise passing.

Checklist

• I have read the CONTRIBUTING document.
• I have run pre-commit run --all-files.
• I have added tests that cover the new fused and fallback behavior.
• No documentation changes are required for this internal kernel path.

[pierre427]

Validated on M5 Max (mlx 0.32.1, NAX active). Built the merge-base
(4367c73b) and PR head (1aee415d) from source with the
MACOSX_DEPLOYMENT_TARGET=26.2 NAX gate; bf16 4096³ GEMM canary confirms NAX in
both. D=256 causal SDPA A/B, GQA 32/8 and 64/8, S∈{16384,32768}, L∈{1..2048},
warmup + 40-iter median, run in both orderings (base-first and head-first) to
control for ambient throttling.

• Win on the large-L prefill shapes the kernel targets — L≥1024 at both GQA
  ratios and both S: head faster in both orderings (≈1.2–1.5×; the fused path
  avoids materializing the [n_heads, qL, kL] score tensor, so the gain grows
  with qL). These are exactly the Qwen3.5/3.6 full-attention (head_dim=256) shapes.
• No regressions on any shape in either ordering; decode (L=1) and mid-L
  (256/512) are within run-to-run noise.
• Correctness fp16-clean vs an fp32 reference (max_abs_err ~1e-4).

Caveat: this run was on a machine with intermittent GPU contention (canary drifted
2.5→4.4 ms across passes), so treat the magnitude as a direction-robust ~1.2–1.5×
rather than a precise figure — a clean-machine re-run would tighten it.

[zcbenz]

@jessegross Would you mind testing how much this is helping in production?

[zcbenz]

The KV cache written for MLX is usually guaranteed to be padded to 256 otherwise it would be too slow:
https://github.com/ml-explore/mlx/blob/main/docs/src/usage/kv_cache.rst

And we make use of this fact when using cuDNN SDPA for decoding:

mlx/mlx/backend/cuda/scaled_dot_product_attention.cpp

Lines 95 to 106 in 9ab977b

	// Check if k/v are slices from fixed-size kv cache.
	auto is_slice = [](const array& kv) {
	// Get pre-sliced sequence length from strides, and check if the buffer
	// belongs to a contiguous kv cache.
	int64_t T_kv = kv.strides(1) / kv.strides(2);
	if (kv.size() / kv.shape(2) * T_kv != kv.buffer_size() / kv.itemsize()) {
	return false;
	}
	// It is possible to use heuristic to check slices, but for now just make
	// mlx-lm work.
	return T_kv % kv_cache_step == 0;
	};

So I think there is no need to pad K/V here, we just need to add heuristics to know it is already padded. The original kernel can also make use of this fact and we can probably do it in another PR.

There is probably no longer need to pad queries as well since unaligned KV might be responsible for most of the slowdown.

[wyanzhao]

Done — the pad copies for Q, K and V are gone.

K/V: when they are slices of a longer buffer (keys[..., :kL, :] from a preallocated cache), the dispatch just tells the kernel a kL widened to the next bk multiple. With qL_off computed from the true lengths, every extra key column lies strictly beyond the causal diagonal of every real query row, so it gets exactly zero probability and its (finite) V row never reaches the output — same assumption the cuDNN decoding path makes about the rows behind a cache slice. Two deliberate differences from use_cudnn_for_decoding: Metal buffers are page-rounded, so the check is offset + extent <= buffer_size rather than equality; and I don't require the cache capacity to be a multiple of 256, only that the head's row pitch (strides[1] / strides[2]) reaches the widened length — after a chunked prefill mlx-lm's KVCache is e.g. 2049 + 1280 = 3329 positions long. Verified directly against mlx_lm.models.cache.KVCache slices (bf16 and fp32/TF32); the tests now include cache slices with large stale rows behind them.
Q: you were right — the unaligned-Q path is free (Q is loaded once, outside the KV loop; 4064x4096 measures 3.61 ms vs 3.46 ms aligned, while padding Q cost 3.84 ms), so Q is left ragged.

Kept this to the split kernel; the base kernel can pick it up in a follow-up as you suggested.

[jessegross]

@zcbenz I can take a look tomorrow but in general I think it is worth fusing head dim 256. It would be nice to have a version for pre M5 devices as well though :)

[zcbenz]

Did some simple inference with Qwen3.8 and I can observe consistent improvement on long prompt prefill, nice improvement!

[pierre427]

Re-validated on M5 Max at the current head. My earlier numbers (2026-07-16) were
against 1aee415d and are superseded by the kernel split and the padding-copy
removal since.

Built PR base db935c094 and head b86d328d6 from source with the
MACOSX_DEPLOYMENT_TARGET=26.2 NAX gate. D=256 causal SDPA, fp16, GQA 32/8 and
64/8, S ∈ {16384, 32768}, L ∈ {1, 8, 256, 512, 1024, 2048}, warmup + 40-iter
median, run in both orderings (base-first and head-first). bf16 4096³ GEMM canary
2.57–2.83 ms across the four passes confirms NAX active.

Noise floor first: same-arm drift between the two passes reaches 28.2% (base) and
23.6% (head) relative to the faster pass. Ratios below roughly 1.1x are not
separable from run-to-run variance on this machine.

Faster in both orderings at ≥1.03x — 8 of 24 shapes:

GQA	S	L	speedup (min of the two orderings)
32/8	16384	1024	1.03x (marginal — inside the drift band)
32/8	16384	2048	1.18x
32/8	32768	1024	1.20x
32/8	32768	2048	1.33x
64/8	16384	1024	1.21x
64/8	16384	2048	1.46x
64/8	32768	1024	1.36x
64/8	32768	2048	1.55x

Under a strict >1.0-in-both-orderings criterion the count is 11 of 24; the extra
three (64/8·16k·L=512 = 1.006, 64/8·32k·L=8 = 1.005, 64/8·32k·L=256 = 1.005) are
well inside the noise floor and I would not claim them.

The robust result is L ≥ 2048, where the four wins (1.18x–1.55x) sit outside the
drift band and grow with both qL and S — consistent with avoiding the
[n_heads, qL, kL] score materialization. The L=1024 wins are directionally
consistent, but only the 64/8 rows clear the noise comfortably.

Slower in both orderings — two shapes:

GQA	S	L	ordering A	ordering B
32/8	16384	256	−7%	−0.4%
32/8	32768	1	−2%	−1.9%

Both are inside the drift band and the −7% does not reproduce in the second
ordering, so I read them as noise rather than real regressions — but they are not
zero.

Elsewhere: mid-L (256/512) spans 0.86–1.13x across orderings with no consistent
direction; decode-ish (L=1/8) spans 0.85–1.24x. Both bands are dominated by drift.

Correctness: fp16 vs an fp32 reference, max_abs_err 2.4e-4 – 4.9e-4 across
GQA 32/8 and 64/8 at L=64 and L=256.

Happy to post the raw per-shape logs for all four passes if that is useful for
recomputing the table independently.

[jessegross]

This is what I see end-to-end on an M5 Max, which are consistent with @wyanzhao's numbers. The longer prompts are are chunked into batches of 2048.

That looks like a pretty worthwhile improvement to me.

qwen3.6:35b (MoE)

prompt tokens	prefill (base)	prefill (PR)	speedup	peak memory
2k	4582 tok/s	4703 tok/s	+2.6%	—
4k	4604 tok/s	4791 tok/s	+4.1%	−0.1 GiB
8k	3896 tok/s	4045 tok/s	+3.8%	−0.3 GiB
16k	3209 tok/s	3491 tok/s	+8.8%	−0.8 GiB

qwen3.8:27b (dense)

prompt tokens	prefill (base)	prefill (PR)	speedup	peak memory
2k	621 tok/s	633 tok/s	+1.9%	—
4k	621 tok/s	637 tok/s	+2.5%	−0.2 GiB
8k	613 tok/s	629 tok/s	+2.5%	−0.5 GiB
16k	593 tok/s	617 tok/s	+4.1%	−1.2 GiB

[zcbenz]

Thank you all for helping testing this PR! 🙏