Fused Metal kernels for SSM (Mamba) and GLA recurrence — MLX-Recurrence

[D-CSIL]

MLX currently lacks native scan/recurrence primitives, which means
SSM (Mamba) and GLA (Gated Linear Attention) models have to fall back
to Python for-loops for the sequential scan — a major bottleneck that
makes training these architectures on Apple Silicon impractical.

I built MLX-Recurrence to fix this. Fused Metal GPU kernels for two
recurrence patterns using mx.fast.metal_kernel:

• Selective Scan (Mamba-style SSM) — forward + backward
• Gated Linear Attention (GLA) — forward + backward

Both include custom VJP implementations so they work with
nn.value_and_grad for training — not just inference.

At seq_len=2048 on M3 Max, kernel-level benchmarks:

• SSM forward+backward: 19x faster (64ms vs 1,225ms)
• GLA forward+backward: 31.8x faster (56ms vs 1,787ms)

End-to-end training throughput improvement is ~3x, as the scan
kernels are one component of the full training step.

Repo: https://github.com/D-CSIL/mlx-recurrence
Install: pip install git+https://github.com/D-CSIL/mlx-recurrence.git

Would love to see native scan primitives in MLX eventually — happy to
take feedback from the MLX team on kernel optimization and share what
I learned about Metal kernel constraints along the way.

For context: as far as I can find, this is the first implementation of
fused Metal backward kernels for any recurrence pattern on MLX. The
existing Mamba/SSM implementations in mlx-lm and other projects use
Metal kernels for inference only — training falls back to pure MLX ops
or Python loops.

MLX-Recurrence provides full forward + backward Metal kernels with VJP
support for both SSM selective scan and GLA. The GLA implementation
appears to be the first on the MLX platform entirely.