TL;DR
We replaced the chunked MLX recurrence path with fused Metal kernels mid-flight in a live SSM+GLA run: peak memory fell 23.88GB → 10.34GB, throughput rose ~1,074 → ~1,500 tok/s, and gradient parity held at ~1e-7. The loss curve never noticed.
Why hot-swap at all
Restarting a multi-week local training run to change kernels costs days. The v2 kernels (ssd_scan, gla_scan) carry custom VJPs and were bit-tested against the chunked reference, so the risky part wasn't correctness — it was doing the swap without disturbing optimizer state.
The numbers
At training shapes (B=3, L=512, H=12, Dh=64): SSD forward 2.30ms and forward+backward 17.34ms at 145MB peak — a 12× memory reduction over v0.1. GLA: 1.41ms / 12.06ms at 81MB (18× reduction). Against a naive Python loop, forward+backward speedups reach 31.8×.
swap at checkpoint boundary
verify: grad parity ~1e-7 (chunked vs fused)
result: 23.88GB -> 10.34GB peak, +40% tok/sWhat made it safe
Three gates, all mandatory: numerical agreement on forward outputs, gradient agreement against the chunked path, and loss-continuity monitoring across the swap boundary. All three passed; the run kept going.