Add D256 to GQA-8 two-pass vector attention

[wyanzhao]

The GQA-8 two-pass kernel already shares K/V loads across query heads for D64 and D128. D256 can use the same kernel with two query heads per simdgroup: its float threadgroup arrays take 16,512 bytes, while four heads would exceed 32 KiB.

This adds that specialization for single-token decode with at least 8192 keys, matching query/value dimensions, no array mask and no sinks. It fits the D256, 16/2-head full-attention layers in Qwen3.5-35B-A3B when using an ordinary floating-point KV cache. Quantized KV cache uses a separate path.

The six measured cells below are faster on M5 Max. Other Metal GPUs and whole-model throughput have not been measured. Tests cover all three dtypes, odd key lengths, batch 2 and sliced KV. From python/tests, MLX_ENABLE_TF32=0 python -m unittest -v test_fast_sdpa passed (28 tests, 2 skipped); uvx pre-commit run --all-files passed.

Times are arm medians; ratios are paired geometric means of main/PR time, with 95% CIs (MLX_ENABLE_TF32=0). The table was measured at b04aea2a7 against 5778a97c0; merge commit 288906a50 was correctness-tested but not retimed.

dtype	B, query/KV heads	qL / kL	Main / PR (µs)	Paired ratio [95% CI]
float16	1, 16/2	1 / 8192	102.4 / 92.4	1.106 [1.099, 1.113]
float16	1, 16/2	1 / 32768	267.6 / 230.8	1.162 [1.158, 1.167]
bfloat16	1, 16/2	1 / 8192	97.5 / 87.7	1.109 [1.103, 1.114]
bfloat16	1, 16/2	1 / 32768	288.7 / 244.6	1.182 [1.174, 1.190]
float32	1, 16/2	1 / 8192	123.1 / 109.6	1.130 [1.120, 1.140]
float32	1, 16/2	1 / 32768	424.9 / 359.3	1.183 [1.172, 1.195]

Benchmark reproduction

Measured builds: main 5778a97c0, candidate b04aea2a7; 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 1 --k 8192
MLX_ENABLE_TF32=0 python sdpa_microbench.py --package /path/to/pr/python --q 1 --k 8192

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 30 fixed main/main calibration pairs followed by 30 main/PR pairs, alternating package order, with 60 seconds of preconditioning and 30 seconds between batches. The float32 cells used a separate session with 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. The float32 k32768 A/B used that retry. 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. All six comparisons passed calibration and drift checks.

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.