Use NAX attention for long unmasked D72/D80 inputs

[wyanzhao]

Padding D72 and D80 inputs to D96 lets them use the NAX attention kernel. The output is then copied back to its original width. On M5 Max, the four cases below are 1.24–1.99× faster, including those copies.

Potential users include the vision encoders in LFM2.5-VL and Qwen3.5-9B, 27B, and 35B-A3B, which use head dimension 72. Their MLX-VLM implementations (LFM2.5-VL, Qwen3.5) use unmasked attention, so fp16/bf16 inputs with at least 512 tokens per vision attention call on NAX devices can take this path.

Padding was slower for short sequences in earlier tests, so I kept their existing path. The new default applies on NAX devices to fp16/bf16 with query and key lengths ≥512, no mask and no sinks.

I benchmarked batch size 1 with 8 query / 2 KV heads, no mask or sinks. Speedups are paired geometric means of padding-off time / padding-on time.

Head dimension	Dtype	Query/key length	Speedup	95% CI
72	fp16	512	1.2414×	[1.2215, 1.2616]
72	bf16	1024	1.8566×	[1.8445, 1.8687]
80	fp16	1024	1.9863×	[1.9741, 1.9986]
80	bf16	512	1.3209×	[1.2965, 1.3459]

test_fast_sdpa and test_fast ran 57 tests: 54 passed and 3 skipped, with TF32 disabled. The opt-in test requires Metal and the GPU as the default device. It sets the internal padding override to 1 for the D72/D80 reference checks, including masks and sinks, then restores the previous environment value.

These measurements are from 8a516f9c8. The latest commit, 96a540459, only changes documentation and tests; I have not rerun the timings. I haven't measured model performance.

Tests and benchmark reproduction

Device: M5 Max, 128 GB, macOS 27.0. Measured commit: 8a516f9c80b35ef4a207ecaa3fb1a93358b9d5dd. Binary: libmlx.dylib:7e284cdf23f0+mlx.metallib:1e4de3f525d3.

Each timed cell passed a 15-pair A/A calibration and the 5% drift check before its A/B result was accepted. Dispatch observations confirmed native D72/D80 with padding off and NAX D96 with padding on. The equivalence tests alone do not prove which route ran.

To run the updated tests, build 96a54045904b4915705b2f17bb92ac0aac03997a from source:

test "$(git rev-parse HEAD)" = 96a54045904b4915705b2f17bb92ac0aac03997a
python -m pip install -e .
cd python/tests
MLX_ENABLE_TF32=0 python -m unittest -v test_fast_sdpa test_fast

For benchmark reproduction, use a separate clean checkout of the measured commit, 8a516f9c80b35ef4a207ecaa3fb1a93358b9d5dd, and build it with python -m pip install -e .. Save the script below as repro_4455.py in that checkout root. Use the same Python environment used to build it, and stop other GPU work before measuring. The script uses the internal MLX_SDPA_PAD_HEAD_DIM override to compare padding off and on.

# Run from the measured-commit checkout root after building it.
python repro_4455.py --source . --smoke --output smoke.json
python repro_4455.py --source . --case all --output results.json

--smoke checks outputs and exercises short timing loops without estimating speedups. Use --help to select a single case. Output files must be new names.

The script covers all four rows, using seed 187 and 48 batches of four dependent calls per sample. It includes the padding and output copies. OFF/ON output comparison checks numerical agreement; it does not replace the pipeline observations above. Padding also needs temporary Q/K/V and output buffers.

This is an independent reproducer of the workloads, not the driver used for the existing table. It warms both arms, runs 15 alternating A/A pairs, then 15 alternating A/B pairs. Each stage gets 60 seconds of saturation, with a 90-second cooldown before A/B. It reports the geometric mean of OFF/ON time ratios and a Student-t 95% interval in log space. A/A must include 1 in its interval; either stage stops on more than 5% within-arm drift, retaining the rejected result without retrying. It does not implement the original sign-test, Wilcoxon or A/A-noise significance checks.

The JSON includes raw sample times, input shapes, numerical checks, source revision, device and binary hashes. New estimates depend on the machine and build; they do not replace the measurements above.

"""Independent MLX PR #4455 reproducer. Requires a built PR checkout and mlx.

Uses public MLX APIs only. Fixed 15-pair estimates are new observations, not
a replay of the original benchmark driver. Stop other GPU work before running.
"""

