Pack gated_delta_seq: 8 value rows per SIMD-group

[wyanzhao]

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: The layout is the one from my mlx-lm#1559. I ported it into this kernel with AI assistance; an AI agent ran the builds, the bitwise A/B, and the benchmarks under my direction. I reviewed all of the code and the data, and I take full responsibility for the change.

Summary

This ports the packed sequential layout from
ml-explore/mlx-lm#1559
into gated_delta_seq on this branch, at @tpegolotti's invitation on
#4020. The change is one kernel plus the two dispatch lines in
eval_gpu: eight value rows share a SIMD-group (four lanes per row,
32 contiguous state elements per lane); the two full-SIMD simd_sums
become an in-lane three-level local tree plus simd_shuffle_xor(1, 2);
the grid / threadgroup go from (32, Dv, B*Hv) / (32, 4, 1) to
(32, Dv/8, B*Hv) / (32, 2, 1). static_assert(Dk == 128) and
Dv % 8 == 0 are added; the API, ABI, instantiations, chunk-8 kernel,
and NAX-16 kernel are untouched.

Bitwise equivalence

Each four-element partial keeps the original lane's sequential
accumulation, so the local work matches the unpacked per-lane loop.
IEEE addition is commutative, which lets the first three xor-tree
levels (1, 2, 4) fold into that in-lane tree without changing
bits; the last two levels (8, 16) become simd_shuffle_xor(1) and
simd_shuffle_xor(2) inside the four-lane row group. The association
is therefore the same butterfly the current simd_sum implements on
this toolchain.

Validation (M5 Max, two builds of this branch at 285be14 — base vs
patched):

• Cross-build bitwise A/B: 77/77 cases mx.array_equal on both
  y and the final state — heads {(24,24),(32,32),(16,32),(16,48)} ×
  {bf16, fp16, fp32} × T ∈ {1, 7, 16, 257, 1111} × B ∈ {1, 2} (B=2 for
  the (16,32) pair), plus chunk-8 and NAX-16 control cases (also
  bitwise-equal, i.e. untouched).
• python/tests/test_fast_gated_delta.py on the patched build: 6
  tests, OK (1 skipped: torch not installed).
• 20,000-case float32 simulation of the packed reduction vs the
  explicit xor-tree: 0 mismatches.
• uvx pre-commit run --all-files: all hooks pass, no reformatting.

Performance

Absolute latency drifted by up to ~90% across rounds on this chassis
(thermal), so the numbers below are within-round paired ratios, not
absolute times. Five interleaved rounds, with the starting side
swapped (three rounds base-first, two patched-first) to cancel order
bias.

benchmarks/python/gated_delta_bench.py, sequential kernel
(GATED_DELTA_CHUNK=0). First three rows: median (min–max) of
per-round paired ratios seq_base / seq_patched. Rows 4–5: range of
per-cell medians.

B	T	speedup
1	2048	1.90× (1.62–2.42)
8	2048	2.16× (2.05–2.37)
16	2048	1.86× (1.61–2.53)
1–16	512–1024	1.78–1.96×
8 / 16	8	1.36–1.43×
1	8	~1.0 (within noise)

Default T ≤ 16 path (decode / short prefill), separate interleaved
probe: B=1 T=1/4/16 faster by ~7–24%, B=8 T=16 ~15%; no case
regressed beyond noise.

This port speeds up the sequential kernel that is the denominator of
the speedup plots on #4020, so those figures will need to be
recomputed or annotated once this lands.

Reproduction

# two installs of this branch (base = unpatched, patched = this PR), then:
python gdn_seq_bitwise_ab.py dump  <dir>    # under the base build
python gdn_seq_bitwise_ab.py compare <dir>  # under the patched build
python benchmarks/python/gated_delta_bench.py --csv   # 3+ interleaved rounds per side

gdn_seq_bitwise_ab.py

"""Cross-build bitwise A/B for the packed gated_delta_seq port (mlx #4020).

Usage:
  Under the BASE build (285be149, unpatched):
      python gdn_seq_bitwise_ab.py dump /tmp/gdn_ab
  Under the PATCHED build (same commit + packed-seq patch):
      python gdn_seq_bitwise_ab.py compare /tmp/gdn_ab

dump saves inputs AND outputs per case; compare reloads the saved inputs,
recomputes, and requires mx.array_equal on y and the final state, so the
check is immune to RNG differences across builds.

Covers every (Hk, Hv) pair use_fallback admits at 285be149, all three
instantiated dtypes, decode/ragged/threshold/spill/long T, and B in {1, 2}.
Chunk-8 and NAX-16 cases are included as controls: the patch must not
change them either.
"""

