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.