Use NAX attention for short causal D256 prefill

[wyanzhao]

Causal fp16/bf16 D256 prefill with 512–1023 query rows currently defaults to unfused attention. This extends the existing NAX split-head-dimension kernel to fp16/bf16 chunks with 512–1023 query rows and at most 1536 keys, where the measurements below favor it. Query/value dimensions must match and query length must not exceed key length.

This window covers the first few 512-token prefill chunks in models such as Qwen3.5-35B-A3B, whose full-attention layers use D256 and 16/2 query/KV heads. It requires an ordinary floating-point KV cache; full-size 2048-token chunks and quantized KV cache do not benefit.

Apple M5 Max, MLX_ENABLE_TF32=0. Times are arm medians; ratios are paired geometric means of main/PR time, with 95% CIs. These are operator measurements, not whole-model throughput. The six window cells were measured at runtime bb0b8a33e; final head 69e88a18a changes only tests and was used for the two control measurements in a separate session.

dtype	B, query/KV heads	qL / kL	Main / PR (µs)	Paired ratio [95% CI]
float16	1, 16/2	512 / 512	213.3 / 170.6	1.237 [1.224, 1.250]
float16	1, 16/2	512 / 1536	487.9 / 460.3	1.067 [1.048, 1.087]
float16	1, 16/2	1023 / 1536	1053.7 / 775.1	1.376 [1.318, 1.438]
bfloat16	1, 16/2	512 / 512	227.9 / 174.4	1.317 [1.304, 1.331]
bfloat16	1, 16/2	512 / 1536	474.2 / 428.8	1.105 [1.084, 1.127]
float16	2, 6/2	768 / 1024	556.9 / 430.3	1.292 [1.267, 1.318]
float16	1, 16/2	512 / 1537	505.9 / 507.7	0.994 [0.973, 1.015]
float32	1, 16/2	512 / 512	542.0 / 543.4	0.997 [0.993, 1.001]

The k1537 and float32 controls both passed calibration and drift checks. Neither showed a significant difference, and both 95% intervals lie within the predeclared main/PR ratio band of 0.95–1.05. This bounds these two controls; other head/batch/sinks combinations remain unmeasured.

Tests cover both dtypes, boundary lengths, sliced KV, batch 2 and sinks. From python/tests, MLX_ENABLE_TF32=0 python -m unittest -v test_fast_sdpa passed (27 tests, 2 skipped). uvx pre-commit run --all-files passed.

Benchmark reproduction

Measured builds: main 5778a97c0, candidate runtime bb0b8a33e (controls at test-only final head 69e88a18a); Apple M5 Max, 128 GiB, macOS 27.0 (26A5425a). Build separate source checkouts with identical Release settings and Python bindings under each checkout's python/ directory. Save the script below as sdpa_microbench.py.

MLX_ENABLE_TF32=0 python sdpa_microbench.py --package /path/to/main/python --q 512 --k 512 --causal
MLX_ENABLE_TF32=0 python sdpa_microbench.py --package /path/to/pr/python --q 512 --k 512 --causal

Set --dtype, --q, --k, --hq, --hk and --batch for each row. The script reports seconds per call using a four-call dependent chain. Keep other GPU work idle; thermal-limit telemetry was unavailable during these measurements.

Use 15 fixed main/main calibration pairs followed by 15 main/PR pairs, except fp16 q512/k1536 and both controls use 30 of each. Alternate package order. Initial preconditioning/cooldown settings were 10s/10s, except fp16 q512/k512 used 60s/10s and fp16 q512/k1536 and each control used 60s/30s after five minutes of initial cooling.

Reject arm-median drift above 5% between session halves; allow one retry after 120 seconds of cooling with doubled preconditioning. Retain all pairs, including flagged outliers, in the paired log-ratio mean and Student-t 95% interval. Do not pool sessions; treat overlap with the calibration interval as unresolved.

import argparse
import statistics
import sys
import time
from pathlib import Path

p = argparse.ArgumentParser()
p.add_argument('--package', required=True, help='Source build python/ directory')
p.add_argument('--dtype', default='float16')
p.add_argument('--q', type=int, required=True)
p.add_argument('--k', type=int, required=True)
p.add_argument('--batch', type=int, default=1)
p.add_argument('--hq', type=int, default=16)
p.add_argument('--hk', type=int, default=2)
p.add_argument('--causal', action='store_true')
a = p.parse_args()
a.package = str(Path(a.package).resolve())
sys.path.insert(0, a.package)
import mlx.core as mx

assert a.package in mx.__file__, mx.__file__
print("BUILD", mx.__file__)

mx.set_default_device(mx.gpu)
mx.random.seed(0)
dtype = getattr(mx, a.dtype)
q = mx.random.normal((a.batch, a.hq, a.q, 256)).astype(dtype)
k = mx.random.normal((a.batch, a.hk, a.k + 32, 256)).astype(dtype)[:, :, :a.k]
v = mx.random.normal((a.batch, a.hk, a.k + 32, 256)).astype(dtype)[:, :, :a.k]
mx.eval(q, k, v)

def chain():
    x = q
    for _ in range(4):
        x = mx.fast.scaled_dot_product_attention(
            x, k, v, scale=1 / 16, mask='causal' if a.causal else None
        )
    mx.eval(x)

for _ in range(8):
    chain()
mx.synchronize()
samples = []
for _ in range(5):
    start = time.perf_counter()
    for _ in range(64):
        chain()
    mx.synchronize()
    samples.append((time.perf_counter() - start) / (64 * 4))
print('RESULT', statistics.median(samples))

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: AI assistance was used in preparing this contribution. I am responsible for the contribution.

[RohanGautam]

Thanks for the PR! Besides the comments below - could you give some end to end prefill numbers using mlx-lm before and after your change too? Once adressed, I'll run it on some more M5 machines to ensure this improves performance more generally.

[RohanGautam]

How were 512/1534 chosen? I understand they could come from the qwen, but we could do a sweep over the shapes and find the point where it the fusion starts being useful, so its more generalisable maybe.

[RohanGautam]

From what I understand, this test would pass regardless of this PR or not, right? What is the reason this test is introduced? Seems to be testing numerics which would be correct regardless of the routing changes