import os
import sys

os.environ["GATED_DELTA_CHUNK"] = "0"  # default; overridden per case below

import mlx.core as mx

HEADS = [(24, 24), (32, 32), (16, 32), (16, 48)]  # use_fallback's list @285be149
DTYPES = {"bf16": mx.bfloat16, "fp16": mx.float16, "fp32": mx.float32}
TS = [1, 7, 16, 257, 1111]
DK = DV = 128


def cases():
    for hk, hv in HEADS:
        for dname in DTYPES:
            for t in TS:
                for b in (1, 2) if (hk, hv) == (16, 32) else (1,):
                    yield ("seq", "0", b, hk, hv, t, dname)
    # controls: chunked paths must be untouched by the patch
    for chunk in ("8", "16"):
        yield ("chunk" + chunk, chunk, 1, 16, 32, 257, "bf16")


def make_inputs(b, hk, hv, t, dtype, seed):
    mx.random.seed(seed)
    q = mx.random.normal((b, t, hk, DK))
    k = mx.random.normal((b, t, hk, DK))
    k = k / (mx.linalg.norm(k, axis=-1, keepdims=True) + 1e-6)
    q = (q / (mx.linalg.norm(q, axis=-1, keepdims=True) + 1e-6)) * DK**-0.5
    g = mx.exp(-mx.random.uniform(shape=(b, t, hv)) * 0.2)
    beta = mx.random.uniform(shape=(b, t, hv))
    v = mx.random.normal((b, t, hv, DV))
    h0 = (mx.random.normal((b, hv, DV, DK)) * 0.3).astype(mx.float32)
    q, k, v, g, beta = (x.astype(dtype) for x in (q, k, v, g, beta))
    mx.eval(q, k, v, g, beta, h0)
    return q, k, v, g, beta, h0


def run(mode, outdir):
    os.makedirs(outdir, exist_ok=True)
    failures = []
    for i, (tag, chunk, b, hk, hv, t, dname) in enumerate(cases()):
        os.environ["GATED_DELTA_CHUNK"] = chunk
        name = f"{tag}_B{b}_Hk{hk}_Hv{hv}_T{t}_{dname}"
        path = os.path.join(outdir, name + ".npz")
        ins = make_inputs(b, hk, hv, t, DTYPES[dname], seed=1000 + i)
        if mode == "dump":
            y, hf = mx.fast.gated_delta_update(*ins[:5], initial_state=ins[5])
            mx.eval(y, hf)
            mx.savez(path, q=ins[0], k=ins[1], v=ins[2], g=ins[3],
                     beta=ins[4], h0=ins[5], y=y, hf=hf)
            print("dumped", name)
        else:
            ref = mx.load(path)
            y, hf = mx.fast.gated_delta_update(
                ref["q"], ref["k"], ref["v"], ref["g"], ref["beta"],
                initial_state=ref["h0"])
            mx.eval(y, hf)
            ok_y = bool(mx.array_equal(y, ref["y"]).item())
            ok_h = bool(mx.array_equal(hf, ref["hf"]).item())
            status = "OK" if (ok_y and ok_h) else "MISMATCH"
            print(f"{status:8s} {name}  y_equal={ok_y} state_equal={ok_h}")
            if status != "OK":
                failures.append(name)
    if mode == "compare":
        print("\n%d case(s) failed" % len(failures) if failures
              else "\nALL CASES BITWISE EQUAL")
        sys.exit(1 if failures else 0)


if __name__ == "__main__":
    if len(sys.argv) != 3 or sys.argv[1] not in ("dump", "compare"):
        sys.exit(__doc__)
    run(sys.argv[1], sys.argv[2])

Out of scope (pre-existing, noted for completeness)

• use_fallback does not yet admit the newly instantiated (16,16) / (16,64)
  head pairs, so those instantiations are currently unreachable.
• No g.ndim guard: a 4-D (vector) gate that passes the shape checks would be
  read by the kernels as [B, T, Hv].