HEADS = ['8a516f9c80b35ef4a207ecaa3fb1a93358b9d5dd']
CASES = {'d72_float16_l512': {'b': 1,
                      'h': 8,
                      'hk': 2,
                      'ql': 512,
                      'kl': 512,
                      'd': 72,
                      'dtype': 'float16'},
 'd72_bfloat16_l1024': {'b': 1,
                        'h': 8,
                        'hk': 2,
                        'ql': 1024,
                        'kl': 1024,
                        'd': 72,
                        'dtype': 'bfloat16'},
 'd80_float16_l1024': {'b': 1,
                       'h': 8,
                       'hk': 2,
                       'ql': 1024,
                       'kl': 1024,
                       'd': 80,
                       'dtype': 'float16'},
 'd80_bfloat16_l512': {'b': 1,
                       'h': 8,
                       'hk': 2,
                       'ql': 512,
                       'kl': 512,
                       'd': 80,
                       'dtype': 'bfloat16'}}
SWITCH="MLX_SDPA_PAD_HEAD_DIM"
VALUES={"off":"0", "on":"1"}
SEED=187
DIRECT_DTYPE=False
OUTER=48
CHAIN=4

import argparse
import hashlib
import json
import math
import os
from pathlib import Path
import platform
import statistics as st
import subprocess
import time


def summary(a, b):
    if len(a) != 15 or len(b) != 15 or any(
        not math.isfinite(x) or x <= 0 for x in a + b
    ):
        raise ValueError("Expected 15 positive finite samples in each arm")
    logs = [math.log(x / y) for x, y in zip(a, b)]
    center = st.mean(logs)
    # Two-sided Student-t 95% interval, df=14; fixed 15 pairs, no early stopping.
    half = 2.1447866879169273 * st.stdev(logs) / math.sqrt(15)
    drift = [st.median(v[7:]) / st.median(v[:7]) for v in (a, b)]
    return dict(ratio=math.exp(center), ci95=[math.exp(center-half), math.exp(center+half)],
                drift=drift, drift_ok=all(abs(x-1) <= .05 for x in drift))


def collect(sample, tags):
    values = [[], []]
    for i in range(15):
        for j in ((0, 1) if i % 2 == 0 else (1, 0)):
            values[j].append(sample(tags[j]))
    return dict(samples_seconds=values, **summary(*values))


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--source", type=Path, required=True, help="built PR checkout")
    p.add_argument("--case", choices=["all", *CASES], default="all")
    p.add_argument("--output", type=Path, required=True)
    p.add_argument("--smoke", action="store_true", help="check outputs and short timing loops; no performance verdict")
    args = p.parse_args()
    source = args.source.resolve(strict=True)
    head = subprocess.check_output(["git", "-C", str(source), "rev-parse", "HEAD"], text=True).strip()
    if head not in HEADS:
        p.error(f"Expected source head in {HEADS}; found {head}")
    dirty = subprocess.check_output(["git", "-C", str(source), "status", "--porcelain", "--untracked-files=no"], text=True)
    if dirty.strip():
        p.error("Use a clean source checkout")
    if os.environ.get("DYLD_INSERT_LIBRARIES"):
        p.error("Remove profiling/dispatch instrumentation before timing")
    # Bind the requested source build instead of silently using another pip package.
    import sys
    sys.path.insert(0, str(source / "python"))
    os.environ["MLX_ENABLE_TF32"] = "0"
    import mlx.core as mx
    if not Path(mx.__file__).resolve().is_relative_to(source):
        p.error(f"MLX imported from the wrong build: {mx.__file__}")
    if not mx.metal.is_available():
        p.error("A Metal GPU is required")
    mx.set_default_device(mx.gpu)
    core = Path(mx.__file__).resolve()
    binaries = [core, core.parent / "lib/libmlx.dylib", core.parent / "lib/mlx.metallib"]
    def fingerprints():
        return {f.name: hashlib.sha256(f.read_bytes()).hexdigest() for f in binaries}
    result = dict(source_head=head, device=mx.metal.device_info(), os=platform.platform(),
                  python=platform.python_version(), binaries=fingerprints(),
                  script_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
                  protocol="Independent fixed-15-pair reproduction, Student-t log CI, 5% drift gate; no retry or sequential stopping",
                  smoke=args.smoke, cases={})
    # Refuse to overwrite an earlier run, including a rejected calibration.
    with args.output.open("x") as output:
        def save():
            output.seek(0)
            json.dump(result, output, indent=2)
            output.write("\n")
            output.truncate()
            output.flush()
        save()
        try:
            for name in CASES if args.case == "all" else [args.case]:
                row = result["cases"][name] = dict(config=CASES[name], status="checking")
                sample, check = workload(mx, CASES[name])
                row["check"] = check
                if args.smoke:
                    sample("off", short=True)
                    sample("on", short=True)
                    row["status"] = "smoke_only_no_performance_result"
                else:
                    for tag in ("off", "on"):
                        for _ in range(3):
                            sample(tag)
                    for stage, tags in (("aa", ("on", "on")), ("ab", ("off", "on"))):
                        if stage == "ab":
                            time.sleep(90)
                        end = time.monotonic() + 60
                        while time.monotonic() < end:
                            for tag in tags:
                                sample(tag)
                        row[stage] = collect(sample, tags)
                        s = row[stage]
                        ok = s["drift_ok"] and (stage != "aa" or s["ci95"][0] <= 1 <= s["ci95"][1])
                        row["status"] = "in_progress" if ok else f"rejected_{stage}"
                        save()
                        if not ok:
                            raise RuntimeError(f"{name}: rejected {stage}; retain this output, do not quote its effect")
                    row["status"] = "completed_independent_reproduction"
                if fingerprints() != result["binaries"]:
                    raise RuntimeError("Binary files changed during the run")
                save()
                print(name, row["status"], row.get("ab", {}).get("ratio", ""), flush=True)
                del sample
        except BaseException as error:
            result["error"] = f"{type(error).__name__}: {error}"
            save()
            raise


