jada42/mlx-mamba3

Mamba-3 is a type of neural network architecture that processes text using a mathematical technique called state-space modeling. Unlike the transformer models that power most modern AI tools, Mamba-3 handles long sequences more efficiently. This repository brings a working version of Mamba-3 to Apple Silicon Macs, using a framework called MLX that runs directly on the Mac's built-in GPU chip. Most existing Mamba-3 code was written for Linux machines with Nvidia graphics cards (CUDA), which means Mac users were locked out of experimenting with the architecture locally. This project rebuilds the full model in pure Python and MLX, so anyone with an M1, M2, or M3 Mac can run it without any Linux tools or extra hardware. The implementation covers three main configurations. SISO mode (single-input, single-output) handles simple channel-by-channel processing. MIMO mode (multi-input, multi-output) uses matrix projections for more expressive mixing. Hybrid mode alternates between Mamba-3 blocks and standard attention layers. All three have been verified to produce mathematically identical results to the original PyTorch code, with a maximum error below 0.00001. Beyond basic inference, the repository includes LoRA fine-tuning support, which lets users adapt a pre-trained model to new text data using mixed-precision training on the Mac GPU. Weights can be saved and loaded in the standard safetensors format. A benchmarking script shows roughly 469 tokens per second during text generation on an M1 Pro machine. The codebase is organized into a core Python package covering model definition, cache management, weight loading, generation loop, and training utilities, plus example scripts for text generation, hybrid model use, and a small fine-tuning demo on a toy dataset. A test suite compares numerical outputs against the PyTorch reference at each correctness boundary. The project is licensed under MIT and is actively maintained with continuous integration on each commit.