Skip to content
D-CSIL

Project File

mlx-recurrence

Fused Metal GPU kernels for linear recurrence on Apple Silicon: Mamba-2 selective scan, gated linear attention, RG-LRU, and rotational LRU — with full training support.

OperationalThe kernels that make local training fast enough to matter.

Overview

MEASURED

Two generations of custom kernels. v2 ships four fused scans — ssd_scan (Mamba-2 MIMO), gla_scan (gated linear attention), rglru_scan (Griffin-style), rotlru_scan (complex rotation) — each with custom VJPs so they train, not just infer. The v2_evolve branch in the D-CSIL repo adds chunk-parallel GLA (128K+ sequences) and a fused projection+SiLU+scan super-kernel, all with numerical gradient tests.

Measured Performance

MEASURED

Against a Python per-step loop on M3 Max at seq 2048: 7.3× (SSM) and 9.1× (GLA) forward; 19.0× and 31.8× forward+backward. v2 cut peak kernel memory 12–18× versus v0.1 at training shapes.

Production Validation

REPLICATED

The kernels were hot-swapped into a live D-CSIL SSM+GLA training run mid-flight: peak memory dropped from 23.88GB to 10.34GB (−57%), throughput rose from ~1,074 to ~1,500 tok/s (+40%), gradient parity held at ~1e-7 relative error, and the loss curve continued cleanly.

Current Limitations

  • −Apple Silicon only (MLX ≥0.22); no CUDA path by design.
  • −rotlru and rglru variants have fewer production miles than ssd/gla.

Next Steps

  • +Serve as the kernel substrate for the HibbieFormer-HX helical core.
  • +Upstream-friendly packaging of the v2_evolve primitives.

Experiment Log