mlx-mamba3
Hey, Built this after hitting some dependency hell with the CUDA/Triton requirements on Colab. Covers SISO, MIMO, and Hybrid Attention-Mamba configs with verified numerical parity against the PyTorch reference (max error < 10⁻⁵, 12 passing tests).
Key things that work: exponential-trapezoidal discretization, complex rotary states, chunked prefill with cache consistency, mixed-precision LoRA fine-tuning, all MLX, no CUDA needed.
If you have any feedback or see any inconsistencies, hmu!
Note: haven't really tested it so far on public trained checkpoints (will do the next days/week) so this is mainly useful for local architecture experimentation and fine-tuning on toy data until the authors drop weights.
u/Traditional_Ad_6304 — 6 days ago