Add M5 NAX prefill kernel for paged GQA attention (#623)

<img width="2496" height="884" alt="bench_nax_prefill"
src="https://github.com/user-attachments/assets/d482a22b-2e37-435c-8fbc-a7097deb2a6d"
/>

## Summary

Adds an automatic, opt-out paged-attention prefill kernel for the NAX
tensor units on M5 GPUs.

- Supports fp16/bf16 MHA/GQA/MQA with head sizes 64, 96, 128, 256, or
512 and block sizes 8, 16, or 32.
- Pure decode, speculative verification windows, TurboQuant, unsupported
hardware/shapes, and missing NAX libraries retain their existing paths.
GDN linear attention is unchanged.
- Mixed batches use NAX for the whole prefill dispatch, including
`q_len=1` decode passengers.
- Set `VLLM_METAL_DISABLE_NAX=1` to force the existing tiled kernel.

The pure-decode kernel is unchanged. End-to-end TPOT can still improve
during chunked prefill because prefill and decode share engine steps.

## Future work

* maybe GDN support
* although this PR's head_dim 256 512 is better than main, it's still
under-utilize the NAX. We could refer to upstream PR for reference on
future work https://github.com/ml-explore/mlx/pull/3842

## Performance

Measured on an M5 Pro with macOS 26.6 and MLX 0.32.0:

| Workload | Metric | This PR vs main |
|---|---|---:|
| Qwen3-8B-bf16, 8 x 4096-token prompts | Mean TTFT | **-21%** |
| Qwen3-8B-bf16, 6 x 7900-token prompts | Mean TTFT | **-23%** |
| Qwen3-0.6B, sonnet 2048 in / 32 out, concurrency 32 | Total tok/s |
**+33%** |
| Same | Mean TTFT | **-41%** |
| Gemma 4 E4B bf16, sonnet 2048 / 32, concurrency 32 | Total tok/s |
**+27%** |
| Same | Mean TTFT | **-24%** |
| Qwen3.8-27B 8-bit | Serving throughput / TTFT | no measurable change |

Qwen3.8's head-size-256 attention kernel is about 2.1x faster, but only
16 of 64 layers use SDPA and 8-bit weight dequantization dominates the
serving step.

## Validation

All 60 M5 specialization pipelines pass parity. Targeted MHA/MQA, fully
masked rows, dual-arm golden generation, AOT packaging, affected tests,
and GitHub CI also pass.

<details>
<summary>Kernel and serving measurements</summary>

### Paged bf16 attention speedup

| Head size | 2K | 4K | 8K |
|---:|---:|---:|---:|
| 96 | 2.77x | 3.14x | |
| 128 | 2.98x | 3.63x | 3.36x |
| 256 | 2.16x | 2.02x | 1.85x |
| 512 | 4.36x | 5.69x | |

At actual head-size-256 GQA geometry, the 4K kernel speedup is 2.09x for
Qwen3.5 (8 query / 2 KV heads) and 2.06x for Qwen3.8 (24 / 4).

Additional serving results:

- Qwen3-8B, 24 x 4096-token prompts at concurrency 4: +57% total tok/s
and -38% mean TTFT.
- Gemma 4 E4B: -22% duration and -22% mean TPOT.
- Qwen3-0.6B 2048/32: -25% mean TPOT. Decode itself does not use NAX;
faster chunked prefill reduces interference with decodes in the same
engine step.

Serving A/B runs use one build and toggle only `VLLM_METAL_DISABLE_NAX`.

</details>

<details>
<summary>Implementation, MLX provenance, and dataflow</summary>

The 449-line shader adapts MLX 0.32's `steel_attention_nax`: BQ=64,
BK=32, four simdgroups, 16x32x16 MPP `matmul2d`, register-fragment Q/K/V
loads, online softmax, relaxed precision, and the head-size-128
scheduling barrier. vLLM extends the schedule to head sizes 96, 256, and
512, using a midpoint PV-loop split for 256/512.

| Portion | Size | Implementation |
|---|---:|---|
| NAX fragment helpers | 55 non-comment lines | MLX lane mapping,
cooperative-tensor packing, and row reduction |
| Paged/varlen attention body | 259 non-comment lines | vLLM addressing,
masking, and online-softmax integration |
| Specializations | 43 non-comment lines | 30 kernels: two dtypes x five
head sizes x three block sizes |

The vLLM body adds block-table gathers, packed varlen queries,
chunked-prefix causal positions, partial-row/page guards, sliding
windows, softcap, sinks, and GQA head mapping.

The dataflow deliberately differs from the classic tiled kernel. Tiled
stages Q/K/V in threadgroup memory and reuses K/V across four
simdgroups; NAX follows MLX and loads register fragments directly,
relying on GPU caches. Controlled ablations found that staging NAX K/V
was 34% to 58% slower, while loading tiled K/V directly was 32% to 63%
slower. Direct loading is therefore specific to the NAX fragment
geometry, not a universal optimization.

The relaxed MPP descriptor changes accumulation order and truncates the
fp32 P operand of PV, so validation uses bounded parity rather than
bitwise identity.

</details>

<details>
<summary>Compatibility, correctness, and release packaging</summary>

NAX requires M5 hardware, macOS 26.2 or later, a successfully loaded NAX
metallib, fp16/bf16, a supported head/block size, and a
prefill-containing batch outside verification-window and TurboQuant
modes. Initialization failures warn and fall back.
`VLLM_METAL_DISABLE_NAX=1` is the emergency override.

Correctness coverage:

- 60 GQA pipelines: two dtypes x five head sizes x three block sizes x
fast/masked regimes. Worst observed divergence remains below 0.9 ULP
max.
- Targeted MHA and MQA fast/masked cases pass.
- Fully masked rows remain finite and match tiled attention.
- Qwen3-0.6B greedy generation produces identical golden tokens with NAX
enabled and disabled.
- Prebuilt and source/JIT initialization paths pass.

Official macOS 15 CI and release jobs select Xcode 26.3. Required native
artifacts target macOS 15.0; the separate NAX metallib targets macOS
26.2. Wheel verification requires the NAX metallib and checks both
native extensions' deployment target.

GitHub's M1 runner compiles and packages NAX. The 60-pipeline execution
matrix remains a manual M5 release check.

</details>

---------

Signed-off-by: Ranran <ranranhaoranzhang@gmail.com>