Every sequence model you have served pays attention’s tax: quadratic cost in sequence length, a KV cache that grows with every token, and memory bandwidth that becomes the bottleneck long before the FLOPs do. Mamba, from Albert Gu and Tri Dao and described in the 2023 paper Mamba: Linear-Time Sequence Modeling with Selective State Spaces, replaced that tax with a recurrence whose state is a fixed-size tensor, selected by the input itself so the model can choose what to remember. The bet landed. The Apache-2.0 repository, now at version 2.3.2.post1, carries three generations of the architecture, its own handwritten CUDA kernels, a full suite of Triton kernels, and pretrained checkpoints trained on 300B tokens of the Pile, and its adoption page lists Hunyuan-TurboS at 560B parameters, AI21’s Jamba at 398B, NVIDIA’s Nemotron-H, TII’s Falcon-H1, IBM’s Bamba, and inference stacks like vLLM and TensorRT-LLM. As always in this series, what follows is an educational tour of published source code.

The repository is deliberately small at the top and deep at the bottom. One Python package, mamba_ssm, holds five modules: model blocks you instantiate, ops that compute their recurrence, a language model that stacks the blocks, utilities for generation and distribution, and the glue for Hugging Face. Underneath sit two compute layers, one handwritten in CUDA under csrc for the Mamba-1 recurrence, and one written in Triton under mamba_ssm/ops/triton for the Mamba-2 SSD algorithm and the Mamba-3 kernels. Three papers, three block classes, and two kernel backends all funnel through the same four-method contract: forward for training, step for decoding, allocate_inference_cache for serving, and a reset path for state. Reading the overview from left to right follows one input from package facade to machine code.

Overview of Mamba: the mamba_ssm package exports three block generations and a language model, the blocks call autograd-wrapped selective scan and Triton SSD ops, a handwritten CUDA core runs the recurrence, and serving utilities handle generation and tensor parallelism.

Reading the overview from left to right:

• The Mamba block is 294 lines: the class at line 31, forward at 119, and a single-token step at 208.
• Mamba2 at line 37 of 383 lines rewrites the block around the chunked SSD algorithm with d_state 64 or 128.
• selective_scan_interface.py wraps the CUDA recurrence in autograd: SelectiveScanFn at 34, mamba_inner_fn at 386.
• The CUDA core is a 497-line dispatcher over a 376-line forward kernel and a 561-line backward kernel.
• ssd_combined.py at 1054 lines drives the chunked scan, with mamba_chunk_scan_combined at line 628.
• ssd_chunk_scan.py is the biggest single kernel at 1887 lines, flanked by chunk state at 1122.
• The language model in 315 lines: MixerModel at 124, MambaLMHeadModel at 221, create_block at 30.
• generation.py is 389 lines of decoding machinery: InferenceParams at 18, DecodingCGCache at 271.

Why You Need This

The first reason is that the recurrence is genuinely selective, and the code shows it plainly. In the Mamba block’s constructor, the dt bias is initialized so that softplus of the bias lands between dt_min and dt_max, which pins the time-constant of each state channel into a targeted range inherited from S4D work; lines 90 through 98 of mamba_simple.py do exactly this, and line 100 flags that bias with _no_reinit so training frameworks that zero all biases do not destroy it. The A matrix comes from an S4D real initialization at lines 103 through 109, stored as A_log and kept in fp32. These are not cosmetic details; they are the difference between a recurrence that trains and one that silently drifts, and having them in readable Python next to the kernel that consumes them is worth more than any amount of paper reading.

The second reason is the two-backend design, which is a masterclass in performance engineering trade-offs. Mamba-1’s recurrence is fused into one handwritten CUDA kernel: the forward at 376 lines, a backward at 561 lines, and a reverse-scan helper at 452 lines, all dispatched from 497 lines of C++ that switch on dtype, dimension, and sequence format through six tiny 9-line instantiation files. Mamba-2 bet that the SSD algorithm’s chunked structure is better expressed in Triton, where kernels like the 1887-line chunk scan can be iterated on without a compiler toolchain. Mamba-3 then grew a third route with siso and mimo Triton kernels plus TileLang and CuTe variants for decode. Watching the same math move across backends as the hardware landscape changed is education you cannot get from a tutorial.

The third reason is that this is one of the few research repositories where inference is a first-class citizen, not an afterthought. Every block exposes a step method that updates conv and SSM states in place for single-token decoding, the 389-line generation module carries an InferenceParams object threaded through the whole stack, and the CUDA build itself is opt-in, since the core package installs pure Python and only pulls the selective_scan_cuda extension when MAMBA_KEEP_CUDA_BUILD is set. The precision troubleshooting section is unusually honest, warning that SSMs are sensitive to their recurrent dynamics and recommending fp32 parameters under AMP when other frameworks that store fp16 weights show instabilities. That is the kind of operational knowledge that only lives in real codebases.

The detailed map: three block generations sharing a Block wrapper and language model, the Python scan ops, the Triton kernel family, the handwritten CUDA core, serving and evaluation tooling, and the parallelism plus build layers.

Start with the block, because the block is the product. Mamba projects the input to x and z with one matrix multiply at line 135, runs a short causal convolution, then feeds the result to the selective scan with A derived as negative exponentials of A_log at line 143. The fast path at line 146 calls mamba_inner_fn, a fully fused CUDA route that requires the external causal-conv1d package and a training shape; otherwise the code takes the modular path with separate conv, scan, and gated normalization. The step method at line 208 is the same math on one token, updating the conv state and the SSM state in place, which is what makes constant-memory decoding possible. Mamba2 at line 37 keeps the interface and swaps the interior: state expansion moves from 16 to 64 or 128, the scan becomes the chunked SSD algorithm, and a gated RMSNorm from the 437-line layernorm_gated module replaces the simpler activation, while the class also inherits PyTorchModelHubMixin so the block itself is hub-pushable.

