Adding metal kernels for the gated delta nets.

[tpegolotti]

Adding kernel support for the Gated Delta Net rule. The math follows the ICLR paper: https://arxiv.org/abs/2412.06464.

Changes

Added three kernels for the forward gated delta rule:

• Sequential: straightforward recurrence
• Simdgroup: chunk parallel implementation using 8x8 simdgroup_matrix operations
• NAX: chunk parallel implementation using 16x16 tiles available on the neural accelerator

Added benchmark and test scripts. The test script also compares against the naive FLA implementation of the gated delta rule using PyTorch ops.

Models Tested

Additionally to the tests added, I validated the kernels by running mlx_lm.evaluate --model <model> --tasks wikitext --num-shots 5 --max-tokens 2048 on the following models mlx-community/Qwen3.5-9B-MLX-4bit, mlx-community/Qwen3.5-35B-A3B-8bit, and mlx-community/Qwen3.5-27B-4bit. In the following table, I report the "word_perplexity" value.

	Sequential	Simdgroup	NAX
Qwen3.5-9B	9.82	9.82	9.82
Qwen3.5-35B	7.45	7.45	7.45
Qwen3.5-27B	7.87	7.87	7.87

As expected, they all match.

Performance

I give micro benchmarks on M1 Max and M5 Max. These can be obtained by running python benchmarks/python/gated_delta_bench.py.

As well as full end to end prompt processing measurements over the 3 models used for validation on the M5 Max.

Future work

• Add backward pass
• Add to mlx-lm
• Add masking to kernels
• Make kernel general in chunk parameter

[wyanzhao]

Hi @tpegolotti — thanks for this PR, and thanks @zcbenz for pointing me at it. The chunked WY-form kernels are a clear step up for batched prefill (the B ≥ 4 curves speak for themselves), and I'd be glad to see this land. @zcbenz asked me to look for possible improvements, so here are my notes. Everything below comes from reading the diff — I have not run the branch yet — so please take it as suggestions and correct me wherever I've misread something. I'm happy to help with any of it (patches against this branch, or measurements on an M5 Max).

1. gated_delta_seq layout (which is also the baseline of the speedup plots). This is the one-SIMD-group-per-value-row kernel from mlx-lm: at Dk = 128 each lane holds 4 state elements and every timestep does two full 32-lane simd_sums. In Add packed gated delta kernel, bitwise-pinned by an explicit-tree comparator mlx-lm#1559 I pack 8 value rows per SIMD-group (4 lanes per row, 32 contiguous state elements per lane, simd_shuffle_xor(1, 2) for the last two reduction levels); it is bitwise-identical to the current kernel and ~1.9–2.0× faster at B = 1, T = 2048–8192 on M5 (independently reproduced in that PR). Reading off your M5 Max plot, B = 1 sits at roughly 1.3–1.45× (simdgroup) and 1.6–2.2× (NAX) for T ≤ 2048, so the packed layout seems to be in the same range at B = 1 — that's a different harness though, so it needs a proper head-to-head before anyone relies on it, and I'd be glad to run that. Independently of the B = 1 question, the port is mechanical and would make the T ≤ 16 path (decode / short prefill) faster while keeping it bitwise-equal to today's kernel. I'd be happy to send it as a PR against your branch if you'd like.

2. Gate precision. gates / beta are cast to out_dtype before the primitive, so for bf16 models the decay is quantized to bf16 before the log + prefix sum. mlx-lm computes g in float32 today (compute_g), so this would be a small precision regression relative to the current mlx-lm path. Since g / beta are only [B, T, H], keeping them in fp32 inside the op looks cheap.

3. dtype guard. Only float and bfloat16_t are instantiated, and use_fallback doesn't look at the dtype, so I believe float16 inputs would look up seq_gated_delta_float16_t_… and fail at kernel lookup. Either instantiating float16_t or routing fp16 to the fallback would cover it.

4. Vector-gate guard. use_fallback doesn't check g.ndim. The fallback handles a [B, T, H, Dk] gate, but when the shape checks pass, such a gate reaches the kernel, which indexes g as [B, T, Hv] — so the result would be silently wrong rather than raising. (mlx-lm's kimi_linear does pass a vector gate; today it only stays out of the kernel because its Dk is 72.)

5. Head list. supported_heads = {(24,24), (32,32), (16,32), (16,48)} misses Qwen3.5-0.8B / 2B (Hk, Hv = 16, 16) and Qwen3.5-122B-A10B / 397B-A17B (16, 64), which would then take the per-timestep fallback (roughly 20–25 graph ops per timestep, i.e. tens of thousands of nodes per layer at T = 2048 — much slower than the current mlx-lm kernel). As far as I can see Hk / Hv only feed strides and the GQA index, so they could be runtime arguments; only Dk / Dv need to be compile-time.

6. Masks take the same fallback. mlx-lm's batched generation passes padding masks, so I think the masked path needs a kernel before mlx-lm can switch over wholesale (the existing mlx-lm masked kernel is a ready reference, and I'm happy to help there too).

7. Smaller things:

   • fill_gpu zero-fills out on every call, but all three kernels write every element, so that's an extra pass + dispatch per layer per token.
   • std::getenv / std::stoi run inside eval_gpu on the decode hot path; caching them in a static would avoid that.
   • The gamma[] threadgroup hand-off (lanes < C write, all lanes read) has no simdgroup_barrier; a simd_shuffle of gamma_val would avoid threadgroup memory altogether.
   • metal::fast::log / exp plus the max(g, 1e-6) clamp change the numerics relative to the sequential kernel, and the NAX test needs atol=1e-1 where the others pass at 1e-4. It might be worth a token-identity check on a real model (greedy decode, sequential vs chunked) in addition to the wikitext perplexity.
   • Docs: the Python docstring shapes are transposed relative to the implementation ([B, H, T, Dk] / [B, H, Dk, Dv] vs actual [B, T, H, Dk] / [B, Hv, Dv, Dk]), gamma is undocumented (it's the linear-space decay, not its log), state() is still a TODO, LOAD_M / LOAD_MT are defined twice in gated_delta_update_impl.h, and there's a duplicated return false; in use_fallback.

Again, happy to turn any of these into patches or measurements — just say which would be most useful. (Also glad to help with the rebase onto current main if that's useful.)

[tpegolotti]

Hi @wyanzhao, thanks for the comments and suggestions! I've added most of them in the latest commits.

1. Layout: Yes, please send your PR.
2. Gate precision: Removing the gamma / beta casts broke the PPL evaluation.
3. Dtype: Added the float16_t instantiation
4. Gate guard: Inside the fallback, I expand the dimensions with a check on the number of dimensions. Is this what you were looking for?
5. Head list: Added the dimensions for the additional Qwen models.
6. Masks: Masks are currently supported only in the fallback. It's in the TODOs for the kernels.
7. Smaller things:
   • Remove the fill_gpu in the eval_gpu call
   • Reading the environment variable is still needed to swap kernels dynamically for testing purposes. I will remove it in the future.
   • Added a simdgroup barrier after initializing gamma.
   • The NAX test has lower tolerance due to the relaxed precision of the matrix multiplication, whereas all other kernels compute the intermediate results in float32.
   • Updated the docs and removed the duplicate code.