Pack eight value rows per SIMD-group in gated_delta_seq

[wyanzhao]

The sequential gated delta kernel now packs eight value rows into each SIMD-group, using four lanes per row. This uses the layout from mlx-lm#1559. GATED_DELTA_PACKED=0 selects the original kernel for comparison.

This applies to supported sequential calls with Dk=128 and Dv divisible by 8. It keeps the existing public shape support and chunk-8/NAX-16 dispatch. For the long-sequence measurements, I set GATED_DELTA_CHUNK=0: M5 would normally use NAX-16 at those lengths.

On M5 Max with (Hk, Hv)=(16,32) and Dk=Dv=128, I got the results below. Speedups are paired geometric means of original time / packed time in one binary. Timing includes public-API graph construction, dispatch, GPU execution and both outputs.

Case	Speedup	95% CI
default, B=8, T=8, float32	1.4144×	[1.4097, 1.4192]
sequential, B=1, T=2048, float32	1.2481×	[1.2394, 1.2568]
sequential, B=8, T=2048, float32	1.8678×	[1.8439, 1.8920]
sequential, B=16, T=2048, float32	1.8877×	[1.8613, 1.9144]
sequential, B=1, T=2048, float16	1.9495×	[1.9380, 1.9610]
sequential, B=1, T=2048, bfloat16	1.9708×	[1.9632, 1.9783]

These are the first session's results. I repeated the measurements in a second session and all six improvements were significant again. Default B1/T1 and B1/T8 fp32, plus the unchanged chunk-8/NAX-16 controls at B1/T2048/fp32, showed no significant difference. I haven't measured model performance.

test_fast_gated_delta passed 6 tests with 1 Torch-related skip; test_fast passed 30 with 1 skip. I used the macOS 26.5 SDK for the build and tests. Building with the macOS 27 SDK failed in unchanged NAX code against its MPP headers.

Dispatch, validation and reproduction

The public fast path still requires GPU, no mask, Dk=Dv=128, and one of (Hk,Hv)=(24,24), (32,32), (16,16), (16,32), (16,48), (16,64). Both GATED_DELTA_CHUNK=0 and =1 select sequential execution. The packed kernel uses grid (32, Dv/8, B*Hv) and threadgroup (32,2,1); the original uses (32,Dv,B*Hv) and (32,4,1).

The regression test compares both outputs for all six head pairs, fp32/fp16/bf16 and both sequential chunk settings, evaluating each result before changing the environment. Bitwise equality is toolchain-dependent and is not required by the test. All pre-commit hooks passed.

Measured commit: 7fb64079288670f88e688af868767a8f1a03a10d. M5 Max, 128 GiB, macOS 27.0 (26A5425a), Python 3.12.13, macOS 26.5 SDK. Binary: libmlx.dylib:4c196a02d790+mlx.metallib:4a03a40e46ac.

Each cell had passing A/A calibration, followed by 15 fixed alternating A/B pairs in each of two sessions. Timing used 60 seconds of initial saturation, at least 90 seconds of cooldown and a 5% within-arm drift limit with one cooled retry. Both sessions passed drift checks; the six reported effects passed the sign-test, Wilcoxon and A/A-noise checks. The sessions are separate estimates, not pooled. This compares the preserved original kernel with the packed kernel in one PR binary, not two separately built checkouts.

With a source build of this branch using the macOS 26.5 SDK, run from the
checkout root:

python python/tests/test_fast_gated_delta.py -v
python python/tests/test_fast.py -v
uvx pre-commit run --all-files

