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	Split D and Dv across two simdgroups
Routing	Causal, no array mask, qL >= 1024, NAX-compatible dtype/TF32 mode
Ragged inputs	Pad to aligned tiles internally, then slice padded query rows
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

The wn=2 configuration 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

For head_dim == 256, the unaligned pipelines (align_Q/align_K = false)
slow every block down by roughly 3x. Ragged causal inputs are therefore
padded to the tile sizes at dispatch:

• qL_off is computed from the true sequence lengths, preserving the
  causal diagonal;
• padded K/V rows are zero-filled;
• true query rows cannot attend padded key columns; and
• padded query rows are removed when copying the result back.

The padding/copy cost is approximately 0.6 ms, avoiding an observed
approximately 49 ms unaligned-path cliff.

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.

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.

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 aligned/ragged
  rectangles.
• 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: 18 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.