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