Fast Attention and Long Context
Mamba and state space models
Mamba carries a fixed-size state through the sequence like a control system, and its key idea is letting the input decide what gets written into that state and what gets forgotten.
- 13 min read
- 3 reading levels
- Updated
Read these first
On this page 10
One lesson, three depths. Pick the one that fits you today — you can switch any time.
Beginner — No maths. Plain English.
The short answer
A state space model carries one small running memory, and Mamba lets each word decide what to keep.
The bucket with a tap
Picture a bucket under a tap, with a small hole in the bottom.
Water pours in, water leaks out. The level in the bucket is a summary of what has poured in recently. Old water has drained away; recent water is still there.
That level is a state — a single running summary that stands in for everything that came before. A state space model is a bucket, and each word pours a little water in.
The idea that made it work
Older versions of this had a fixed hole. Everything drained at the same rate, whether it mattered or not.
Mamba's change is that each word gets to adjust the tap and the hole. A word that matters is poured in hard, and the hole is closed so the level holds. A word that does not matter is barely poured, and nothing is disturbed.
That is what "selective" means in this design, and it is the whole idea. The model chooses what to remember while it reads, instead of following one fixed rule.
Why anyone wanted this
A transformer keeps every word it has read and searches through all of them for each new word. Costs grow with the square of the length, and the stored notes grow forever.
A bucket has one level. It costs the same at word ten as at word ten million. Memory does not grow at all.
transformer state space model
----------- -----------------
keep every word keep one running state
search all of them update it, read it
cost grows with length cost is flat
exact recall fuzzy recallWhere it comes from
This is not new mathematics. It is the standard way engineers describe a system that changes over time. A moving vehicle, a heating room, a circuit.
The surprise was that the same equations could model language. It took learned settings, and a clever way to run them fast on a graphics card.
The honest limit
A bucket cannot hold a list. Ask a pure state space model to copy a long random string it saw earlier. It does noticeably worse than a transformer.
This is not a bug to be fixed. A fixed-size summary cannot store an unbounded number of exact items. That is arithmetic, not engineering.
What is actually being built now
Almost every shipped model in this family is a mixture. Most layers are the cheap state space kind, and a few are full attention.
The cheap layers handle the flow of the text. The few attention layers handle the moments when something must be looked up exactly. It works better than either alone.
Where you have already seen this
- Long-audio and long-genome models, where sequences run into millions of steps.
- Hybrid open models from NVIDIA, AI21 and others that advertise faster long-context throughput.
- On-device models with a flat memory budget however long the session runs.
Remember this
- A state space model carries one fixed-size state, like the level in a bucket.
- Mamba's contribution is letting the input control what is written and what is forgotten.
- Flat cost and flat memory, at the price of exact recall. Real models mix in a few attention layers.
What to learn next
- Ring attention and sequence parallelism — the other way to reach a million tokens, using many machines.
- Recurrent neural networks — the older family these designs return to.
- LSTM — gates that decide what to forget, invented for the same reason.
Developer — Code and libraries.
Setup
pip install numpyWritten against NumPy 1.26.4, Python 3.10. CPU, instant. Nothing is trained here; the parameters are set by hand so the mechanism is visible.
Selectivity, in twenty lines
import numpy as np
# A one-dimensional state space model, discretised with zero-order hold.
# continuous: h'(t) = A h(t) + B x(t), y(t) = C h(t)
# discrete: h_t = Abar h_{t-1} + Bbar x_t, y_t = C h_t
A, C = -1.0, 1.0
def step(h, x, delta):
Abar = np.exp(delta * A) # decay over a step of size delta
Bbar = (Abar - 1) / A # zero-order-hold discretisation of B=1
return Abar * h + Bbar * x
# The sequence: one value worth keeping at t=1, then 14 distractors.
x = np.array([0.0, 5.0, 0.3, -0.2, 0.4, 0.1, -0.3, 0.2, 0.5, -0.4, 0.1, 0.3, -0.1, 0.2, 0.0, 0.1])
important = np.zeros(len(x)); important[1] = 1.0
print(f"{'t':>3}{'x_t':>7}{'fixed delta=0.5':>18}{'selective delta':>18}")
h_fixed = h_sel = 0.0
for t, xt in enumerate(x):
h_fixed = step(h_fixed, xt, 0.5) # same delta every step (this is S4)
delta_t = 1.0 if important[t] else 0.001 # delta chosen FROM the input (this is Mamba)
h_sel = step(h_sel, xt, delta_t)
print(f"{t:>3}{xt:>7.1f}{C*h_fixed:>18.4f}{C*h_sel:>18.4f}")
print()
print("after 14 distractor steps the fixed model has decayed to", round(C*h_fixed, 4))
print("the selective model still holds", round(C*h_sel, 4))
print()
print("state size is fixed, whatever the sequence length")
for L in (100, 10_000, 1_000_000):
print(f" length {L:>9,}: SSM state = 1 number | attention KV cache = {2*L:>9,} numbers")
print()
print("Mamba-130m, from state-spaces/mamba-130m-hf config.json")
for k, v in [("d_model", 768), ("state_size (N)", 16), ("expand", 2),
("d_inner = expand * d_model", 1536), ("conv_kernel", 4),
("n_layer", 24), ("time_step_rank", 48)]:
print(f" {k:<28}{v}")
print(f" state per layer = d_inner * N = {1536*16:,} numbers, constant in sequence length")t x_t fixed delta=0.5 selective delta 0 0.0 0.0000 0.0000 1 5.0 1.9673 3.1606 2 0.3 1.3113 3.1577 3 -0.2 0.7166 3.1544 4 0.4 0.5921 3.1516 5 0.1 0.3984 3.1486 6 -0.3 0.1236 3.1451 7 0.2 0.1537 3.1422 8 0.5 0.2899 3.1396 9 -0.4 0.0185 3.1360 10 0.1 0.0506 3.1330 11 0.3 0.1487 3.1301 12 -0.1 0.0508 3.1269 13 0.2 0.1095 3.1240 14 0.0 0.0664 3.1209 15 0.1 0.0796 3.1179 after 14 distractor steps the fixed model has decayed to 0.0796 the selective model still holds 3.1179 state size is fixed, whatever the sequence length length 100: SSM state = 1 number | attention KV cache = 200 numbers length 10,000: SSM state = 1 number | attention KV cache = 20,000 numbers length 1,000,000: SSM state = 1 number | attention KV cache = 2,000,000 numbers Mamba-130m, from state-spaces/mamba-130m-hf config.json d_model 768 state_size (N) 16 expand 2 d_inner = expand * d_model 1536 conv_kernel 4 n_layer 24 time_step_rank 48 state per layer = d_inner * N = 24,576 numbers, constant in sequence length
Reading the output
Column three decays to 0.0796. Column four holds 3.1179. Same equations, same inputs, same starting state. The only difference is that delta is a constant in one and a function of the input in the other.
That difference is the entire contribution of Mamba over the state space models that preceded it. Everything else in the paper is machinery to make it run fast.
Why delta = 0.001 preserves the state. Abar = exp(-0.001) is 0.999, so almost nothing decays. Bbar = (Abar - 1)/A is about 0.001, so almost nothing is written. A tiny step size means "this token barely happened".
Why delta = 1.0 writes hard. Abar = exp(-1) = 0.368 and Bbar = 0.632. The old state is largely flushed and the new input dominates.
So one scalar per token, per channel, controls both remembering and writing. That coupling is deliberate. It is the same idea as a forget gate in an LSTM, derived from a different starting point.
The state size table is the selling point. One number at length 100 and one number at length 1,000,000.
The real model's state is 24,576 numbers per layer. Constant, but not tiny. Mamba widens the channel dimension precisely because a small state is the binding constraint on what it can recall.
What is missing from this toy
Three things a real Mamba block has:
Per-channel parameters. A, B, C and delta are vectors, and A is initialised as a structured (HiPPO-derived) matrix rather than a scalar, so different channels decay at different rates. That is what lets one layer hold both fast and slow information.
A short causal convolution before the scan. conv_kernel: 4 in the config. It mixes each channel with its three predecessors, which the scan alone cannot do.
A parallel scan. The for t loop above is unusable for training. Real implementations use an associative scan, quadratic within chunks and recurrent across them, with the state kept in on-chip memory. That is the "hardware-aware" half of the paper, and it is the same trick as FlashAttention: fuse, tile, keep it off the slow memory.
Common mistakes
Writing the recurrence as a Python loop and benchmarking it. You will conclude SSMs are slow. Use mamba-ssm or the fused kernels in transformers.
Assuming constant memory means unlimited memory. A 24,576-number state at layer 12 is what the model has, whether the context is 1k or 1M tokens.
Installing mamba-ssm and expecting it to work everywhere. It needs a CUDA build with matching toolkit versions and does not support all GPUs. transformers ships a slow pure-PyTorch fallback that is correct and much slower; check which one you got before drawing conclusions about speed.
Testing only on perplexity. SSMs match transformers on aggregate next-token loss and lose on exact-copy and long-range retrieval tasks. Those are the tests that separate the architectures.
Calling every hybrid model "Mamba". Jamba, Nemotron-H and Qwen3-Next mix in different proportions with different recurrent blocks. The layer schedule matters more than the block name.
Try it yourself
Make delta_t a smooth function of the input instead of a switch, for example delta_t = 1 / (1 + np.exp(-abs(xt) + 1)). That is closer to what a trained model learns. Then set two channels with different A values and watch one hold information the other has already lost.
What to learn next
- Ring attention and sequence parallelism — the other way to reach a million tokens, using many machines.
- Recurrent neural networks — the older family these designs return to.
- LSTM — gates that decide what to forget, invented for the same reason.
Researcher — Mathematics and papers.
The model
A continuous linear time-invariant system:
$$ h'(t) = A h(t) + B x(t), \qquad y(t) = C h(t) $$
$h(t) \in \mathbb{R}^{N}$ is the hidden state, $x(t)$ the scalar input for one channel, $A \in \mathbb{R}^{N \times N}$, $B \in \mathbb{R}^{N \times 1}$, $C \in \mathbb{R}^{1 \times N}$. Discretise with step size $\Delta$ under zero-order hold:
$$ \bar{A} = \exp(\Delta A), \qquad \bar{B} = (\Delta A)^{-1}(\exp(\Delta A) - I)\,\Delta B $$
$$ h_t = \bar{A} h_{t-1} + \bar{B} x_t, \qquad y_t = C h_t $$
Because the system is linear and time-invariant, the whole sequence is also a convolution with kernel $\bar{K} = (C\bar{B}, C\bar{A}\bar{B}, C\bar{A}^2\bar{B}, \dots)$, computable by FFT in $O(N \log N)$. Training runs as a convolution, inference as a recurrence. That duality is the basis of S4 (Gu, Goel and Ré, 2021, Efficiently Modeling Long Sequences with Structured State Spaces, arXiv:2111.00396), which makes $A$ tractable via a structured HiPPO parameterisation.
Selection breaks the duality on purpose
Gu and Dao (2023), Mamba: Linear-Time Sequence Modeling with Selective State Spaces (arXiv:2312.00752) make $\Delta$, $B$ and $C$ functions of the input:
$$ \Delta_t = \operatorname{softplus}(W_{\Delta} x_t), \qquad B_t = W_B x_t, \qquad C_t = W_C x_t $$
The system is now time-varying, so no fixed convolution kernel exists and the FFT trick is gone. The paper's argument is that this loss is worth it: an LTI system processes every token identically, and content-based reasoning requires the opposite. Their motivating tasks — selective copying and induction heads — are unsolvable by any LTI system and solved by the selective one.
The replacement for the FFT is a hardware-aware parallel scan: the recurrence is an associative operation, so a Blelloch scan parallelises it, with $h$ kept in SRAM and $\bar{A}, \bar{B}$ recomputed in the backward pass rather than stored. Structurally identical to FlashAttention's argument.
The duality with linear attention
Dao and Gu (2024), Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality (arXiv:2405.21060, ICML 2024) show that selective SSMs with scalar-times-identity $A$ and masked kernel attention are two views of the same class of semiseparable-matrix transforms.
Concretely, Mamba-2's recurrence
$$ h_t = a_t h_{t-1} + B_t x_t $$
is gated linear attention with a scalar decay. The practical payoff was a rewrite that uses matmul-heavy chunked algorithms, giving a reported 2–8x speed-up over Mamba-1 with larger state sizes.
This is the single most useful fact in the area: the SSM literature and the linear attention literature are studying the same object from two directions. Papers in one field routinely rediscover results from the other.
Where the limits are
Jelassi et al. (2024), Repeat After Me: Transformers are Better than State Space Models at Copying (arXiv:2402.01032) prove and demonstrate the expected separation: copying a string of length $n$ requires state $\Omega(n)$, so a fixed-state model fails beyond its capacity while a transformer with a growing cache does not. The result generalises to any task requiring unbounded exact recall.
This is not a training problem. It is an information-theoretic bound on fixed-size state, and it is why nobody serious proposes replacing attention outright.
The hybrid consensus
| Model | Composition |
|---|---|
| Jamba (AI21, 2024) | Mamba layers interleaved with attention layers and MoE |
| Nemotron-H (NVIDIA, 2025) | Hybrid Mamba-Transformer at 8B, 47B and 56B; reported up to 3x throughput over similarly sized transformers |
| Qwen3-Next | 3 Gated DeltaNet layers per 1 full-attention layer, 48 layers |
The recurring design is a small minority of attention layers carrying the exact-recall load. Ratios between 1-in-4 and 1-in-8 are typical, and the ablations behind those numbers are among the more useful public results in architecture research.
Current state
Lahoti, Li, Chen, Wang, Bick, Kolter, Dao and Gu (2026), Mamba-3: Improved Sequence Modeling using State Space Principles (arXiv:2603.15569) add a more expressive discretisation-derived recurrence, a complex-valued state update for richer state tracking, and a multi-input multi-output formulation. They report matching Mamba-2 perplexity at half the state size, and a 1.8 point accuracy gain at 1.5B scale.
Halving the state at equal quality is the right axis to be optimising, since state capacity is the binding constraint. It does not remove the copying bound; it moves the constant.
For deployment, note that long-context SSM inference is compute-bound rather than memory-bound, which inverts the hardware selection logic in memory-bound vs compute-bound. The right accelerator for a hybrid is not the right accelerator for a transformer.
What to learn next
- Ring attention and sequence parallelism — the other way to reach a million tokens, using many machines.
- Recurrent neural networks — the older family these designs return to.
- LSTM — gates that decide what to forget, invented for the same reason.