Save the script below as repro_4409.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_4409.py --source . --smoke --output smoke.json
python repro_4409.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 the six reported rows and four controls. It fixes the input seed to 20260905, checks both returned outputs, and records the chunk setting and invocation count for every case. Build with the macOS 26.5 SDK, as above.

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 #4409 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 = ['7fb64079288670f88e688af868767a8f1a03a10d']
CASES = {'control_auto_b1_t2048_f32': {'b': 1,
                               't': 2048,
                               'dtype': 'float32',
                               'chunk': 'auto',
                               'calls': 128},
 'control_chunk8_b1_t2048_f32': {'b': 1,
                                 't': 2048,
                                 'dtype': 'float32',
                                 'chunk': 8,
                                 'calls': 128},
 'default_b1_t1_f32': {'b': 1, 't': 1, 'dtype': 'float32', 'chunk': 'auto', 'calls': 2048},
 'default_b1_t8_f32': {'b': 1, 't': 8, 'dtype': 'float32', 'chunk': 'auto', 'calls': 2048},
 'default_b8_t8_f32': {'b': 8, 't': 8, 'dtype': 'float32', 'chunk': 'auto', 'calls': 2048},
 'seq_b16_t2048_f32': {'b': 16, 't': 2048, 'dtype': 'float32', 'chunk': 0, 'calls': 8},
 'seq_b1_t2048_bf16': {'b': 1, 't': 2048, 'dtype': 'bfloat16', 'chunk': 0, 'calls': 128},
 'seq_b1_t2048_f16': {'b': 1, 't': 2048, 'dtype': 'float16', 'chunk': 0, 'calls': 128},
 'seq_b1_t2048_f32': {'b': 1, 't': 2048, 'dtype': 'float32', 'chunk': 0, 'calls': 128},
 'seq_b8_t2048_f32': {'b': 8, 't': 2048, 'dtype': 'float32', 'chunk': 0, 'calls': 16}}

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):
    b, length, dtype = c["b"], c["t"], getattr(mx, c["dtype"])
    mx.random.seed(20260905)
    q = mx.random.normal((b, length, 16, 128))
    q = (q / (mx.linalg.norm(q, axis=-1, keepdims=True)+1e-6) / 128**.5).astype(dtype)
    k = mx.random.normal((b, length, 16, 128))
    k = (k / (mx.linalg.norm(k, axis=-1, keepdims=True)+1e-6)).astype(dtype)
    v = mx.random.normal((b, length, 32, 128)).astype(dtype)
    g = mx.exp(-mx.random.uniform(shape=(b, length, 32))*.2).astype(dtype)
    beta = mx.random.uniform(shape=(b, length, 32)).astype(dtype)
    state = (mx.random.normal((b, 32, 128, 128))*.1).astype(mx.float32)
    mx.eval(q, k, v, g, beta, state)
    def select(tag):
        os.environ["GATED_DELTA_THRESH"] = "16"
        if c["chunk"] == "auto":
            os.environ.pop("GATED_DELTA_CHUNK", None)
        else:
            os.environ["GATED_DELTA_CHUNK"] = str(c["chunk"])
        os.environ["GATED_DELTA_PACKED"] = "0" if tag == "off" else "1"
    def call():
        return mx.fast.gated_delta_update(q, k, v, g, beta, initial_state=state)
    outputs = []
    for tag in ("off", "on"):
        select(tag)
        value = call()
        mx.eval(*value)
        if not all(bool(mx.all(mx.isfinite(x)).item()) for x in value):
            raise RuntimeError("Non-finite output or final state")
        outputs.append(value)
    if not all(bool(mx.allclose(a, b, atol=1e-4, rtol=1e-4).item()) for a, b in zip(*outputs)):
        raise RuntimeError("Original/packed output or final state differs")
    def sample(tag, short=False):
        select(tag)
        for _ in range(4):
            mx.async_eval(*call())
        mx.synchronize()
        count = 1 if short else c["calls"]
        start = time.perf_counter()
        for _ in range(count):
            mx.async_eval(*call())
        mx.synchronize()
        return (time.perf_counter()-start) / count
    return sample, dict(finite=True, allclose=True, atol=1e-4, rtol=1e-4,
                        bitwise_equal=all(bool(mx.array_equal(a,b).item()) for a,b in zip(*outputs)))


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.

[tpegolotti]

Thanks for the PR! The performance is indeed faster. The only open question is if we should leave the previous kernel for the cases where Dk != 128.