Chapter 21 · Section 21.4 · 80 min

Mamba and selective state spaces

Mamba is a recurrent layer that decides, token by token, what to keep in its memory and what to forget. It was the first layer without attention to match a Transformer of the same size on language, and its descendants now sit inside the hybrid models that labs ship. This section builds it from scratch, proves why it works, and counts what it costs on a GPU.

Start from what came before. A state-space model (SSM), the layer of §21.3, keeps a small vector of numbers called its state. At every token it shrinks the state by a fixed factor and adds a fixed multiple of the token. The rule is the same at every step, so a token's influence on later outputs depends only on how long ago it arrived, never on what it is. Mamba lets the token decide. Each token sets its own step size, the number that decides how much of the old state is kept and how much of the new token is written. It also sets which parts of the state it writes to and which parts the output reads.

Here is a task that shows why that change is needed. A sequence of 4,096 tokens is mostly one filler token, and 16 data tokens sit among the filler at random positions. At the end, the model must output the 16 data tokens in the order they appeared. Gu & Dao (2023) trained two-layer models of width 64 on this selective copying task. With S4, the time-invariant SSM of §21.3, the accuracy was 18.3%. With the selective layer of this section it was 97.0%, and with the full Mamba block 99.8%.

That gap should surprise you. S4 was the first model to solve Path-X, a benchmark that requires carrying a signal across 16,384 steps. Remembering 16 tokens is a much smaller job. What S4 cannot do is choose which 16 to remember.

Here is the plan. We first prove that no layer that is linear in its input can solve even a three-token version of the task. Then we let each token choose its step size, and prove that this turns an SSM into a recurrent network with an input-dependent gate. The price is the fast convolution of §21.3. So we implement the scan that replaces it, count the memory traffic that forces the scan into one fused GPU kernel, build and count the full Mamba block, and decode with a cache whose size never changes. The section ends with what Mamba still cannot do, which is where §21.5 and §21.6 pick up.

Why a time-invariant layer cannot select

A time-invariant layer applies the same rule at every step. §21.3 showed that a time-invariant SSM is a long convolution. Take one channel with input $x_t$ . Write $\bar{\mathbf{A}}$ for the matrix that shrinks the state at each step, $\bar{\mathbf{B}}$ for the vector that writes the input into the state, and $\mathbf{C}$ for the vector that reads the output from it. Then the output is

y_t = \sum_{j\ge 0}\bar K_j\,x_{t-j}, \qquad \bar K_j = \mathbf{C}\bar{\mathbf{A}}^j\bar{\mathbf{B}} .

In plain words, the output is a weighted sum of past inputs, and the weight on the input $j$ steps back is a fixed number $\bar K_j$ . The weight with which token $s$ reaches output $t$ is $\bar K_{t-s}$ . It depends on the distance $t - s$ and on nothing else.

That is enough for the plain copying task, where the tokens to copy sit at fixed positions. A kernel with a spike at the right distance copies them, and convolutions as long as the sequence are known to solve it (Romero et al. 2021). Selective copying randomizes the spacing. The right distance changes from sequence to sequence, and only the content of the tokens says where the data are.

Shrink the task to its smallest form, recall the last data token. Each token is either filler or data, and each data token carries a value. The target at step $t$ is the value of the most recent data token at or before $t$ . For the values $(3, \text{filler}, 7, \text{filler})$ the targets are $(3, 3, 7, 7)$ . A layer that solves it must give the newest data token weight 1 and every older token weight 0, and which token is newest depends on the input. The obstacle already shows with three tokens.

∴ No layer that is linear in its input recalls the last data token

Let every token $w$ be embedded as a vector $\mathbf{e}(w) \in \R^m$ (any embedding), and let a layer compute its output at step $t$ as an affine function (linear plus a constant) of the embeddings at steps $1, \dots, t$ ,

y_t = \beta_t + \sum_{s=1}^{t}\mathbf{k}_{t,s}^\top\mathbf{e}(w_s),

with any vectors $\mathbf{k}_{t,s} \in \R^m$ and offsets $\beta_t$ . This covers every time-invariant SSM and long convolution (where $\mathbf{k}_{t,s}$ depends only on $t - s$ ), with any state size, followed by a linear readout. If there are two data tokens with different values, no such layer outputs the value of the most recent data token at step 3 for every sequence of three tokens.

Proof. Let $\nu$ be the filler token and $p$ , $u$ , $w$ data tokens with $\operatorname{val}(u) \ne \operatorname{val}(p)$ . Look at step $t = 3$ for four sequences,

S_1 = (p, u, w), \quad S_2 = (p, \nu, w), \quad S_3 = (p, u, \nu), \quad S_4 = (p, \nu, \nu),

whose targets are $\operatorname{val}(w)$ , $\operatorname{val}(w)$ , $\operatorname{val}(u)$ , and $\operatorname{val}(p)$ . The pairs $(S_1, S_2)$ and $(S_3, S_4)$ each differ only at step 2, where $u$ replaces $\nu$ . So in both pairs the output changes by the same amount,

y_3(S_1) - y_3(S_2) = \mathbf{k}_{3,2}^\top\big(\mathbf{e}(u) - \mathbf{e}(\nu)\big) = y_3(S_3) - y_3(S_4).

An exact solution needs the left side to be $\operatorname{val}(w) - \operatorname{val}(w) = 0$ and the right side to be $\operatorname{val}(u) - \operatorname{val}(p) \ne 0$ , a contradiction. $\square$

In plain words, in a linear layer the contribution of token $u$ is the same whether or not a newer data token follows it. Recall needs that contribution to vanish when $w$ arrives and to survive when filler arrives. Only a layer whose mixing weights depend on the tokens themselves can do both.

Notice what the proof did not use. It never assumed time invariance, so a linear layer with different weights at every position fails too. The defect is linearity in the input.

Two kinds of layer escape the theorem, and neither escapes it well. The first adds nonlinear MLPs after the linear layers. §21.1 quoted a theorem that a linear recurrence followed by an MLP can approximate any well-behaved causal sequence map. The theorem says nothing about how large the state must be. Here the linear state would have to keep enough of the whole history for the MLP to find the newest data token in it, and that grows with the sequence length. The second is a gated layer, which multiplies the output of a time-invariant SSM or long convolution by a projection of the input. Two such layers met in §21.3 are H3 (Fu et al. 2022) and Hyena (Poli et al. 2023). The product is nonlinear, but it acts at one position at a time. It never changes how far a token's information travels along the sequence. In the same selective copying experiment, H3 around an S4 layer reached 57.0% and Hyena 30.1%.

If an interviewer asks why S4 struggled on language, this is the answer they want. A time-invariant layer weighs tokens by position, and language needs weights that depend on content.

Letting each token choose its step size

An SSM starts life as a system in continuous time, and §21.3 turned it into a recurrence over tokens. Zero-order hold (ZOH) treats each token as an input held constant for a stretch of time of length $\Delta$ , the step size, and solves the system exactly over that stretch. Each state number (a mode) has a rate $\lambda < 0$ , and over one step it shrinks by the factor $\bar a = e^{\Delta\lambda}$ . So $\Delta$ sets a timescale. A small $\Delta$ keeps the state, since $\bar a \approx 1$ . A large $\Delta$ forgets it, since $\bar a \approx 0$ . In S4 the same $\Delta$ is used at every step.

Mamba's idea fits in one sentence. Let each token choose its own $\Delta$ . A token the model wants to ignore picks $\Delta \approx 0$ , which leaves the state untouched. A token it wants to store picks a large $\Delta$ , which flushes the old state and writes the token.

Start with a case you can do by hand, the leaky bucket of §21.3. Its level $h$ drains in proportion to how full it is and fills from the input, $\dot h = -h + x$ , where $\dot h$ is the rate of change. So there is one state ( $N = 1$ ), with $A = -1$ and $B = C = 1$ . ZOH with a step $\Delta_t$ that may differ at every token gives

\bar a_t = e^{-\Delta_t}, \qquad \bar b_t = 1 - e^{-\Delta_t}, \qquad h_t = \bar a_th_{t-1} + \bar b_tx_t, \qquad y_t = h_t .

In plain words, each token keeps the share $\bar a_t$ of the old state and writes itself with the weight $\bar b_t$ , and the two shares add up to one.

✏️ Selective copying of one value, by hand

