I’ve been developing mlx-recurrence, a plug-in framework of fused Metal GPU kernels for linear recurrences on Apple silicon—roughly analogous to flash linear attention for MLX.
Sequential recurrences are difficult for MLX to fuse automatically. Architectures such as state-space models, gated linear attention, and diagonal RNNs ordinarily require a loop across the sequence length. When that loop is implemented in Python, a sequence of length L can require L separate Python-to-Metal dispatches.
These kernels instead execute the entire recurrence in a single Metal dispatch.
The training path uses segment checkpointing with recomputation during the backward pass. In validated M3 Max tests, the checkpoint-and-recompute kernels reduced peak recurrent-state memory by approximately 12–18× at the kernel level and lowered total training peak memory from 23.88 GB to 10.34 GB. At the same batch size, end-to-end training throughput improved by roughly 1.4×, while individual fused forward-and-backward kernels ran approximately 1.5–1.9× faster than the original full-state implementations. Results will vary with recurrence type, sequence length, state dimensions, batch size, datatype, model architecture, and hardware.
Current kernels:
ssd_scan
Mamba-2-style, head-wise SSD selective scan.
Intended for Mamba-2 and other SSM hybrid architectures.
State shape:
[B, H, Dh, N]
gla_scan
Gated Linear Attention with a scalar forget gate and outer-product write.
Intended for GLA and linear-attention hybrid architectures.
State shape:
[B, H, Dh, Dh]
rglru_scan
RG-LRU diagonal recurrence.
Intended for Griffin and RecurrentGemma-style architectures.
State shape:
[B, D]
rotlru_scan
Rotational LRU using a complex-diagonal recurrence, a magnitude gate, and a per-step rotation of two-dimensional channel pairs.
Intended for complex-LRU and S4-style oscillatory memory architectures.
State shape:
[B, D], represented as interleaved channel pairs.
Each kernel is implemented as a self-contained plug-in on a shared chassis located at:
mlx_recurrence._chassis
The chassis provides:
Segment checkpoint-and-recompute infrastructure
Shape and argument validation
VJP integration
Forward and gradient parity-test helpers
Common recurrence plug-in handling
Adding another recurrence therefore requires implementing its Metal forward and backward source pair and connecting its VJP. The checkpointing, validation, and testing infrastructure does not need to be rebuilt for each operator.
The original version 0.1 kernels remain available under:
mlx_recurrence.legacy
They are also re-exported at the package’s top level for backward compatibility.
I’m interested in feedback from developers working with MLX or custom Metal compute kernels, particularly around:
Preferred APIs for packaging reusable MLX Metal extensions
Threadgroup and memory-layout strategies across Apple GPU generations
Numerical stability expectations for long recurrent sequences
Benchmarking fused scans against MLX-native implementations
Additional recurrent operators that would be valuable to support
I would also be interested to know whether others are developing similar fused recurrence primitives for MLX and whether a shared interface for these operations would be useful.
My setup is a M3 MAX Macbook Pro with 36GB Ram and I am running on macOS 26.4.1 (25E253).