Chunkwise-parallel gated delta rule for training

[adityak74]

Problem

The gated delta layers fall back to gated_delta_ops whenever the Metal kernel is unavailable, and that includes every training step — use_kernel=not self.training, because the kernel has no VJP. gated_delta_ops is a sequential loop over timesteps, so autograd retains one state per token.

Measured on one layer (Hk=16, Hv=32, Dk=Dv=128), forward and backward, M4 Pro:

seq	peak memory	time
512	4.29 GB	0.36 s
1024	10.18 GB	1.16 s
2048	31.44 GB	4.20 s
4096	108.33 GB	39.02 s

About 8 MB per token per layer, growing linearly, and superlinear in time. Five model families route through this during training: qwen3_5, qwen3_next, kimi_k3, kimi_linear, bailing_moe_v3.

Change

gated_delta_ops becomes a thin dispatcher: scalar gating without a mask takes a new gated_delta_chunkwise, everything else goes to gated_delta_sequential, which is the existing loop under a name that says what it is. The existing kernel tests now compare against gated_delta_sequential so they keep testing what they tested before.

gated_delta_chunkwise is the chunkwise parallel form of the gated delta rule (arXiv:2406.06484, arXiv:2412.06464). Splitting the sequence into chunks of 64 and solving one unit lower-triangular system per chunk keeps O(T / chunk) states instead of O(T) and replaces the per-step loop with matmuls.

Same setup, with the new path:

seq	peak memory	time	vs sequential
512	0.33 GB	0.04 s	13x mem, 10x speed
1024	0.97 GB	0.09 s	10x mem, 12x speed
2048	1.74 GB	0.23 s	18x mem, 19x speed
4096	4.73 GB	0.60 s	23x mem, 65x speed

Inference is untouched: that path uses the Metal kernel and never reaches this code.

Scope

The dispatch prefers the chunkwise path only for scalar gating without a mask. Vectorized gating (g.ndim == 4) and masked sequences continue to use the reference loop, so the change cannot alter results for cases it does not cover.

Two implementation notes

mx.linalg.tri_inv is CPU-only and has no VJP, so it cannot appear in a training graph. The unit-triangular inverse is instead built from matmuls by 2x2 block recursion in log2(n) levels, which autodiff handles natively and which parallelizes across the batch — for the batched 64x64 case here that also measured faster than the LAPACK path (5.9 ms vs 7.4 ms for 1024 matrices, before counting the device round trip).

The textbook derivation divides by the cumulative decay G_j, which underflows and sends the intermediate writes to infinity on strongly decaying heads. This implementation carries bounded ratios G_j / G_i (at most 1 for i <= j) instead. The tests cover that regime explicitly.

Testing

tests/test_gated_delta.py adds a TestGatedDeltaChunkwise case covering:

• agreement with gated_delta_ops at mild (0.2) and strong (4.0) decay
• sequence lengths that are and are not a multiple of the chunk size
• grouped-query heads (Hv // Hk == 4)
• gradients through all five inputs

Outputs match to ~1e-6 relative in fp32 and gradients to ~1e-4. The existing kernel tests pass unchanged (11 passed, 32 subtests).

Verified to merge cleanly into current main, and the merged tree passes the suite (training dispatch agrees with the kernel to 1.5e-6).

[Goekdeniz-Guelmez]

this is already implemented in MLX-LM-LoRA which already provides chunked and efficient long context training for all linear models within MLX-LM. you can also manually change the chunk size via:

from mlx_lm_lora.trainer.sft_trainer import SFTTrainingArgs
args = SFTTrainingArgs(..., recurrence_chunk_size=512, ...)

works with the other algorithms like DPO, ORPO etc. too.

[adityak74]

Thanks for the pointer — I hadn't seen recurrent_patch.py. Having read it, I think the two are complementary rather than duplicates, and it's worth being precise about why.

enable_memory_safe_recurrences keeps _gated_delta_step_ops, the sequential per-token loop, and wraps each block of chunk_size steps in mx.checkpoint. That is activation checkpointing of the existing recurrence: memory falls to O(T / chunk) retained states, the arithmetic is unchanged, and the backward pass recomputes the token loop inside each block. Time stays superlinear in T, as in the table above (4.20 s at 2,048 → 39.02 s at 4,096 for one layer), and checkpointing adds a forward recompute on top.

This PR changes the arithmetic. gated_delta_chunkwise is the chunkwise-parallel form of the gated delta rule (arXiv:2406.06484 §4, arXiv:2412.06464): within a chunk the recurrence is solved as one unit-lower-triangular system and the intra- and inter-chunk contributions are matmuls, so there is no per-token loop in either direction. Same outputs to bf16 tolerance (the tests compare against gated_delta_sequential), O(T / chunk) states without recomputation, and the 135× at 4,096 in the table is from removing the loop, not from checkpointing it.

The two compose: your checkpoint wrapper could sit around gated_delta_chunkwise blocks just as it sits around the step loop now, which would give the memory of both and the speed of this one. Happy to add a benchmark against recurrent_patch at 2,048 and 4,096 if the maintainers want it for the record.