Feed five tokens with values $x = (3, 9, 9, 7, 9)$ , where tokens 1 and 4 are data and the 9s are filler. Let data tokens choose $\Delta = 5$ and filler tokens $\Delta = 0$ . On data, $\bar a = e^{-5} = 0.0067$ and $\bar b = 0.9933$ . On filler, $\bar a = 1$ and $\bar b = 0$ . Starting from $h_0 = 0$ ,

h_1 = 0.9933\cdot 3 = 2.980, \quad h_2 = h_3 = 2.980, \quad h_4 = 0.0067\cdot 2.980 + 0.9933\cdot 7 = 6.973, \quad h_5 = 6.973 .

The output tracks the target $(3, 3, 3, 7, 7)$ to within 0.7%, and the two filler 9s never touch the state. The small error is the share $e^{-5}$ of the old value that survives a write, and it vanishes as the data token's $\Delta$ grows. No single fixed $\Delta$ can do this, because the filler would then leak in with weight $1 - e^{-\Delta} > 0$ .

Python · cell 1

Now scale the same comparison up. The next cell draws 400 sequences of 48 tokens. Each token is data with probability 0.2 and then carries a value from 1 to 8. Filler is embedded as 0, which already removes it from every weighted sum, so the kernel only has to decide how much weight each older data token gets. The cell fits the best possible causal kernel by least squares, which gives the time-invariant layer every advantage, and compares it with the selective rule.

Python · cell 2

This is selection in its simplest form, and the first thing to say if someone asks what Mamba changed relative to S4. The step size stops being a fixed timescale and becomes a decision that each token makes.

The selective state-space model

A real layer runs many channels side by side and keeps more than one state number per channel. Here is the general model in words first. Each channel keeps $N$ state numbers, and each of them decays at its own fixed rate. At every token, three things are computed from the token itself. A step size per channel decides how much to keep and how much to write. A write vector $\mathbf{B}_t$ decides which state numbers receive the token. A read vector $\mathbf{C}_t$ decides which state numbers form the output.

