[jvmenen]

Why

When MTPLX checks a block of 9 to 32 proposed tokens at once, attention gets several times slower than for 8 tokens. Such blocks come from context-copy, which proposes a whole span of text found earlier in the prompt. At head size 256, MLX's fast attention kernel handles at most 8 query rows; for more rows MLX falls back to a route that first writes the full table of attention scores to memory and reads the KV cache several times. On Qwen3.8-27B with an 80K-token conversation, attention for a 13-row block takes 11.5 ms per layer, against 2.6 ms for an 8-row block on MTPLX's own kernel, and a whole verify step of 9 rows takes 286 ms against 156 ms for 8 rows. A longer copied span, which should save time, therefore makes each step much more expensive.

How it works, in plain terms

In a verify step the model computes, for every proposed token (a row), how much attention it pays to each earlier token in the conversation. A flash kernel does that in one pass over the cache and keeps the intermediate scores on the chip, instead of writing the whole score table to memory and reading it back. MTPLX already has such a kernel for blocks of up to 8 rows. With MTPLX_NAX_FLASH_WIDE=1, blocks of 9 to 32 rows go to the same kernel, which now processes the rows of each KV head in groups of 32 within one GPU launch. Blocks of 8 rows or fewer do not change. For 9 or more rows the rounding differs slightly from the stock route; at every tested position the chosen token was the same, except where two tokens were exactly tied in bf16.

What changes

MTPLX_NAX_FLASH_WIDE=1 (off by default) sends q_len 9 to 32 to the existing head-dim-split flash kernel (sdpa_nax_flash_dsplit). The kernel now tiles the query rows: the 6 × q_len rows of a KV head are processed in groups of 32 rows within one launch, on a new grid axis, and split over the KV length as before.

• mtplx/kernels/sdpa_nax_flash_dsplit.py: new template parameter TG_M (16-row blocks per threadgroup) and a row-group axis. Tail-causal masking uses the global row index. The shared gate and dispatch serve sdpa_nax_flash_dsplit (contract unchanged, up to 32 rows) and the new sdpa_nax_flash_dsplit_wide (up to 192 rows).
• mtplx/attention_split.py: with the switch on, the packed-GQA route admits q_len up to 32 and routes 9 or more to the wide kernel. Rows 6 to 8 stay on sdpa_nax_flash, which is as fast there. A bail falls back to stock SDPA.
• kernel_selfcheck.py: new lane nax_flash_dsplit_wide_sdpa, probed only when the switch is on.
• CHANGELOG line under [Unreleased].

Results

M5 Pro 64 GB, stock MLX 0.32.2, v2.12.2 plus this branch, Qwen3.8-27B Optimized-Speed.

Attention op at an 80K prefix, 16 layers chained, ms per layer:

q_len	stock	sdpa_nax_flash	wide
9	9.83	3.30	2.70
13	11.54	bails	3.61
16	6.80	bails	3.69
24	7.88	bails	5.67

The error against an fp32 reference stays within the stock route's error at every shape tested (rel. 2.7e-3 against 4.4e-3).

Whole verify forward (target_verify, median of 7, ms, switch off → on):

Rows	20K	50K	80K
9	154 → 124	221 → 140	286 → 164
13	176 → 138	264 → 161	325 → 189
16	147 → 139	192 → 163	241 → 190
24	219 → 215	272 → 250	336 → 294

Rows 4 to 8 are bit-identical. For 9 to 25 rows the argmax is equal at every position at 20K and 80K. At 50K it differs at 1 or 2 positions per block, and every such position is an exact bf16 tie (top-2 margin of 0 or one bf16 step).

End to end (generate_mtpk, depth 3, context-copy on, 600 tokens, greedy, median of 3 interleaved, one process per series):

Prefix	Workload	off tok/s	on tok/s	ms per verify
50K	edit inside a file	60.0	66.7 (+11%)	192 → 176
50K	code edit	38.4	40.7 (+6%)	153 → 145
80K	edit inside a file	32.7	34.8 (+6%)	157 → 147
80K	code edit	39.1	44.1 (+13%)	192 → 170

Tests

tests/test_sdpa_nax_flash_dsplit_wide.py checks the kernel against an fp32 tail-causal reference and the stock route for q_len 1 to 32. It covers several context lengths, 1, 2 and 4 blocks per threadgroup, a tensor offset, block counts and the contract bails. Existing dsplit tests pass unchanged (q_len up to 5 is bit-identical with the switch off). Full suite passes.

Notes

• Off by default.
• With a large cache buffer and a short prefix (2K), the wide kernel is slower than stock above 16 rows. Routing keys on cache capacity, which follows the prefix length in normal serving.
• MTPLX_NAX_FLASH_WIDE_TGM sets the row blocks per threadgroup (default 2).

🤖 Generated with Claude Code