Mamba: what the selective state space paper actually says
Reading Gu & Dao (2023) after watching a long-context inference job allocate 40 GB of KV cache for a single batch.
The KV cache problem is straightforward to state and annoying to solve. Each token you generate extends the key and value tensors that attention needs to attend over. At 32K tokens, a 70B model's KV cache across 40 layers, with 8 KV heads at 128 dimensions, in FP16, costs roughly 20 GB per request. Scale to 10 concurrent long-context requests and you've filled an H100. The model weights aren't the bottleneck — you've been running 70B for months. It's the workspace that grows with every token you generate.
PagedAttention (vLLM) helps by managing KV cache memory like virtual memory pages instead of pre-allocating contiguous blocks. FlashAttention reduces attention's memory footprint during the forward pass. Neither addresses the root issue: in an attention-based architecture, every token you've generated has to be stored in memory for every future token you generate. The KV cache is O(N) per request, unavoidably, by design.
Mamba — "Mamba: Linear-Time Sequence Modeling with Selective State Spaces", Gu and Dao, 2023 — proposes a fundamentally different approach. Instead of attending over the entire history, Mamba compresses it into a fixed-size hidden state. Generating the next token requires only that state plus the new input — not the full history. The memory cost per request is constant regardless of sequence length. The KV cache doesn't exist.
This sounds too good. Understanding where it actually works requires understanding what the compression loses, and what the paper's "selective" innovation actually changes.
What state space models are, and why they weren't good enough before
State space models (SSMs) are a classical control-theory formalism for mapping a one-dimensional input sequence to an output sequence through a latent state. The continuous-time formulation:
ḣ(t) = Ah(t) + Bx(t)
y(t) = Ch(t)
where x(t) is input, h(t) is the hidden state, y(t) is output, and A, B, C are learned matrices. Discretized with a step size Δ (the Zero-Order Hold method, in the paper), this becomes:
h_t = Āh_{t-1} + B̄x_t
y_t = Ch_t
where Ā = exp(ΔA), B̄ = (exp(ΔA) - I)A⁻¹BΔ
This is a recurrent form: each step's state is a function of the previous state plus the current input. At inference time, this is ideal — you maintain a hidden state of fixed dimension and update it one token at a time. Constant memory, O(1) per token.
The problem with naive SSMs for language modeling is that A, B, and C are time-invariant: they don't depend on the input. The same transition matrix A applies to every token regardless of what that token is. This means the model treats every position in the sequence identically — it can't decide to "pay attention" to some tokens and "forget" others based on content. It's a linear, input-agnostic compression of history.
The S4 paper (Gu et al., 2021) made SSMs practical at scale by using a diagonal-plus-low-rank structure for A (the HiPPO initialization) and exploiting the fact that time-invariant SSMs can be computed as a convolution during training — enabling parallelization across the sequence dimension. S4 achieved strong results on long-range benchmarks like Path-X. But on language modeling, it lagged behind transformers. The time-invariant compression isn't expressive enough for the highly content-dependent structure of language.
The selective scan: Mamba's actual contribution
Mamba's key insight is that the time-invariance is the fundamental limitation, and you can break it without losing the inference-time efficiency of the recurrent form.
In Mamba, B, C, and the step size Δ are functions of the input x_t. They're not learned constants — they're computed from the current token at each timestep. The discrete transition becomes:
h_t = Ā(x_t)h_{t-1} + B̄(x_t)x_t
y_t = C(x_t)h_t
This is a "selection mechanism": the model learns how to update its state in a content-dependent way. It can learn to let important tokens update the state significantly (large Δ, meaning the state tracks the input closely) and to let unimportant tokens pass through without changing the state (small Δ, meaning the state mostly carries forward). The paper shows this is analogous to gating in LSTMs, but derived from first principles in the SSM framework.
The price: making B, C, and Δ input-dependent breaks the convolution equivalence. Time-varying SSMs can't be computed as a convolution — you're back to needing to run the recurrence explicitly, which is sequential.
The paper solves this with a hardware-aware parallel scan algorithm. Prefix scans (also called parallel prefix sums) are a classical parallel algorithm that computes all prefix results of an associative operation in O(log N) steps using a tree-structured reduction. The recurrence h_t = Āh_{t-1} + B̄x_t can be reformulated as a prefix scan over tuples (Ā, B̄x) — the operation (Ā₁, b₁) ∘ (Ā₂, b₂) = (Ā₁Ā₂, Ā₁b₂ + b₁) is associative, which is all you need.
The practical problem with a naive parallel scan: it requires materializing all intermediate hidden states simultaneously — O(N × d_state × d_model) memory, where d_state is typically 16 and d_model might be 4096. This is a significant HBM footprint for long sequences.
Mamba handles this with the same IO-aware strategy as FlashAttention: kernel fusion and selective recomputation. The expanded states (the high-dimensional d_state × d_model tensors) are kept in SRAM during the scan and never written to HBM. During the backward pass, the intermediate states are recomputed from the inputs rather than stored. The result is a training algorithm that's O(N) in compute but avoids the HBM reads and writes that would make it slow in practice.
What the performance numbers actually show
The paper compares Mamba to Transformer++ (the Pythia / GPT-NeoX recipe: rotary embeddings, SwiGLU FFN, grouped-query attention, rmsNorm) across parameter counts from 130M to 2.8B on language modeling benchmarks.
At 370M and 1.4B parameters, Mamba roughly matches Transformer++ on perplexity. At 2.8B, Mamba outperforms a same-size transformer across the Pile subsets. The paper frames this as Mamba matching "a Transformer with twice the number of parameters" — meaning Mamba-1.4B ≈ Transformer-2.8B in perplexity. Whether that holds up at scale beyond 7B has been contested in follow-up work, but the result at the paper's scale is clean.
The throughput gap is where Mamba's advantage is most pronounced. The paper reports:
- At 2K sequence length: Mamba is 4.8× faster than a Transformer of comparable size in generation (autoregressive token generation)
- At 16K sequence length: Mamba is ~5× faster
This is the KV cache effect in practice. The transformer's generation throughput degrades with sequence length because each new token requires attending over an increasingly long KV cache. Mamba's generation throughput is flat — it performs the same fixed-size state update regardless of how many tokens came before.
Prefill (processing the prompt) is a different story. For a prompt of length N, both models are O(N) in compute, but the constant factors differ and the transformer is better optimized (FlashAttention, fused kernels, etc.). At short lengths (<2K), the transformer's prefill advantage combined with the overhead of Mamba's scan kernel means the transformer is often faster overall.
Production tradeoffs that don't make it into the benchmark post
The state is fixed-size regardless of context, but you trained at a fixed length. Mamba compresses sequence history into a fixed-size state, which means the amount of information retained about any particular token degrades as more tokens are processed — the state is being overwritten continuously. If you trained on sequences up to 8K tokens, the model learned to compress 8K tokens of context into the state. That training objective doesn't guarantee the state is useful at 32K or 128K tokens. Empirically, Mamba-based models trained at shorter context lengths generalize worse to longer contexts than transformers with RoPE. You can train on longer sequences, but there's no position embedding trick that buys you free length generalization.
Prompt processing isn't free. The KV cache advantage only applies during token generation, not during prefill. If you're building a system where the prompt is 50K tokens (a document ingested for QA) and the response is 512 tokens, the generation speedup is negligible compared to the prefill cost, and your serving infrastructure still has to handle the O(N) prefill step. The paper's 5× throughput gain is for generation-dominated workloads.
Serving infrastructure resets. vLLM, TGI, TensorRT-LLM, and the broader serving ecosystem are built around transformer assumptions: KV cache management, continuous batching at the token level, PagedAttention for memory efficiency. Mamba's state is simpler (one fixed-size tensor per layer per request) but the existing frameworks don't manage it natively. As of 2024-2025, Mamba serving requires either running vanilla Hugging Face with no serving optimizations, or substantial custom work to integrate state management into a production serving framework. The 5× throughput advantage can easily disappear if you're running Mamba on HF Transformers and Transformers on vLLM.
Hybrid architectures shift the equation. Jamba (AI21 Labs, 2024) and similar hybrid architectures alternate Mamba layers with sparse attention layers. A typical ratio is 7 Mamba layers per 1 attention layer. This gives you most of Mamba's generation-time efficiency (the KV cache only grows at attention layers) while preserving attention for the retrieval-heavy operations that SSMs handle poorly. The KV cache is ~7× smaller than a pure transformer, not eliminated. Hybrid architectures are arguably more production-pragmatic than pure Mamba because they sidestep the associative recall failure mode while still capturing most of the memory savings.
You can't prompt-cache across requests. Transformer serving systems support KV cache prefix sharing: if 100 requests start with the same system prompt, you compute that prompt's KV cache once and reuse it across all 100 requests. With Mamba, the state after processing a prompt is an opaque compressed summary — it's not cacheable in the same way because combining a cached state with a new request requires actually running the model on the prefix again. Systems built on prefix caching (which includes most production deployments with long system prompts) lose that optimization entirely.
Failure modes in practice
Associative recall. This is the fundamental limitation the paper acknowledges. Tasks that require "look up what value was associated with this key, somewhere in the history" — like in-context learning with many examples, or document QA where you need to retrieve a specific fact stated 10K tokens ago — perform noticeably worse in pure Mamba models. The fixed-size state compresses history by necessity, and rarely-accessed information gets overwritten. Attention handles this naturally because the QKV mechanism is explicit retrieval over the full history. The Zoology paper (Arora et al., 2023) and subsequent synthetic benchmarks quantify this gap precisely: on Multi-Query Associative Recall (MQAR) and similar tasks, SSMs significantly lag transformers at the same parameter count.
In-context learning at scale. Transformers' ability to do few-shot learning from examples in the prompt relies on the ability to precisely retrieve patterns from those examples during generation. Mamba's state compression makes this less reliable at high shot counts or with examples that are structurally similar. For tasks where the examples differ substantially, this matters less. For tasks that require precise mimicry of a formatting or reasoning pattern across many examples, the recall degradation is measurable.
The "forgetting curve" on long sequences. In practice, Mamba models tend to lose information about the beginning of long sequences as more tokens are processed. This isn't a bug — it's a consequence of compressing unbounded history into fixed state. But it manifests in ways that can be surprising: a model that handles 4K-token contexts well may give inconsistent answers about facts stated in the first 500 tokens of an 8K-token context. The distribution of what gets retained is learned during training and not easily introspectable.
When not to use Mamba
Retrieval-heavy workloads. If your application requires reliably surfacing specific facts from long documents — "what was the revenue figure mentioned on page 7?", "what was the decision made in the third paragraph?" — pure Mamba architectures are the wrong choice. Either use a RAG system, a hybrid architecture, or a transformer. The compression-based state can't guarantee precise retrieval.
Short contexts where your stack is already transformer-optimized. If you're running inference at 1K-4K context lengths on well-optimized transformer infrastructure (vLLM, continuous batching, FlashAttention-2), you're probably leaving the transformer's efficiency gains on the table unnecessarily. Mamba's advantage materializes at long context and generation-dominated workloads. At short contexts, the transformer often wins on both throughput and quality due to better tooling.
When you need few-shot prompting as a primary feature. Few-shot examples in the prompt rely on attention-based retrieval to work reliably. If your application is built around "prepend examples, then generate," the in-context learning degradation in SSMs will affect output quality in ways that are hard to characterize without careful per-task evaluation.
When you need a reliable serving stack in the next six months. As of mid-2025, the Mamba serving ecosystem lags the transformer ecosystem by 2-3 years in production maturity. FlashMamba and similar optimized kernels exist but aren't as battle-tested. The memory management tooling for Mamba's recurrent states in a multi-request serving context is nascent. If you're deploying to production under real traffic, the pragmatic choice is transformers with mature tooling unless you have specific evidence that context length is your bottleneck and you have the engineering bandwidth to build the serving layer yourself.
When your prompt is most of your context. The KV cache advantage only shows up during generation. If you have a 20K-token prompt and generate 200 tokens, Mamba's per-token generation cost is lower, but the prompt processing is still O(N) and that's where most of your inference latency lives. Profile your token budget split before assuming Mamba will solve a memory problem.
What the paper actually gives you
Mamba makes two contributions worth separating.
The first is the selective mechanism: making SSM parameters input-dependent via B, C, and Δ as functions of x. This is the theoretical advance — it gives SSMs the content-dependent filtering that time-invariant models lack, and the paper demonstrates this closes most of the language modeling gap with transformers. The mechanism is elegant and the connection to gating mechanisms in LSTMs is explicitly drawn.
The second is the hardware-aware scan: the IO-aware parallel scan that makes selective SSMs trainable efficiently despite the broken convolution equivalence. This is in the same lineage as FlashAttention — reframe the bottleneck from compute to IO, then design the kernel to minimize HBM reads and writes using SRAM reuse and recomputation. It's the engineering contribution that makes the theoretical advance practical.
The production bet Mamba asks you to make is: the information you need from context can be represented in a fixed-size compressed state, and you can train the model to do that compression well for your task. For many language tasks, this bet pays off — the quality is competitive. For tasks that require precise recall of specific past content, the bet is wrong by design.
The long-context inference economics are real. If you're running 70B+ models with 100K+ context lengths and paying for KV cache memory at $0.03/GB-hour, the elimination of KV cache growth is a meaningful cost reduction. But that assumes you've solved the serving infrastructure problem, that your task is generation-dominated, and that your application doesn't require the precise retrieval that SSMs can't provide.
The 40 GB KV cache job that started this post? We moved it to a Jamba-style hybrid architecture — Mamba layers for most processing, attention every N layers for retrieval-critical positions. The KV cache dropped to about 6 GB. The quality on our document QA task held because the attention layers preserved retrieval for the positions that needed it. We also spent three months building the serving infrastructure to make it work under production traffic.
That ratio — months of infrastructure work for a 7× memory reduction — is probably the honest framing for where Mamba sits in 2025. The numbers are real. The engineering lift is also real.
Mamba: Linear-Time Sequence Modeling with Selective State Spaces — Gu, Dao. arXiv 2312.00752, 2023. Zoology: Measuring and Improving Recall in Efficient Language Models — Arora, Eyuboglu, Timalsina, Johnson, Poli, Zou, Rudra, Ré. arXiv 2312.04927, 2023. Jamba: A Hybrid Transformer-Mamba Language Model — AI21 Labs. arXiv 2403.19887, 2024.