The scan ops are where autograd meets raw kernels. selective_scan_interface.py defines SelectiveScanFn at line 34 with forward at 37 and backward at 71, then wraps it as selective_scan_fn at 118, keeping a pure PyTorch selective_scan_ref at 127 as ground truth for tests. The fused MambaInnerFn at 196 chains convolution, scan, and gating in one backward-friendly call, exposed as mamba_inner_fn at 386. On the SSD side, ssd_combined.py offers MambaChunkScanCombinedFn at 596 behind mamba_chunk_scan_combined at 628, plus MambaSplitConv1dScanCombinedFn at 820 and the split conv1d scan entry at 983 that fuses convolution, scan, normalization, and projection for the simple Mamba-2 variant. Underneath, the 1887-line ssd_chunk_scan kernel does the state-times-input heavy lifting, the 1122-line chunk state module moves activations into and out of recurrence states, a 350-line state passing kernel carries states across chunks, and a 264-line bmm module computes the CB outer products, each with a reference implementation you can diff against.

The CUDA core rewards a slow read. selective_scan.cpp uses a 25-line static_switch header to dispatch on every binary axis, dtype, whether B and C are variable or shared, whether D exists, whether z gating is present, and whether the sequence is a single strand or variable-length strands, and the six 9-to-10-line .cu files are the only place dtype meets compilation. The forward kernel at 376 lines fuses the whole recurrence into one pass with careful thread-level state distribution, the backward kernel at 561 lines recomputes and produces every gradient including dx through the 452-line reverse scan helper, and all of it is built by a 407-line setup.py whose FORCE_BUILD and KEEP_CUDA_BUILD flags at lines 41 and 42 implement the opt-in CUDA install. Mamba3 at line 43 of 526 lines then shows where the architecture is heading: a MIMO mode with a mimo_rank parameter, an angle state carried across tokens, preprocessing at 281 and postprocessing at 304 around the recurrence, and kernels in three styles under ops/triton/mamba3, ops/tilelang/mamba3, and ops/cute/mamba3, with the paper’s inner module also available as the 103-line ssd_minimal reference.

The model and serving layers are thin on purpose. mixer_seq_simple.py builds each layer with create_block at line 30, initializes weights at 92, wraps the stack in MixerModel at 124, and adds the language head in MambaLMHeadModel at 221, a GenerationMixin subclass whose from_pretrained at 293 pulls config and weights through the 23-line utils/hf.py helpers. The Block wrapper at line 10 of block.py uses the add-then-norm order with a fused_add_norm path at line 57 that merges the residual add into the layer norm kernel, and the optional mixer slots accept the 294-line MHA module for hybrid models like mamba2attn and transformerpp or the 34-line GatedMLP. Decoding lives in generation.py, where InferenceParams at line 18 tracks sequence offsets and caches, get_logits at 173, sample_tokens at 197, and should_stop at 205 implement sampling, and a DecodingCGCache at 271 manages CUDA graphs for fast decode; the 296-line tensor_parallel module adds ColumnParallelLinear at 111 and RowParallelLinear at 151 for sharding, and the 43-line usage.md lists who runs all this in production.

From Install to a Running Block

The core install is pure Python, with the CUDA extension as an explicit opt-in on Linux, Python 3.10+, PyTorch 1.12+, and CUDA 11.6+:

pip install mamba-ssm --no-build-isolation
# with the fused conv package and the CUDA extension:
pip install "mamba-ssm[causal-conv1d]" --no-build-isolation
MAMBA_KEEP_CUDA_BUILD=TRUE pip install mamba-ssm --no-build-isolation

Then use the block exactly as the README shows:

import torch
from mamba_ssm import Mamba

batch, length, dim = 2, 64, 16
x = torch.randn(batch, length, dim).to("cuda")
model = Mamba(
    d_model=dim, # Model dimension d_model
    d_state=16,  # SSM state expansion factor
    d_conv=4,    # Local convolution width
    expand=2,    # Block expansion factor
).to("cuda")
y = model(x)
assert y.shape == x.shape

A block uses roughly 3 times expand times d_model squared parameters; swap in Mamba2 with d_state=64 for the SSD variant, or Mamba3 with headdim=64 and is_mimo=True for the newest recurrence. For a full language model, MambaLMHeadModel.from_pretrained loads any of the state-spaces checkpoints, from mamba-130m at 24 layers and 768 dimensions up to mamba-2.8b at 64 layers and 2560 dimensions, and the generation benchmark script times top-p sampling against them. Zero-shot evaluation goes through lm-eval at version 0.4.2, with tasks like lambada_openai, hellaswag, piqa, and arc_easy, either from the command line or the evals/lm_harness_eval.py wrapper.

Honest limits: the CUDA extension is Linux-only, so Windows development means WSL or the pure-Python fallback, and ROCm users on 6.0 need a patch from the rocm_patch directory that 6.1 and later do not. The pretrained checkpoints are base models trained on the Pile without instruction tuning, so they are research baselines, not assistants. The layer counts double a Transformer’s at equal size because two Mamba blocks replace one attention-plus-MLP layer, which matters when comparing parameter budgets. And the kernels themselves assume shapes and dtypes the dispatchers know, so exotic configurations fall back to reference implementations that are correct but slow. Read the block first, then the scan, then the kernels; that order turns a famous paper into code you can actually modify.

Enjoyed this post? Never miss out on future posts by following us