Add D512 support to Metal vector attention

[wyanzhao]

D512 attention can now use the existing single-pass and two-pass Metal vector kernels. In the five cases below, this is 1.11–1.23× faster than the unfused fallback on M5 Max.

Potential users include Gemma 4 31B, 26B-A4B, and E2B, whose global attention layers use head dimension 512 and a query-to-KV head ratio of 8. Single-token decoding with at least 1024 keys and no array mask can take this path.

I've enabled it by default for one query, GQA factor 8, no array mask and at least 1024 keys. Other supported D512 vector shapes can use it with MLX_SDPA_D512_MIN_KL=0; setting the threshold above the key length selects the fallback.

I measured c6dc20d95. The final head 5528d75b9 only adds test changes and has the same runtime source, but I haven't retimed it.

All rows use D=512, one query, 32 query / 4 KV heads and 1024 keys. Speedups are paired geometric means of fallback time / fused time.

Case	Speedup	95% CI
B1, fp16	1.1105×	[1.1031, 1.1179]
B2, fp16	1.1572×	[1.1421, 1.1724]
B1, fp16, sinks	1.2095×	[1.1916, 1.2276]
B1, bf16	1.1125×	[1.1037, 1.1213]
B1, fp32	1.2335×	[1.2139, 1.2534]

With TF32 disabled, test_fast_sdpa passed 22 tests with 2 skipped. These cover both kernels and the default dispatch restrictions. I haven't measured model performance; the results here are for the five operator cases above.

Controls and reproduction

At D=512 and kL=1024, the excluded fp16 controls were: GQA4 (B1, qL1, hq32/hk8) 1.0400 [1.0179, 1.0627], with improvement magnitude unresolved; MHA (B1, qL1, hq8/hk8) 0.9555 [0.9382, 0.9732]; and multi-query GQA8 (B1, qL4, hq32/hk4) 1.0277 [1.0208, 1.0346]. They retain the fallback by default. Array-mask calibration failed the drift check, so no timing claim is made for it.

Device: M5 Max, 128 GB. The local release build adds 102,944 bytes to mlx.metallib. Each reported cell passed A/A calibration and the 5% drift check.

Measured binary: libmlx.dylib:1bec8a11f32a+mlx.metallib:ba939f89e04b, at c6dc20d9575bbde819ae731b73b49a53cafcf5b0. The final-head test build has the same library hashes.

Build the reviewed commit from source and run the tests:

test "$(git rev-parse HEAD)" = 5528d75b9b7b8c5818242abe6593ceac92401262
python -m pip install -e .
cd python/tests
MLX_ENABLE_TF32=0 python -m unittest -v test_fast_sdpa
MLX_ENABLE_TF32=0 DEVICE=cpu python -m unittest -v \
  test_fast_sdpa.TestFastSDPA.test_sdpa_vector_head_dim_512

Save the script below as repro_4459.py in the built checkout root. Use the same Python environment used to build MLX, and stop other GPU work before measuring.

# Run from the checkout root, after building the pinned commit above.
python repro_4459.py --source . --smoke --output smoke.json
python repro_4459.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 includes the five reported cases and the three excluded controls. It uses seed 0, 250 batches of eight dependent calls, and checks fused admission with force_fused. The table used adaptive A/B sampling (9–12 pairs); this script uses 15 fixed pairs.

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 #4459 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 = ['c6dc20d9575bbde819ae731b73b49a53cafcf5b0', '5528d75b9b7b8c5818242abe6593ceac92401262']
CASES = {'b1_fp16': {'b': 1,
             'h': 32,
             'hk': 4,
             'ql': 1,
             'kl': 1024,
             'd': 512,
             'dtype': 'float16',
             'sinks': False},
 'b2_fp16': {'b': 2,
             'h': 32,
             'hk': 4,
             'ql': 1,
             'kl': 1024,
             'd': 512,
             'dtype': 'float16',
             'sinks': False},
 'sinks_fp16': {'b': 1,
                'h': 32,
                'hk': 4,
                'ql': 1,
                'kl': 1024,
                'd': 512,
                'dtype': 'float16',
                'sinks': True},
 'b1_bf16': {'b': 1,
             'h': 32,
             'hk': 4,
             'ql': 1,
             'kl': 1024,
             'd': 512,
             'dtype': 'bfloat16',
             'sinks': False},
 'b1_fp32': {'b': 1,
             'h': 32,
             'hk': 4,
             'ql': 1,
             'kl': 1024,
             'd': 512,
             'dtype': 'float32',
             'sinks': False},
 'gqa4_fp16': {'b': 1,
               'h': 32,
               'hk': 8,
               'ql': 1,
               'kl': 1024,
               'd': 512,
               'dtype': 'float16',
               'sinks': False},
 'mha_fp16': {'b': 1,
              'h': 8,
              'hk': 8,
              'ql': 1,
              'kl': 1024,
              'd': 512,
              'dtype': 'float16',
              'sinks': False},
 'q4_fp16': {'b': 1,
             'h': 32,
             'hk': 4,
             'ql': 4,
             'kl': 1024,
             'd': 512,
             'dtype': 'float16',
             'sinks': False}}
SWITCH="MLX_SDPA_D512_MIN_KL"
VALUES={"off":"999999", "on":"0"}
SEED=0
DIRECT_DTYPE=True
OUTER=250
CHAIN=8

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.