This is the diagonal SSM of §21.3 (S4D, in which each channel's state matrix is diagonal) with three of its parameters computed from the input at every step. Gu and Dao call it S6, for S4 with selection, computed with a scan.

📐 Selective state-space model (S6)

With $d_{\text{in}}$ channels and state size $N$ per channel, the layer maps an input sequence $\mathbf{x}_1, \dots, \mathbf{x}_T \in \R^{d_{\text{in}}}$ to outputs $\mathbf{y}_1, \dots, \mathbf{y}_T \in \R^{d_{\text{in}}}$ . Write $x_t^{(c)}$ for channel $c$ of $\mathbf{x}_t$ . Its parameters are a log-rate matrix $\boldsymbol\ell \in \R^{d_{\text{in}}\times N}$ , which sets the state eigenvalues $\lambda_{cn} = -e^{\ell_{cn}} < 0$ ; projections through a small rank $r_\Delta$ , $\mathbf{W}_{\downarrow} \in \R^{d_{\text{in}}\times r_\Delta}$ , $\mathbf{W}_{\uparrow} \in \R^{r_\Delta\times d_{\text{in}}}$ , $\mathbf{b}_\Delta \in \R^{d_{\text{in}}}$ , and $\mathbf{W}_B, \mathbf{W}_C \in \R^{d_{\text{in}}\times N}$ ; and a skip vector $\mathbf{D} \in \R^{d_{\text{in}}}$ . At every step $t$ it computes

\boldsymbol\Delta_t = \operatorname{softplus}\!\big(\mathbf{W}_{\uparrow}^\top\mathbf{W}_{\downarrow}^\top\mathbf{x}_t + \mathbf{b}_\Delta\big) \in \R^{d_{\text{in}}}, \qquad \mathbf{B}_t = \mathbf{W}_B^\top\mathbf{x}_t \in \R^N, \qquad \mathbf{C}_t = \mathbf{W}_C^\top\mathbf{x}_t \in \R^N,

with $\operatorname{softplus}(v) = \log(1 + e^v)$ applied entrywise, so every step size is positive. Then, for every channel $c$ and state index $n$ , starting from $h_{0,cn} = 0$ ,

\bar a_{t,cn} = e^{\Delta_{t,c}\lambda_{cn}}, \qquad h_{t,cn} = \bar a_{t,cn}\,h_{t-1,cn} + \Delta_{t,c}\,B_{t,n}\,x_t^{(c)}, \qquad y_t^{(c)} = \sum_{n=1}^{N} C_{t,n}\,h_{t,cn} + D_c\,x_t^{(c)} .

The state $\mathbf{H}_t = (h_{t,cn}) \in \R^{d_{\text{in}}\times N}$ holds $N$ numbers per channel. The step sizes are per token and per channel. The vectors $\mathbf{B}_t$ and $\mathbf{C}_t$ are per token and shared by all $d_{\text{in}}$ channels (they are written as capitals to match the SSM literature, although they are vectors). The eigenvalues are real, learned, and the same at every step.

Look at one channel $c$ and one state index $n$ . The update is the scalar linear recurrence of §21.1, $h_t = a_th_{t-1} + u_tx_t$ , with decay $a_t = \bar a_{t,cn}$ and input weight $u_t = \Delta_{t,c}B_{t,n}$ . What is new is that these coefficients are now functions of $\mathbf{x}_t$ . They depend on the input and not on the state, so the recurrence stays linear in its state. Every tool of §21.1 still applies, the loop, the mixing matrix, and the scan.

In matrix form the whole layer is two lines,

\mathbf{H}_t = \exp\!\big(\boldsymbol\Delta_t\boldsymbol\lambda\big)\odot\mathbf{H}_{t-1} + (\boldsymbol\Delta_t\odot\mathbf{x}_t)\,\mathbf{B}_t^\top, \qquad \mathbf{y}_t = \mathbf{H}_t\mathbf{C}_t + \mathbf{D}\odot\mathbf{x}_t ,

where $\boldsymbol\Delta_t\boldsymbol\lambda$ means the matrix with entries $\Delta_{t,c}\lambda_{cn}$ , $\exp$ acts entrywise, and $\odot$ is the entrywise product.

In plain words, each step shrinks every entry of a $d_{\text{in}}\times N$ state at its own rate, adds the outer product of the step-scaled input with $\mathbf{B}_t$ , and reads the state against $\mathbf{C}_t$ . Linear attention (§21.2), attention with the softmax removed so that it runs as a recurrence, has the same shape. It keeps a matrix state, adds the outer product of a key and a value at every token, and reads the state with a query. Here $\mathbf{B}_t$ plays the key, $\mathbf{C}_t$ the query, and $\boldsymbol\Delta_t\odot\mathbf{x}_t$ the value, with a decay on every entry. (In §21.2's layout the state would be $\mathbf{H}_t^\top$ , with rows indexing keys.) §21.6 turns that resemblance into a theorem.

The discretization the code actually uses

The definition writes the input weight as $\Delta_{t,c}B_{t,n}$ , not as the ZOH formula. That is deliberate, because it is what the released code computes. For one mode with eigenvalue $\lambda < 0$ , write $\kappa = \Delta|\lambda|$ , the step measured in units of that mode's timescale. ZOH on both parameters gives

\bar a = e^{-\kappa}, \qquad \bar b_{\text{ZOH}} = \frac{e^{\Delta\lambda} - 1}{\lambda}B = \frac{1 - e^{-\kappa}}{\kappa}\,\Delta B .

The Mamba paper states this ZOH rule. The released kernels compute something slightly different. The reference function selective_scan_ref in the authors' repository, and the CUDA kernel tested against it, keep the exact decay $\bar a = e^{\Delta\lambda}$ but use the input weight $\Delta B$ . That is the first term of the expansion $(1 - e^{-\kappa})/\kappa = 1 - \kappa/2 + \kappa^2/6 - \cdots$ . The Mamba-3 paper later derived this rule and named it exponential-Euler (Lahoti et al. 2026).

Neither rule is wrong, since a network learns its weights under whichever rule it is trained with. But they define different models. A checkpoint trained with one gives different outputs if run with the other, mostly in the modes with $\Delta|\lambda|$ near 1 or above. That is the practical reason to know which one your code uses.

🔬 Where exponential-Euler comes from, and what it changes

The numbers behind the fact. A fresh layer uses the eigenvalues $\lambda_n = -n$ for $n = 1, \dots, 16$ (the paper writes the same values as $-(n+1)$ for $n = 0, \dots, 15$ ), and step sizes between 0.001 and 0.1. So $\kappa$ ranges from 0.001 to 1.6, and at $\kappa = 1.6$ the ratio $(1 - e^{-\kappa})/\kappa$ is 0.50.

The derivation. §21.3 solved a linear system exactly over one step by the variation-of-constants formula. For one mode, the update is $h_t = e^{\Delta\lambda}h_{t-1} + \big(\int_0^\Delta e^{\lambda\tau}\,d\tau\big)Bx_t$ , where $\tau$ is the time remaining until the end of the step. ZOH evaluates the integral exactly. Exponential-Euler keeps the state's exponential exact and approximates the integral by Euler's rule, the integrand's value at the end of the step ( $\tau = 0$ , where $e^{\lambda\cdot 0} = 1$ ) times the step length $\Delta$ . That gives $\Delta B$ . Mamba-3 replaces this one-point rule by a data-dependent blend of both ends of the step, a trapezoid-like rule that §21.6 describes. It is second-order accurate when the blend weight stays near one half, a constraint the authors found works better left unenforced.

What a huge step means. Under ZOH the write saturates, $\bar b \to B/|\lambda|$ as $\Delta\to\infty$ , so a huge step means "replace the state by the token". Under exponential-Euler the write keeps growing like $\Delta B$ , so a huge step means "replace the state by $\Delta$ times the token", and $\mathbf{C}_t$ or the next layer must cope with that scale.

Selection is a gate

The step size behaves exactly like a gate in the recurrent networks of §3.3. A gated recurrent unit (GRU) updates its state as a blend, $h_t = (1 - g_t)h_{t-1} + g_t\tilde h_t$ . The update gate $g_t \in (0, 1)$ is a sigmoid of the input and the previous state, and $\tilde h_t$ is a candidate new state. A gate near 0 keeps the old state, and a gate near 1 replaces it. (§3.3 writes the GRU with the roles of $g_t$ and $1 - g_t$ swapped, so check the convention of any formula you meet.)

Now put the leaky bucket through ZOH with $\Delta_t = \operatorname{softplus}(z_t)$ , where $z_t$ is a number computed from the token, and try a number. Take $z_t = 0$ . Then $\Delta_t = \log 2 = 0.693$ , the decay is $\bar a_t = e^{-0.693} = 0.5$ , and the write is $\bar b_t = 1 - 0.5 = 0.5$ . The logistic sigmoid of 0 is also 0.5. Take $z_t = 2$ . Then $\Delta_t = \log(1 + e^2) = 2.127$ and $\bar a_t = e^{-2.127} = 0.119$ , while $1 - \sigma(2) = 1 - 0.881 = 0.119$ . The coincidence is exact, and it is the reason Mamba passes $\Delta$ through a softplus.

∴ Selection is a gate (Gu & Dao 2023, Theorem 1)

Take one channel and one state, $N = 1$ , $\lambda = -1$ , $B = 1$ , and step size $\Delta_t = \operatorname{softplus}(z_t)$ , where $z_t = \mathbf{w}^\top\mathbf{x}_t + b$ is any affine function of the input. Discretize both parameters by ZOH. Then

h_t = (1 - g_t)\,h_{t-1} + g_t\,x_t, \qquad g_t = \sigma(z_t),

with $\sigma$ the logistic sigmoid. The selective SSM is a gated recurrence with an input-dependent gate.

Proof. The decay is $\bar a_t = e^{\Delta_t\lambda} = e^{-\operatorname{softplus}(z_t)} = e^{-\log(1 + e^{z_t})} = \frac{1}{1 + e^{z_t}}$ . Since $\sigma(z) = \frac{1}{1 + e^{-z}} = \frac{e^z}{1 + e^z}$ , we have $1 - \sigma(z) = \frac{1}{1 + e^z}$ , so $\bar a_t = 1 - \sigma(z_t)$ . The ZOH input weight is $\bar b_t = \frac{e^{\Delta_t\lambda} - 1}{\lambda}\cdot B = \frac{\bar a_t - 1}{-1} = 1 - \bar a_t = \sigma(z_t)$ . Substituting both into $h_t = \bar a_th_{t-1} + \bar b_tx_t$ gives the claim. $\square$

In plain words, the softplus and the exponential cancel into a sigmoid, so the discretized leaky bucket is a GRU's update gate. Unlike a GRU's, this gate is computed from the current input only, and the candidate is the input itself. The bias $b$ plays the role of $\mathbf{b}_\Delta$ in the definition, the learned constant part of $\Delta$ .

With the kernels' exponential-Euler rule, the decay is the same $1 - \sigma(z_t)$ , but the write becomes $\bar b_t = \Delta_t = \operatorname{softplus}(z_t)$ . For very negative $z_t$ both $\sigma(z_t)$ and $\operatorname{softplus}(z_t)$ are about $e^{z_t}$ , so the two rules agree there. For large $z_t$ the softplus grows like $z_t$ . The code's layer is a gate whose write is not capped at 1.

The theorem is the standard answer when an interviewer asks how Mamba relates to an LSTM or a GRU. It also explains a design choice. Mamba computes $\Delta$ through a projection of low rank $r_\Delta$ . It squeezes $\mathbf{x}_t$ to $r_\Delta$ numbers with $\mathbf{W}_{\downarrow}$ and expands them back to one step size per channel with $\mathbf{W}_{\uparrow}$ . The squeeze makes the channels share their decision, because if a token should be ignored, every channel should ignore it. The paper motivates this with a projection of $\mathbf{x}_t$ to a single number broadcast to all channels. The code generalizes it to rank $r_\Delta = \lceil d/16\rceil$ for model width $d$ , which can still express one shared decision and costs few parameters. Perplexity, the exponential of the average next-token loss (lower is better), measures the effect. In the paper's ablations on models of about 350M parameters and width 1,024, it was 9.12 with no input-dependent $\Delta$ , 8.97 with rank 1, and 8.71 with rank 64, the default $\lceil d/16\rceil$ at that width (Gu & Dao 2023). Rank 1 already gives over a third of the gain, and the rest comes gradually as the rank grows, for under 4% more parameters.

Warm-up

Selection is a gate, numerically

Implement two functions.

zoh_scalar(delta, A, B) returns the pair (a_bar, b_bar) of the ZOH discretization of the scalar system $\dot h = Ah + Bx$ with $A < 0$ . Its inputs may be NumPy arrays of the same shape.
selective_gate_scan(z, x) runs the setting of the theorem. Given arrays z and x of length $T$ , use $\Delta_t = \operatorname{softplus}(z_t)$ , $A = -1$ , $B = 1$ , and ZOH for both parameters. Return the states $h_1, \dots, h_T$ as an array of shape (T,), starting from $h_0 = 0$ .

Extreme inputs such as $z_t = \pm 1000$ must give finite states. The tests check that your scan equals the gated recurrence of the theorem.

import numpy as np

def zoh_scalar(delta, A, B):
    raise NotImplementedError

def selective_gate_scan(z, x):
    raise NotImplementedError

Tests, verified solutions, and AI help, open once you add your API key.

What each parameter does

With the theorem in hand, each selective parameter has a clear job.

The step size $\Delta$ sets the balance between keeping the state and taking in the current token. As $\Delta_{t,c}\to 0$ the decay tends to 1 and the write to 0, so the token is ignored and the state passes through untouched. As $\Delta_{t,c}\to\infty$ the decay tends to 0, so the old state is flushed and the token written. Values in between give a partial update.

A large $\Delta$ also lets the layer reset. Training packs several documents into one sequence, and the state should forget everything at each document break (§4.4 built masks for the same job in attention). One token with a large $\Delta$ does that. Its decay near 0 zeroes every entry of the mixing matrix that links the two documents, as in the reset example of §21.1.

The eigenvalues $\lambda_{cn}$ enter only through the products $\Delta_{t,c}\lambda_{cn}$ . Making them input-dependent too would add little that $\Delta$ does not already provide, so the paper keeps them fixed. What they provide is a spread of timescales for each step size. With the initialization $\lambda_{cn} = -n$ (called S4D-Real), one channel's $N$ modes forget at rates $1, 2, \dots, N$ times $\Delta$ . A stored value in the mode with $\lambda = -n$ halves after $\ln 2/(n\Delta)$ steps, its half-life (§21.3).

The vectors $\mathbf{B}_t$ and $\mathbf{C}_t$ decide what enters the state and what is read out. A selective $\mathbf{B}_t$ can write a token into some state entries and not others, and a selective $\mathbf{C}_t$ can read the entry that matters now. They are the content side of selection, and they become important once $N > 1$ , as the ablations later in this section show.

The widget lets you play the recall game yourself on a strip of 24 tokens. In time-invariant mode the layer uses one step size for every token and up to four modes. Its readout $\mathbf{C}$ is fitted by least squares to the current strip, the best those fixed modes can do. In selective mode, filler and data tokens get their own step sizes. Filler tokens enter the layer as 0, and the row of green bars shows how much of the state each token replaces, $1 - \bar a = 1 - e^{-\Delta}$ . Click any token to switch it between data and filler, and press Play to watch the layer read the strip one token at a time.

Selection breaks the convolution

Selection has a price, and the mixing matrix of §21.1 shows exactly what it is. Recall what that section proved. A scalar linear recurrence $h_t = a_th_{t-1} + u_tx_t$ with readout $y_t = c_th_t$ can be unrolled, so each output is a weighted sum of all past inputs. The weight on input $s$ is its input weight $u_s$ , times the product of every decay applied after it, $a_{t:s} = a_ta_{t-1}\cdots a_{s+1}$ (the segment product), times the readout weight $c_t$ . These weights fill a lower-triangular mixing matrix $\mathbf{M}$ with $\mathbf{y} = \mathbf{M}\mathbf{x}$ .

Apply that to one channel $c$ and one mode $n$ of the selective SSM, with decay $\bar a_{t,cn}$ and input weight $\Delta_{t,c}B_{t,n}$ , and sum over modes with the readout weights $C_{t,n}$ .

∴ The mixing matrix of a selective SSM

For each channel $c$ , the selective SSM computes $y_t^{(c)} = \sum_{s\le t}\mathbf{M}^{(c)}_{ts}x_s^{(c)} + D_cx_t^{(c)}$ with the lower-triangular matrix

\mathbf{M}^{(c)}_{ts} = \sum_{n=1}^{N} C_{t,n}\Big(\prod_{r=s+1}^{t}\bar a_{r,cn}\Big)\Delta_{s,c}B_{s,n} = \sum_{n=1}^{N} C_{t,n}\exp\!\Big(\lambda_{cn}\sum_{r=s+1}^{t}\Delta_{r,c}\Big)\Delta_{s,c}B_{s,n}, \qquad s\le t .

If $\Delta$ , $\mathbf{B}$ , and $\mathbf{C}$ are the same at every step, $\mathbf{M}^{(c)}_{ts}$ depends only on $t - s$ , and the layer is the convolution of §21.3. Otherwise it is not a convolution.

Proof. For fixed $c$ and $n$ , the recurrence $h_{t,cn} = \bar a_{t,cn}h_{t-1,cn} + (\Delta_{t,c}B_{t,n})x_t^{(c)}$ is the scalar linear recurrence. Its unrolled form is $h_{t,cn} = \sum_{s\le t}a_{t:s}\,\Delta_{s,c}B_{s,n}\,x_s^{(c)}$ with the segment product $a_{t:s} = \prod_{r=s+1}^{t}\bar a_{r,cn}$ . Multiply by $C_{t,n}$ and sum over $n$ . A product of exponentials is the exponential of the sum, which gives the second form. If every parameter is constant, the segment product is $e^{\lambda_{cn}\Delta(t-s)}$ , and the whole entry is a function of $t - s$ . $\square$

In plain words, the weight from token $s$ to output $t$ now depends on three things. It depends on what token $s$ wrote ( $\Delta_s\mathbf{B}_s$ ), on how much every token in between let the state decay (the sum of their step sizes), and on what output $t$ reads ( $\mathbf{C}_t$ ). That is what selection needs, and it is exactly what kills the convolution.

A convolution is a Toeplitz matrix, one that is constant along each diagonal. The FFT algorithm of §21.3 applies it in $O(T\log T)$ time by transforming its single kernel. The selective layer's matrix depends on $(t, s)$ , not only on $t - s$ , so there is no single kernel to transform.

Two ways of computing the layer survive. The recurrence runs in $O(TNd_{\text{in}})$ time and keeps only the current state. The parallel scan of §11.4 gives the same result in $O(\log T)$ parallel steps. A scan computes every prefix of an associative operation, and here each step is the affine map $h \mapsto a h + b$ . Applying $(a_1, b_1)$ and then $(a_2, b_2)$ is again an affine map, $(a_1, b_1)\bullet(a_2, b_2) = (a_1a_2, a_2b_1 + b_2)$ , and this rule is associative. The scan runs on $d_{\text{in}}N$ independent pairs at once, one per channel and state entry.

The matrix $\mathbf{M}^{(c)}$ also helps interpretation. It is a data-dependent, lower-triangular mixing matrix, like attention's but with no softmax. It is also highly structured, since every block that lies within its lower triangle has rank at most $N$ , a property §21.6 calls semiseparable. Ali, Zimerman & Wolf (2024) plot it as Mamba's "hidden attention" and apply attention-style interpretability tools to it. The next cell builds the matrix for a small example, checks it against the scan, and tests whether it is a convolution.

Python · cell 3

The selective scan from scratch

The function selective_scan_ref above is the whole layer, and it maps line by line to the definition. The decay factor is exp(delta[t][:, None] * A), one number per channel and state entry. The write is the outer product of $\boldsymbol\Delta_t\odot\mathbf{x}_t$ with $\mathbf{B}_t$ , and the read is $\mathbf{H}_t\mathbf{C}_t$ plus the skip term. It costs $O(Td_{\text{in}}N)$ arithmetic and $O(d_{\text{in}}N)$ memory, with a Python loop over $T$ that the GPU kernel replaces. Writing it yourself, with a batch dimension and the kernels' exact discretization, is a standard interview problem.

Interview

The selective scan, from scratch

Implement selective_scan(x, delta, A, B, C, D), the selective SSM of the definition with the discretization of the released kernels. The shapes are as follows.

x and delta have shape (T, d_in), or (batch, T, d_in). The step sizes are already positive, because the softplus has been applied.
A has shape (d_in, N) and holds the negative eigenvalues $\lambda_{cn}$ .
B and C have shape (T, N), or (batch, T, N), and are shared by all channels.
D has shape (d_in,).

Start from $\mathbf{H}_0 = \mathbf{0}$ . Compute $\mathbf{H}_t = \exp(\boldsymbol\Delta_t\boldsymbol\lambda)\odot\mathbf{H}_{t-1} + (\boldsymbol\Delta_t\odot\mathbf{x}_t)\mathbf{B}_t^\top$ and $\mathbf{y}_t = \mathbf{H}_t\mathbf{C}_t + \mathbf{D}\odot\mathbf{x}_t$ . The input weight is exactly $\Delta B$ , not the ZOH formula. Return y with the shape of x. A loop over $T$ is fine, and everything else should be vectorized.

import numpy as np

def selective_scan(x, delta, A, B, C, D):
    raise NotImplementedError

Tests, verified solutions, and AI help, open once you add your API key.

The theorem at the start of this section proved that a linear layer cannot recall the last data token. The next exercise turns that proof into a measurement and asks you to build the selective layer that solves the task exactly.

Research

Why a fixed kernel cannot select

The setup provides make_recall(n, T, p, seed). It returns X (the values, 0 for filler), is_data (booleans), and the target Y, each of shape (n, T). Y[i, t] is the value of the most recent data token at or before t in sequence i, and 0 before the first one. Implement two functions.

best_lti_kernel(X, Y) fits the best causal kernel $\mathbf{K} \in \R^T$ , with $\hat y_t = \sum_{j=0}^{t}K_jx_{t-j}$ , by least squares over every sequence and step. It returns (K, rel_err), where rel_err is $\sum(\hat y - y)^2/\sum y^2$ .
selective_solution(X, is_data) returns arrays (a_bar, b_bar) of shape (n, T) such that the recurrence $h_t = \bar a_th_{t-1} + \bar b_tx_t$ , $h_0 = 0$ , outputs $h_t = y_t$ exactly. The coefficients may depend on whether a token is data, never on its position.

In a comment, explain with the four-sequence argument of this section why no kernel length or data set size can drive rel_err to zero.

import numpy as np

def best_lti_kernel(X, Y):
    raise NotImplementedError

def selective_solution(X, is_data):
    raise NotImplementedError

Tests, verified solutions, and AI help, open once you add your API key.

Why the scan must be fused

A GPU has two kinds of memory that matter here. HBM (high-bandwidth memory) is the large main memory, 80 GB on an H100, read at about 3.35 TB/s. On-chip memory (registers and shared memory, often called SRAM) holds only a few hundred kilobytes per streaming multiprocessor, but it is far faster (§11.1). The roofline model of §11.2 says an operation takes at least the longer of two times, its bytes moved divided by the bandwidth, and its arithmetic divided by the peak rate. An elementwise operation does a few operations per byte, so it is limited by the bytes.

On paper the recurrence is cheap, $O(BTd_{\text{in}}N)$ arithmetic for a batch of $B$ sequences. (Italic $B$ is the batch size here and below, unrelated to the bold vector $\mathbf{B}_t$ .) The trouble is memory traffic. Over a whole sequence the states form a tensor of shape $(B, T, d_{\text{in}}, N)$ , which is $N$ times the size of the input $(B, T, d_{\text{in}})$ .

Count the bytes of the straightforward implementation, the one you would write with PyTorch operations. It computes the decays $\bar{\mathbf{A}}$ and the writes $\Delta\mathbf{B}x$ as two tensors of shape $(B, T, d_{\text{in}}, N)$ and stores them in HBM. It runs a scan that reads both and writes all the states $\mathbf{H}_t$ back to HBM. Then it reads the states again to multiply by $\mathbf{C}_t$ . That is six passes over a state-sized tensor. A fused kernel does everything in one pass instead. It reads $\Delta$ , $\mathbf{x}$ , $\mathbf{B}$ , $\mathbf{C}$ , and $\boldsymbol\lambda$ from HBM into on-chip memory, discretizes and runs the recurrence there, and writes only $\mathbf{y}$ .

∴ HBM traffic of the selective scan

Let $S = BTd_{\text{in}}N\cdot s_{\text{state}}$ be the bytes of one state-sized tensor in fp32 ( $s_{\text{state}} = 4$ ), and $I = BTd_{\text{in}}\cdot s_{\text{io}}$ the bytes of one input-sized tensor in bf16 ( $s_{\text{io}} = 2$ ). Ignore the small tensors $\mathbf{B}$ and $\mathbf{C}$ (of size $BTN$ ) and $\boldsymbol\lambda$ . Then the forward pass moves

\text{unfused: } 6S + 3I, \qquad \text{fused: } 3I \quad (\text{read } \mathbf{x} \text{ and } \Delta, \text{ write } \mathbf{y}),

so the fused kernel moves $1 + 2Ns_{\text{state}}/s_{\text{io}} = 1 + 4N$ times fewer bytes, a factor of 65 at $N = 16$ .

Proof. Count each tensor once per read or write. The unfused path writes $\bar{\mathbf{A}}$ and the write term ( $2S$ ), reads both in the scan ( $2S$ ), writes the states ( $S$ ), and reads them for the readout ( $S$ ). Both paths also read $\mathbf{x}$ and $\Delta$ and write $\mathbf{y}$ ( $3I$ ). The ratio is $(6S + 3I)/3I = 1 + 2S/I = 1 + 2Ns_{\text{state}}/s_{\text{io}}$ . $\square$

In plain words, storing the states costs bytes in proportion to $N$ , and fusing removes all of them. What remains is arithmetic on data that never leave the chip, and it comes in two kinds. Per state element, the kernel does about 6 floating-point operations on the ordinary fp32 units (the product $\Delta\lambda$ , the write $\Delta x\cdot B$ , a multiply-add for the decay, and a multiply-add for the readout). It also computes one exponential, $e^{\Delta\lambda}$ .

The exponential does not run on the ordinary units. As §11.1 showed, exp runs on the special-function units, which on an H100 deliver only about $3.9\times10^{12}$ results per second in total. That is 17 times fewer than the 67 TFLOP/s of the fp32 units. So one scan takes at least

t \ \ge\ \max\Big(\frac{\text{bytes}}{3.35\times10^{12}},\ \frac{\text{FLOPs}}{67\times10^{12}},\ \frac{\text{exponentials}}{3.9\times10^{12}}\Big) \text{ seconds}.

This is the roofline of §11.2 with a third limit for the special functions. Mamba's scan needs one exponential for every sequence, token, channel, and state entry, $BTd_{\text{in}}N$ in all, because every entry of the diagonal state matrix has its own decay.

✏️ One layer of Mamba-2.8B, batch 8 at 2,048 tokens

Take $B = 8$ , $T = 2{,}048$ , $d_{\text{in}} = 5{,}120$ , and $N = 16$ . One fp32 state tensor is $S = 8\cdot 2048\cdot 5120\cdot 16\cdot 4 = 5.37$ GB, and one bf16 input tensor is $I = 0.168$ GB.

Unfused, the forward pass moves about $6S + 3I = 32.7$ GB, which takes at least 9.8 ms at the H100's 3.35 TB/s.
Fused, it moves about 0.5 GB, which is 0.15 ms of traffic.
Its $6\cdot BTd_{\text{in}}N = 8.1\times10^{9}$ fp32 operations need 0.12 ms at 67 TFLOP/s.
Its $BTd_{\text{in}}N = 1.34\times10^{9}$ exponentials need at least 0.34 ms at $3.9\times10^{12}$ per second.

So the fused scan takes at least 0.34 ms, set by the exponentials, not by memory or ordinary arithmetic. The unfused scan is about 28 times slower than that floor.

Python · cell 4

The table shows why Mamba could afford a larger state than earlier recurrent models. Unfused, every doubling of $N$ doubles the time. Fused, the time stays at the 0.15 ms of traffic up to $N \approx 7$ . Above that the exponentials set it. They also grow with $N$ , but at the rate of the special-function units rather than the memory. At $N = 128$ the exponentials alone take 2.8 ms per layer for this batch (the fp32 arithmetic 1.0 ms). That is the regime where Mamba-2 takes over (§21.6), with one decay per head instead of one per state entry and its products on tensor cores.

The widget below draws the same bound for this layer at every state size $N$ from 1 to 256.

The backward pass has the same problem in a sharper form. Autograd needs every intermediate state to compute gradients, and storing them costs $S$ bytes per layer, 5.37 GB in the example above, for each of 64 layers. Mamba's kernel stores none of them. It keeps only the inputs $\mathbf{x}$ , $\Delta$ , $\mathbf{B}$ , and $\mathbf{C}$ . In the backward pass it reloads them into on-chip memory and recomputes the states. FlashAttention, the fused attention kernel of §12.3, uses the same trick for its attention weights. Recomputation adds arithmetic but removes the $O(BTd_{\text{in}}N)$ reads of stored states, so it is faster as well as smaller.

The paper applies recomputation to the rest of the block as well, to the convolution and activation outputs, which are cheap to recompute. Two Mamba layers then need about as much activation memory as one FlashAttention layer plus one MLP (§11.7 covers fusion and recomputation in general). The backward pass of a linear recurrence is itself a linear recurrence run in reverse, which §21.5 derives.

In an interview about Mamba's speed, these are the numbers to give. Fusing saves a factor of about $4N$ in memory traffic, and after fusing, the exponentials set the floor.

The Mamba block

The selective SSM mixes information along the sequence. A network also needs to mix information across channels, which an MLP does in a Transformer. Mamba wraps the SSM in one block that does both, and it stacks that one block everywhere with no separate MLP. Write $\mathbf{u}_t \in \R^d$ for the block's input at step $t$ . It is the normalized residual stream, the running vector that every layer reads and adds its output to (§5.1).

In words, the block expands the input to twice the width and splits it into two halves. One half goes through a short convolution over the last four tokens and then the SSM. The other half becomes a gate that multiplies the result. A final projection maps back to width $d$ .

📐 The Mamba block

With model width $d$ , inner width $d_{\text{in}} = 2d$ , state size $N$ , and a causal depthwise convolution of width $k = 4$ , the block maps $\mathbf{u}_1, \dots, \mathbf{u}_T \in \R^d$ to outputs in $\R^d$ by

\begin{aligned} [\tilde{\mathbf{x}}_t;\ \mathbf{z}_t] &= \mathbf{W}_{\text{in}}^\top\mathbf{u}_t, && \mathbf{W}_{\text{in}} \in \R^{d\times 2d_{\text{in}}},\\ \mathbf{x}_t &= \operatorname{SiLU}\Big(\sum_{i=0}^{k-1}\mathbf{w}_i\odot\tilde{\mathbf{x}}_{t-i} + \mathbf{b}_{\text{conv}}\Big), && \mathbf{w}_i, \mathbf{b}_{\text{conv}} \in \R^{d_{\text{in}}},\ \tilde{\mathbf{x}}_{s} = \mathbf{0} \text{ for } s \le 0,\\ \mathbf{y}_t &= \operatorname{S6}(\mathbf{x}_1, \dots, \mathbf{x}_t), && \text{the selective SSM, with } r_\Delta = \lceil d/16\rceil,\\ \mathbf{o}_t &= \mathbf{W}_{\text{out}}^\top\big(\mathbf{y}_t\odot\operatorname{SiLU}(\mathbf{z}_t)\big), && \mathbf{W}_{\text{out}} \in \R^{d_{\text{in}}\times d}, \end{aligned}

where $\operatorname{SiLU}(v) = v\,\sigma(v)$ entrywise. A Mamba layer is the pre-norm residual update $\mathbf{r}_t \leftarrow \mathbf{r}_t + \mathbf{o}_t$ with $\mathbf{u}_t = \operatorname{RMSNorm}(\mathbf{r}_t)$ . The network is a token embedding, $L$ such layers, a final RMSNorm, and an output head tied to the embedding.

Read the block as two branches. The main branch expands to $d_{\text{in}}$ , mixes each channel with its few predecessors through the convolution, and then mixes along the whole sequence through the SSM. The gate branch $\operatorname{SiLU}(\mathbf{z}_t)$ multiplies the result entrywise. A SwiGLU MLP (§5.3) has the same shape. It computes two projections of its input and multiplies one by an activated copy of the other. The Mamba block is that gated MLP with an SSM inserted into one branch. RMSNorm is the normalization of §5.2, which rescales a vector by its root mean square.

Gu and Dao built this block by merging two earlier designs into one, the H3 block (an SSM layer between gates) and the gated MLP. With a selective layer inside both, the merged block is simpler than H3 and slightly better (perplexity 8.69 against 8.95 in the paper's architecture ablation).

The short convolution deserves a word, because it is easy to dismiss. It is depthwise, meaning each channel is convolved with its own four weights and channels do not mix. It gives every channel direct access to the last three tokens before the SSM sees them, and $\Delta_t$ , $\mathbf{B}_t$ , and $\mathbf{C}_t$ are computed after it, so they can use that local context. An induction head (§6.8) needs exactly such information. It completes a pattern $[A][B]\dots[A] \to [B]$ and so must know which token came just before each position. A Transformer gets that from a previous-token head, an attention head that looks one position back. H3 got it from a shift SSM, whose kernel only delays its input by a few steps. In Mamba the four-tap convolution does the same job for almost no parameters.

✏️ Counting Mamba-2.8B, term by term

Mamba-2.8B has $d = 2{,}560$ , $L = 64$ , and $N = 16$ , so $d_{\text{in}} = 5{,}120$ and $r_\Delta = \lceil 2560/16\rceil = 160$ . Its vocabulary has 50,280 entries (Hugging Face config of state-spaces/mamba-2.8b-hf). The input projection has no bias, the convolution has one, the projection to $(\Delta, \mathbf{B}, \mathbf{C})$ has none, and the $\Delta$ up-projection has one. Per layer,

tensor	shape	parameters
$\mathbf{W}_{\text{in}}$	$d\times 2d_{\text{in}}$	26,214,400
convolution weights and bias	$d_{\text{in}}\times k$ , $d_{\text{in}}$	25,600
$[\mathbf{W}_{\downarrow}, \mathbf{W}_B, \mathbf{W}_C]$	$d_{\text{in}}\times(r_\Delta + 2N)$	983,040
$\mathbf{W}_{\uparrow}$ and $\mathbf{b}_\Delta$	$r_\Delta\times d_{\text{in}}$ , $d_{\text{in}}$	824,320
$\boldsymbol\ell$ (the eigenvalues)	$d_{\text{in}}\times N$	81,920
$\mathbf{D}$	$d_{\text{in}}$	5,120
$\mathbf{W}_{\text{out}}$	$d_{\text{in}}\times d$	13,107,200
RMSNorm gain	$d$	2,560
total per layer		41,244,160

Sixty-four layers give 2,639,626,240. The tied embedding adds $50{,}280\cdot2{,}560 = 128{,}716{,}800$ and the final norm 2,560, for 2,768,345,600 in all. That is exactly the count Hugging Face lists for the checkpoint. The same formula gives 129,135,360 for Mamba-130M ( $d = 768$ , $L = 24$ , $r_\Delta = 48$ ), again an exact match.

The projections dominate. $\mathbf{W}_{\text{in}}$ and $\mathbf{W}_{\text{out}}$ hold $3dd_{\text{in}} = 6d^2$ parameters. Everything specific to the SSM ( $\Delta$ , $\mathbf{B}$ , $\mathbf{C}$ , $\boldsymbol\ell$ , $\mathbf{D}$ ) is $d_{\text{in}}(2r_\Delta + 3N + 2) \approx d^2/4 + 6Nd$ . For Mamba-2.8B that is 1.9 million parameters, under 5% of the layer.

A Transformer layer has $4d^2$ parameters in attention and $8d^2$ in an MLP four times as wide, $12d^2$ in all (§6.2). So two Mamba layers stand in for one Transformer layer. That is why Mamba-2.8B has 64 layers at the width where the Transformer Pythia-2.8B has 32.

📌 How a Mamba layer is initialized

In the reference implementation the eigenvalues start at S4D-Real, $\lambda_{cn} = -n$ for $n = 1, \dots, N$ in every channel, stored as $\ell_{cn} = \log n$ in fp32. The bias $\mathbf{b}_\Delta$ is the inverse softplus of step sizes drawn log-uniformly from $[0.001, 0.1]$ , and $\mathbf{W}_{\uparrow}$ is uniform in $\pm r_\Delta^{-1/2}$ . $\mathbf{D}$ starts at 1. Both $\boldsymbol\ell$ and $\mathbf{D}$ are excluded from weight decay.

Real eigenvalues were a deliberate change from S4's complex ones. On language, S4D-Real gave perplexity 8.71 against 9.16 for the complex S4D-Lin initialization, and Mamba used complex values only for raw audio (Gu & Dao 2023).

Core

Count Mamba's parameters and its decode state

Implement mamba_param_count(d, n_layer, N, vocab, d_conv=4, expand=2, dt_rank=None, tied=True). It returns the exact parameter count of a Mamba-1 language model laid out as in the table above. dt_rank=None means $r_\Delta = \lceil d/16\rceil$ . With tied=False, the output head is a separate $d\times$ vocab matrix without bias.

Then implement mamba_state_numbers(d, n_layer, N, d_conv=4, expand=2). It returns the number of values one sequence must keep between decoding steps, which is the SSM state plus the last d_conv - 1 convolution inputs of every channel in every layer. Both functions return Python integers.

def mamba_param_count(d, n_layer, N, vocab, d_conv=4, expand=2, dt_rank=None, tied=True):
    raise NotImplementedError

def mamba_state_numbers(d, n_layer, N, d_conv=4, expand=2):
    raise NotImplementedError

Tests, verified solutions, and AI help, open once you add your API key.

Decoding with a constant-size cache

Generation is where the fixed-size state pays off. To produce each new token, a Transformer reads the keys and values of every earlier token, its KV cache (§6.5), and that cache grows with the context. A Mamba layer needs only two small things from the past, and both have a fixed size.

The convolution cache, the last $k - 1 = 3$ inputs to the convolution, $\tilde{\mathbf{x}}_{t-3}, \tilde{\mathbf{x}}_{t-2}, \tilde{\mathbf{x}}_{t-1}$ , which is $3d_{\text{in}}$ numbers.
The SSM state $\mathbf{H}_{t-1}$ , which is $d_{\text{in}}N$ numbers.

One decoding step runs the block's definition once. Project $\mathbf{u}_t$ to $\tilde{\mathbf{x}}_t$ and $\mathbf{z}_t$ . Append $\tilde{\mathbf{x}}_t$ to the convolution window, and apply the convolution and SiLU. Compute $\Delta_t$ , $\mathbf{B}_t$ , and $\mathbf{C}_t$ , update the state entrywise, read it, gate it, and project out. Then drop the oldest entry of the window. Apart from the projections, the step costs $O(d_{\text{in}}N)$ arithmetic, whatever the context length. Prefill, the processing of the prompt, runs the scan over the whole prompt and keeps only the final state and the last three convolution inputs.

✏️ Mamba-2.8B against Pythia-2.8B in decoding

Mamba-2.8B keeps $64\cdot5{,}120\cdot(16 + 3) = 6{,}225{,}920$ numbers per sequence, 12.5 MB in bf16 or 24.9 MB in fp32, at every context length. Pythia-2.8B (Biderman et al. 2023) is a Transformer of the same width and size trained on the same data, with 32 layers, $d = 2{,}560$ , and 32 attention heads that each keep their own keys and values. It keeps $2\cdot32\cdot2{,}560 = 163{,}840$ numbers per token of context, 320 KiB in fp16, so 640 MiB per sequence at its 2,048-token training context.

The Transformer's cache passes Mamba's entire bf16 state after 38 tokens. Beside about 5.5 GB of bf16 weights, an 80 GB GPU holds 5,979 Mamba sequences but only 110 Pythia sequences at 2,048 tokens, 54 times fewer, and 3 at 65,536 tokens. By the bandwidth model of §21.1, a decoding step costs about the time to read the weights plus every sequence's cache from HBM, so these batch sizes are what set decoding throughput.

Python · cell 5

This is where Mamba's measured speed comes from. Gu and Dao generated 128 tokens after a 2,048-token prompt on one A100 80GB, at batch sizes from 1 to 128. Mamba's throughput was 4 to 5 times that of a Transformer of similar size, because without a KV cache it can run much larger batches. By their measurement, even an untrained Mamba-6.9B would generate faster than a Transformer five times smaller, at 1.3B parameters.

Interview

One decoding step with conv and state caches

The setup provides make_params(d, N, d_conv, expand, seed), the helpers silu and softplus, and mamba_block_forward(params, U). The last one runs the block, without norm or residual, over a whole sequence U of shape (T, d). Read it first, because it is this section's definition in code. Then implement two functions.

init_cache(params) returns the cache before the first token. It is a dict with "conv" of shape (d_conv - 1, d_in), the previous inputs to the convolution with the oldest first, and "h" of shape (d_in, N).
mamba_block_step(params, cache, u) takes one input u of shape (d,) and returns (out, new_cache), where out has shape (d,).

Running mamba_block_step over the tokens of U one at a time must reproduce mamba_block_forward(params, U) at every step. Each step must take O(1) time in the sequence length, with no recomputation over the past.

import numpy as np

def init_cache(p):
    raise NotImplementedError

def mamba_block_step(p, cache, u):
    raise NotImplementedError

Tests, verified solutions, and AI help, open once you add your API key.

⚠️ Where Mamba implementations go wrong

Silent changes to the layer. Applying the softplus to $\Delta$ twice, or adding the bias $\mathbf{b}_\Delta$ in both the projection and the scan, gives a different layer that still trains. The bug survives until you load someone else's checkpoint. So does using the exact ZOH input weight in code meant for the kernels' $\Delta\mathbf{B}$ , or the reverse.
Reading too early. Reading the state before writing the current token shifts every output by one step.
The convolution window. In decoding, the window must shift every step and hold the inputs to the convolution, not its outputs.
Left padding. In batched prefill, a padding token is not masked out of the state the way attention masks it. Zeroing the pad embedding is not enough, because the convolution bias makes $\mathbf{x}_t = \operatorname{SiLU}(\mathbf{b}_{\text{conv}}) \ne \mathbf{0}$ and the state absorbs it. Zero the SSM input at padded positions too.
Precision. Keep $\boldsymbol\ell$ , $\mathbf{b}_\Delta$ , the decays, and the state in fp32, for the reasons of §21.1. A bf16 decay is either at most $1 - 2^{-8}$ or exactly 1, so every half-life between about 177 steps and forever is lost.

Mamba at real scale

Gu and Dao trained Mamba language models from 130M to 2.8B parameters on 300B tokens of the Pile, an 825 GiB collection of diverse English text (Gao et al. 2020). They used the same tokenizer, data, and number of tokens as the Pythia suite of Transformers, which makes the comparison unusually clean.

📌 Mamba at scale (Gu & Dao, December 2023)

In scaling experiments from about 125M to 1.3B parameters on the Pile, Mamba was the first attention-free model to match a strong modern Transformer recipe (rotary embeddings, SwiGLU, RMSNorm), and its advantage grew with sequence length. On zero-shot evaluations, Mamba-2.8B averaged 63.3 over six tasks (LAMBADA, HellaSwag, PIQA, ARC-Easy, ARC-Challenge, WinoGrande). Pythia-2.8B averaged 59.1, Pythia-6.9B 61.7, and GPT-J-6B 63.0. Mamba-2.8B's Pile validation perplexity was 6.22, against 6.51 for Pythia-6.9B.

On the synthetic induction-heads task, a two-layer Mamba trained at length 256 stayed perfect up to length $2^{20}$ , 4,000 times longer, while no attention variant generalized beyond twice its training length. That result is for one synthetic task with a single item to find. It does not carry over to language modeling, where Mamba models degrade past their training length (the last limit in the next part).

The paper's ablations, on models of about 350M parameters, show where the gains come from. They are worth knowing in detail, because interviewers ask which part of selection matters. Each row makes a different set of parameters depend on the input.

selective parameters	perplexity
none (a time-invariant layer)	10.93
$\mathbf{B}$ only	10.15
$\mathbf{C}$ only	9.98
$\Delta$ only	9.81
$\Delta$ , $\mathbf{B}$ , and $\mathbf{C}$	8.71

The step size alone captures half the gain, as the gate theorem suggests, and the three together do much better than any one. The state size tells the same story from the other side. With $\mathbf{B}$ and $\mathbf{C}$ fixed, growing $N$ from 1 to 16 moved perplexity only from 9.88 to 9.81. With $\mathbf{B}$ and $\mathbf{C}$ selective, the same growth moved it from 9.73 to 8.71, more than a full point, for about 1% more parameters. A larger state helps only if the layer can choose what goes into it.

Mamba layers soon appeared elsewhere. Jamba (Lieber et al. 2024) interleaves Mamba-1 layers with attention and mixture-of-experts layers. Its authors found that pure Mamba models struggled to follow the output format of few-shot prompts, and that a hybrid with one attention layer in eight fixed it. MambaByte (Wang et al. 2024) models raw bytes with no tokenizer (§2.7). Vision Mamba (Zhu et al. 2024) scans image patches in both directions, a bidirectional use of the layer that competes with the vision Transformer on long sequences of patches and has been applied to video models too.

	S4 and S4D (§21.3)	Mamba (S6)
parameters $\Delta$ , $\mathbf{B}$ , $\mathbf{C}$	fixed after training	computed from each token
mixing matrix	Toeplitz, a convolution	depends on $(t, s)$ , not a convolution
training computation	FFT convolution, $O(T\log T)$	fused scan, $O(TN)$ per channel
state matrix	complex diagonal (diagonal plus low rank in S4)	real diagonal
state per channel	$N = 64$ is typical, cheap through the kernel	$N = 16$ , fully materialized on chip
decoding	recurrence, constant memory	recurrence, constant memory, plus a 3-token conv cache
strengths	continuous signals, very long fixed-range dependencies, resampling	discrete data (text, DNA), content-based filtering, resets
weaknesses	cannot select by content	no tensor cores, real eigenvalues, recall limited by the state

Where Mamba breaks

Mamba has four limits, and each one drives a later part of the chapter.

The scan does not use tensor cores. Every operation inside it is entrywise, so its arithmetic runs at the 67 TFLOP/s of ordinary fp32 units, not the 989 TFLOP/s of the tensor cores, the GPU's bf16 matrix-multiply units (§21.1). The arithmetic grows with $N$ . Each state entry also needs its own exponential, $BTd_{\text{in}}N$ of them, on special-function units 17 times slower still, and at $N = 16$ these already set the fused kernel's time. Together these costs are why Mamba kept $N = 16$ .

Mamba-2 fixes this by giving each head, a group of $P$ channels, a single scalar decay. That turns most of the work into matrix products and allows $N = 64$ to 256 (§21.6). §21.5 develops the chunked algorithm that makes it possible.

Positive real decays cannot count. Every decay $\bar a_{t,cn} = e^{\Delta\lambda}$ lies in $(0, 1]$ , so the state can shrink but never flip sign or rotate. Tracking the parity of a bit string needs exactly a flip. §21.6 states the theorem that finite-precision linear recurrences with only positive eigenvalues cannot compute parity, and shows the negative and complex eigenvalues that fix it.

A fixed state cannot recall everything. Mamba-2.8B's state is 6.2 million numbers, which is large, but the update rule decides what to keep before it knows what will be asked. Waleffe et al. (2024) compared 8B-parameter models trained on the same 3.5T tokens. Pure Mamba and Mamba-2 models matched or beat a Transformer on many tasks. They lagged on tasks that need strong copying or in-context learning, such as 5-shot MMLU and phonebook lookup (given a long list of names and numbers, return one person's number). A hybrid with 7% attention layers beat the Transformer on all 12 standard tasks. §21.6 states the recall and copying theorems and designs such hybrids.

A state tuned to the training length. A recurrent layer can read sequences of any length, but it has been trained on one length, and its decays are tuned to it. Mamba and Mamba-2 language models trained at 2K to 8K tokens degrade beyond that length, sometimes sharply. Two mechanisms have been identified, and they are opposite failures. Some layers forget too fast, so the start of a long input fades out of reach. Other layers forget too little, so their state keeps growing past anything seen in training. In both cases passkey retrieval, finding a number hidden in a long filler text, fails past the training length.

🔬 The two length failures, measured

Forgetting too fast. The weight from token $s$ to output $t$ carries the factor $\exp(\lambda_{cn}\sum_{r=s+1}^{t}\Delta_{r,c})$ , which falls exponentially in the total step size between them. Ben-Kish et al. (2024) studied Mamba models trained at 2K tokens. The layers that carry long-range information learn step sizes whose sum over 2K tokens still leaves the start visible. With step sizes like those seen in training, the sum over 16K tokens is about 8 times larger, and the start's weights vanish. So the effective receptive field is set by the training length, and passkey retrieval fails past it.

Forgetting too little. Chen et al. (2024) studied the official Mamba-2 checkpoints, trained at 8K. Some heads learned to keep nearly everything. In one layer, three of the first eight heads still gave the first token a weight above 0.997 after 8K steps. The reason is that 8K tokens never fill a state that large, so forgetting was never needed. Past the training length a few channels of those states grow far beyond anything seen in training, and the state's mean and variance explode. Passkey retrieval that was near perfect within 8K falls to poor or zero beyond 16K, a failure the paper's first version called state collapse. The shortest training length at which models learned to forget grew linearly with the state size.

The fixes follow from the mechanisms. Train, or continue training, at longer lengths. With continued pre-training on longer sequences, Chen et al. brought a Mamba-2 370M model to near-perfect passkey retrieval at 256K tokens. At inference, force more forgetting or cap the state's norm. Or, as Ben-Kish et al.'s DeciMamba does, keep only the tokens with the largest $\Delta_t$ in the long-range layers during prefill, so that the total step size stays near the trained range. Better still, use layers whose update rules control forgetting more directly, the gated and delta-rule recurrences of §21.6. In practice, measure any recurrent model past its training length before trusting it there.

In PyTorch

The block below is the whole Mamba mixer in PyTorch, with a plain loop for the scan, which is also how the reference repository's selective_scan_ref computes it. It checks three things any implementation should pass. The output is causal, the parameter count equals the formula, and step-by-step decoding with the two caches reproduces the full forward pass. In production you would call the fused kernels instead, selective_scan_fn from the mamba_ssm package and causal_conv1d_fn from the causal-conv1d package, or the Hugging Face MambaForCausalLM, which uses them when they are installed.

🔥 In PyTorchfor a machine with PyTorch

🧭 What to remember

No layer that is linear in its input, time-invariant or not, can recall the last data token, because a token's contribution cannot depend on whether a newer one follows it. Selection means making the mixing weights depend on the content.
The selective SSM (S6) computes from the input a step size $\Delta_t = \operatorname{softplus}(\cdot)$ per token and channel, and vectors $\mathbf{B}_t, \mathbf{C}_t \in \R^N$ per token. It updates $\mathbf{H}_t = e^{\Delta_t\lambda}\odot\mathbf{H}_{t-1} + (\Delta_t\odot\mathbf{x}_t)\mathbf{B}_t^\top$ and reads $\mathbf{y}_t = \mathbf{H}_t\mathbf{C}_t + \mathbf{D}\odot\mathbf{x}_t$ . The kernels use the exponential-Euler write $\Delta\mathbf{B}$ , not exact ZOH.
With one state, ZOH, and $\Delta = \operatorname{softplus}(z)$ , selection is exactly the gate $h_t = (1 - \sigma(z_t))h_{t-1} + \sigma(z_t)x_t$ . A small $\Delta$ ignores a token, and a large $\Delta$ flushes the state and writes it.
Input-dependent parameters make the mixing matrix depend on $(t, s)$ , so the FFT convolution is gone. The layer is still linear in its state, so it runs as a recurrence or a parallel scan.
Storing the $(B, T, d_{\text{in}}, N)$ states costs about $1 + 4N$ times the necessary HBM traffic, so the scan is fused in on-chip memory and the backward pass recomputes the states. The fused scan is 20 to 40 times faster than a PyTorch scan, and its floor is set by $BTd_{\text{in}}N$ exponentials on the slow special-function units.
A Mamba block has about $6d^2$ parameters, mostly in its projections, so two replace one Transformer layer. Mamba-2.8B has exactly 2,768,345,600 parameters and decodes with 6.2 million cached numbers at any context length, which buys 4 to 5 times higher throughput.
Mamba's limits are elementwise work that limits $N$ , positive real decays that cannot count, a fixed state that limits recall, and decays tuned to the training length. Next, §21.5 develops the scans and the chunked algorithm behind fast recurrent kernels, which §21.6 then uses to put Mamba on tensor cores.

Gu & Dao (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. The paper. Section 3 for selection and Theorem 1 (proof in Appendix C), Appendix D for the fused scan and recomputation, Section 4.6 for the ablations quoted here.
state-spaces/mamba. Reference implementation. selective_scan_ref is the readable definition of the kernel, including its $\Delta\mathbf{B}$ discretization, and mamba_simple.py holds the block, its initialization, and the decoding step with both caches.
Rush. Mamba: The Hard Way. Builds the fused selective scan step by step in Triton, the best way to see what lives in on-chip memory.
Lahoti et al. (2026). Mamba-3: Improved Sequence Modeling using State Space Principles. Section 3 derives the kernels' discretization as exponential-Euler and generalizes it.
Ali, Zimerman & Wolf (2024). The Hidden Attention of Mamba Models. Mamba's mixing matrix read as data-dependent attention, with explainability methods built on it.
Waleffe et al. (2024). An Empirical Study of Mamba-based Language Models. The controlled 8B comparison of Mamba, Mamba-2, Transformers, and hybrids on the same 3.5T tokens.
Ben-Kish et al. (2024). DeciMamba: Exploring the Length Extrapolation Potential of Mamba. Mamba's limited effective receptive field beyond the training length, measured through its hidden attention, and a training-free fix that drops low- $\Delta$ tokens.
Chen et al. (2024). Stuffed Mamba: Oversized States Lead to the Inability to Forget. State collapse in Mamba-2 past its training length, and the training length needed to learn forgetting as a function of state size. The first version, titled "State Collapse and State Capacity of RNN-Based Long-Context Modeling", also tests training-free fixes such as state normalization.
Lieber et al. (2024). Jamba: A Hybrid Transformer-Mamba Language Model. The first large hybrid built on Mamba-1 layers, with the in-context learning ablation and the RMSNorm needed inside Mamba layers at scale.

Made for larger screens

Mamba and selective state spaces

Why a time-invariant layer cannot select

Letting each token choose its step size

The selective state-space model

The discretization the code actually uses

Selection is a gate

Selection is a gate, numerically

What each parameter does

Selection breaks the convolution

The selective scan from scratch

The selective scan, from scratch

Why a fixed kernel cannot select

Why the scan must be fused

The Mamba block

Count Mamba's parameters and its decode state

Decoding with a constant-size cache

One decoding step with conv and state caches

Mamba at real scale

Where Mamba breaks

In PyTorch

Read more

Watch next