Mamba: what the selective state space paper actually says
Reading Gu and Dao (CMU / Princeton, December 2023) while investigating why KV cache growth was killing our long-context throughput.
The request came from the product team: support 32K-context conversations without quadrupling the serving cost. We were on LLaMA-13B with PagedAttention already deployed. The KV cache math was brutal — at 32K tokens, a single conversation occupied roughly 26 GB of KV cache on an 80 GB A100. Two concurrent long-context sessions and you were OOM. Short of buying eight times the hardware or moving to a 4-bit quantized model and hoping quality held, there wasn't a clean answer within the Transformer paradigm.
That's when I started taking state space models seriously.
The paper is "Mamba: Linear-Time Sequence Modeling with Selective State Spaces", by Albert Gu (Carnegie Mellon) and Tri Dao (Princeton, now at Together.AI), published on arXiv in December 2023 (arXiv:2312.00752). Dao is also the first author of FlashAttention, which is worth noting — the hardware-aware thinking that made FlashAttention work is directly applied here. The contribution is a new sequence model architecture that processes tokens in O(N) time and O(1) memory per decoding step, is competitive with Transformers on language modeling quality at the same parameter count, and does it using a mechanism that has nothing to do with attention.
The problem the paper is actually solving
To understand what Mamba is doing, you need to understand why its predecessor (S4) wasn't good enough.
State space models have been the dominant framework for continuous-time dynamical systems for decades. In the discrete sequence setting, an SSM describes a latent state h evolving over time:
h_t = A h_{t-1} + B x_t
y_t = C h_t
Where x_t is the input at step t, h_t is the hidden state (a vector of dimension N), and y_t is the output. A (N×N), B (N×1), and C (1×N) are learned parameter matrices. This is structurally identical to an RNN — a fixed recurrence relation applied sequentially.
The structured state space model S4 (Gu et al., 2022, the same first author) showed that if you constrain A to a specific structure (diagonal plus low-rank), you get two useful properties. First, the recurrence can be parallelized as a global convolution during training — the entire sequence can be processed in O(N log N) using FFT rather than step-by-step. Second, the structured A enables stable long-range dependency learning, which vanilla RNNs struggle with.
S4 matched or beat Transformers on structured sequence tasks (audio, time series) and came within striking distance on the Long Range Arena benchmark. The problem was language modeling. On text, S4 trailed Transformers substantially — not because of parameter count or training compute, but because of something more fundamental about what SSMs could represent.
The limitation was that S4 uses Linear Time Invariant (LTI) dynamics. The parameters A, B, and C are fixed once trained — they don't depend on the input. Every token goes through the same state transition, regardless of content. This works for continuous signals where the dynamics of the system don't depend on what you're measuring (audio waveforms, video frames). It doesn't work well for text, where whether you should remember or forget context depends entirely on what the tokens say.
Consider the sentence: "The lawyer the man who had studied at Harvard graduated from did pro bono work." Whether the model should maintain "lawyer" in its state while processing "who had studied at Harvard" depends on understanding that this clause modifies a person referenced earlier, not the lawyer. The right state transition is content-dependent. LTI dynamics can't express this — they apply the same transition regardless.
Mamba's core contribution is making the SSM parameters input-dependent.
The selective mechanism
Mamba calls its approach the Selective State Space Model, or selective SSM. The change is surgical: instead of fixed B, C, and a scalar Δ (the discretization step size that converts continuous-time parameters to discrete), these become functions of the input x.
B_t = Linear_B(x_t) # [N × 1], input-dependent
C_t = Linear_C(x_t) # [1 × N], input-dependent
Δ_t = softplus(Linear_Δ(x_t)) # scalar, input-dependent
A stays fixed (input-independent). The discretization is then applied using the zero-order hold method to get the effective discrete Ā and B̄:
Ā_t = exp(Δ_t · A)
B̄_t = (Δ_t · A)^{-1} (Ā_t - I) · Δ_t · B_t
The role of Δ is the key intuition. When Δ is large, Ā_t ≈ 0 and B̄_t ≈ B_t — the model essentially resets its state and focuses on the current input. When Δ is small, Ā_t ≈ I and B̄_t ≈ 0 — the model ignores the current input and preserves its state. This is a learned, input-conditioned interpolation between "remember context" and "attend to current token."
By making B and C input-dependent as well, the model can selectively control what information gets written into the state (via B) and what gets read out (via C).
The paper frames this as a selection mechanism analogous to what softmax attention does: attention selects which past tokens to retrieve based on content similarity; selective SSM selects what to remember based on content. The mechanism is structurally different but serves the same purpose.
Why this broke the efficient computation
S4's LTI property enabled the convolution trick. Because A, B, C were fixed, you could precompute a global convolutional kernel K = [CB, CAB, CA²B, ...] and apply it to the entire input sequence in parallel using FFT. Training was efficient.
Making parameters input-dependent destroys this. The kernel K_t changes at every step because B_t and C_t change. You can't precompute it. You're back to sequential recurrence, which is slow to train on GPUs — GPUs are designed for parallel, not sequential, computation.
This is the same problem FlashAttention solved for attention by avoiding materialization. Mamba solves it with a different technique: the parallel scan algorithm (also called prefix sum or scan reduction, from Blelloch 1990).
The parallel scan computes all prefix values of an associative operation in O(log N) depth and O(N) total work using a tree-reduction pattern, rather than O(N) depth for the sequential approach. The recurrence h_t = Ā_t h_{t-1} + B̄_t x_t has exactly the form required — it's an associative operation on (Ā, B̄x) pairs.
Critically, Mamba also applies kernel fusion: rather than writing intermediate states to HBM (slow) and reading them back, the entire scan runs inside SRAM with the intermediate hidden states never materialized to HBM. This is the FlashAttention insight applied to recurrence. The hidden state h_t is N-dimensional (N=16 in the paper's main experiments). For N=16, float16: 16 × 2 bytes = 32 bytes per token. Even for a 32K-sequence, the total hidden state stream is 32K × 32 bytes = 1 MB — trivially fits in SRAM.
The result: Mamba trains with the parallelism of a modern GPU, and runs inference with pure O(1)-memory recurrence.
The Mamba block architecture
The full Mamba block is simpler than a Transformer block. There's no separate attention sublayer and MLP sublayer with residual connections between them. The Mamba block is:
z, x = split(Linear(input)) # project + gate split
x = depthwise_conv1d(x) # short local context
x = silu(x)
x = SSM(x) # selective state space model
output = x * silu(z) # gated output
output = Linear(output) # project back
The 1D convolution adds a small window of local context (kernel size 4) before the SSM, which helps in practice. The gated structure (multiplying by silu(z)) is adapted from H3 (Fu et al. 2023). There's no positional encoding — the SSM has implicit positional awareness through its sequential processing.
A full Mamba model stacks these blocks with residual connections and normalization, similar to a Transformer's block stacking. The parameter count is comparable to a Transformer with the same embedding dimension and depth.
What the performance numbers actually show
The main language modeling results compare Mamba against GPT-NeoX (a clean Transformer baseline) on The Pile at various scales, from 130M to 2.8B parameters.
Quality: Mamba 1.4B reaches lower perplexity on The Pile than GPT-NeoX 1.3B — a Transformer with more parameters. The gap is notable: Mamba consistently beats parameter-matched Transformers across scales. The scaling curves show similar slope, suggesting the quality advantage holds into larger regimes.
Throughput: This is where the architecture difference is most dramatic. The paper reports inference throughput (tokens/second) vs. sequence length for Mamba vs. Transformers:
- At 2K tokens: Mamba is ~5× faster throughput than a Transformer of equal size
- At 16K tokens: the gap grows, because Transformer attention is O(N²) while Mamba is O(N)
The throughput advantage comes from two sources: no KV cache to manage (inference is pure recurrence), and no attention computation scaling quadratically with context length.
Memory: During inference, Mamba's state is fixed-size regardless of sequence length. The hidden state per layer is N × D (state dimension × model dimension). For a Mamba-1.4B with N=16, D=2048: 16 × 2048 × 2 bytes ≈ 64KB per layer × 48 layers = ~3MB total. This is independent of sequence length. Compare to PagedAttention's ~820KB per token per layer for LLaMA-13B — a 32K-context conversation occupies 26 GB; Mamba occupies 3MB.
For the long-context serving problem that motivated this reading, Mamba looked like a fundamental win.
The production implications
Inference architecture changes completely. With Transformers, inference has two phases: prefill (process the prompt in parallel) and decode (generate tokens one by one, maintaining a growing KV cache). With Mamba:
- Prefill is sequential per-token recurrence — slower than Transformer prefill for the same sequence length
- Decode is pure recurrence — no KV cache, O(1) memory, and simpler to implement
- There is no equivalent to prefix caching (a major vLLM/SGLang optimization for repeated system prompts)
At short sequences, Transformer prefill parallelism wins. The crossover point (where Mamba's decode advantage overcomes its prefill disadvantage) depends on the output/input length ratio. For generation-heavy workloads with long outputs relative to inputs, Mamba is faster. For prompt-heavy workloads (long system prompts, short completions), Transformer prefill + fast KV cache lookup wins.
The fixed state is both the feature and the bug. The bounded memory is great for serving. But the state has to encode everything the model needs to condition on — the entire conversational context compresses into a few megabytes. For most token predictions in a well-structured conversation, this works. For tasks requiring exact recall of something said earlier, it doesn't.
Batching efficiency is different. With Transformers, decode is memory-bandwidth bound regardless of batch size; batching helps by amortizing weight loads. With Mamba, decode processes each token sequentially through the state — weight loads per token are lower (smaller parameter matrices than full Transformer), and batch processing is also straightforward. The GPU utilization profile differs but is not obviously worse.
Fine-tuning and adaptation: The LoRA / QLoRA ecosystem targets linear layers, and Mamba's architecture is mostly linear projections — adapters should work similarly. In practice, the community has adapted LoRA to Mamba without fundamental issues.
When NOT to use Mamba
When you need in-context learning on rare or arbitrary formats. Transformers can attend directly to any prior token in context. If a user demonstrates a novel format in the prompt ("I want output structured like: [example]"), a Transformer can copy that pattern precisely because it can retrieve the example via attention. Mamba compresses that example into state — it can learn from it, but verbatim copying from distant context is unreliable. Tasks requiring explicit retrieval from a long prompt are where Mamba fails most visibly.
Needle-in-a-haystack retrieval. The paper evaluates on selective copying tasks and shows Mamba can fail at retrieving specific tokens from long context when they were introduced early. This shows up in "needle in a haystack" benchmarks where a key fact is buried in a 50K-token document and the query requires recalling it exactly. Selective SSMs have a theoretical information bottleneck — the state dimension N bounds what can be retained.
When you need prefix caching. Many production deployments cache the KV cache for system prompts that repeat across requests. With Transformer-based serving, a 2K-token system prompt processed once can be cached and reused for thousands of subsequent requests. Mamba has no analog — the recurrent state is sequence-specific, and there's no equivalent to KV cache prefix sharing. If your workload has repeated prefixes (chatbots, API wrappers with fixed system prompts), Mamba loses a significant serving efficiency advantage that Transformers have.
When you're integrating with an existing fine-tuned model ecosystem. The Transformer ecosystem — thousands of fine-tuned models on HuggingFace, LoRA adapters, GGUF quantizations — doesn't transfer to Mamba. You're starting from scratch or from the small set of Mamba base models. For most production teams, this is the real blocker, independent of the architecture merits.
For short sequences. The quadratic scaling of attention becomes a practical problem above roughly 4–8K tokens. Below that, Transformer inference is fast, KV cache fits in memory, prefix caching works well, and the ecosystem support is orders of magnitude better. Mamba's throughput advantage is most pronounced at 16K+ token sequences.
What happened after publication
Mamba's architecture spawned immediate follow-on work that is probably more production-relevant than pure Mamba itself.
Jamba (AI21 Labs, March 2024) interleaves Mamba blocks and Transformer attention blocks in a single model. The resulting hybrid gets Mamba's long-context memory efficiency combined with Transformer attention's strong in-context retrieval at critical layers. Jamba-1.5 (7B active / 52B total with MoE) achieves competitive quality with significantly better long-context throughput than a pure Transformer of equivalent active parameters.
Mamba 2 (Gu and Dao, May 2024) reformulated the selective SSM using Structured State Space Duality (SSD), showing a mathematical equivalence between a class of SSMs and a restricted form of attention (with a specific structured mask). This unification lets Mamba 2 use matrix multiplications (rather than the custom parallel scan) for the core computation, achieving better GPU utilization. Mamba 2 is ~2× faster than Mamba 1 in practice and has more principled theoretical footing.
The dominant production trajectory appears to be hybrid attention-SSM models rather than pure Mamba. The intuition from Jamba and follow-on work: you don't need attention everywhere, but you do need it somewhere — periodic attention layers give the model the exact-retrieval capability that pure Mamba lacks, while SSM layers handle the bulk of sequence processing cheaply. The optimal ratio of attention-to-SSM layers is an active research question, but even 1-in-6 or 1-in-8 attention layers appears to recover most of Transformer's in-context learning quality.
The deeper implication: Mamba didn't replace the Transformer. It demonstrated that state-based alternatives can reach competitive quality, which opened the design space for hybrid architectures. If you're building a new foundation model for long-context applications, the Mamba paper is required reading — not because you'll implement pure Mamba, but because it gives you the conceptual tools to think about what attention is doing and what you can replace it with.
For the original problem — 32K-context serving at reasonable cost — the answer we landed on wasn't pure Mamba deployment, but the paper changed how we thought about the architecture space. The KV cache isn't a fundamental requirement. It's a consequence of one particular approach to building sequence models, and there are alternatives.