def workload(mx, c):
    dtype = getattr(mx, c["dtype"])
    b, h, hk, ql, kl, d = (c[k] for k in ("b", "h", "hk", "ql", "kl", "d"))
    mx.random.seed(SEED)
    def normal(shape):
        x = mx.random.normal(shape, dtype=dtype) if DIRECT_DTYPE else mx.random.normal(shape).astype(dtype)
        return x
    q, k, v = normal((b, h, ql, d)), normal((b, hk, kl, d)), normal((b, hk, kl, d))
    sinks = 10 * normal((h,)) if c.get("sinks") else None
    mx.eval(q, k, v, *([] if sinks is None else [sinks]))
    def select(tag):
        os.environ[SWITCH] = VALUES[tag]
    def call(x, force=False):
        return mx.fast.scaled_dot_product_attention(x, k, v, scale=d**-0.5,
                                                    sinks=sinks, force_fused=force)
    outputs = []
    for tag in ("off", "on"):
        select(tag)
        value = call(q)
        mx.eval(value)
        if not bool(mx.all(mx.isfinite(value)).item()):
            raise RuntimeError("Non-finite output")
        outputs.append(value)
    equal = bool(mx.array_equal(*outputs).item())
    atol, rtol = (1e-4, 1e-4) if c["dtype"] == "float32" else ((.005, .005) if d != 512 else (.02, 0.0))
    if not bool(mx.allclose(*outputs, atol=atol, rtol=rtol).item()):
        raise RuntimeError("OFF/ON numerical comparison failed")
    if d == 512:
        select("on")
        mx.eval(call(q, force=True))
        select("off")
        try:
            mx.eval(call(q, force=True))
        except ValueError as e:
            if "force_fused=True but no fused kernel is available" not in str(e):
                raise
        else:
            raise RuntimeError("D512 fallback unexpectedly admitted a fused kernel")
    elif equal:
        # Conservative guard: a non-NAX device/build can ignore this switch.
        # Different outputs are NOT a kernel-name witness.
        raise RuntimeError("Padding arms are identical; verify NAX engagement before benchmarking")
    def sample(tag, short=False):
        select(tag)
        mx.synchronize()
        start = time.perf_counter()
        outer = 1 if short else OUTER
        for _ in range(outer):
            x = q
            for _ in range(CHAIN):
                x = call(x)
            mx.eval(x)
        mx.synchronize()
        return (time.perf_counter()-start) / (outer*CHAIN)
    return sample, dict(finite=True, allclose=True, atol=atol, rtol=rtol, bitwise_equal=equal,
                        route_check="force_fused admission" if d == 512 else "output difference only; no pipeline witness")


if __name__ == "__main__":
    main()

---

• ☑️ 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.

[wyanzhao]

@zcbenz Thanks for the review. Addressed all three comments in 96a5404:

• Removed the environment-variable documentation.
• Made the opt-in test Metal-only and skip when the default device is CPU.
• Removed the repeated default/OFF comparisons. The test now sets MLX_SDPA_PAD_HEAD_DIM=1, reuses the D72/D80 reference checks—including masks and sinks—and restores the previous environment value.

test_fast_sdpa and test_fast: 54 passed, 3 skipped, with TF32 disabled. I also updated the PR description to reflect these changes.