Mamba & State Space Models

The previous chapter on Mixture of Experts cut feed-forward work by routing each token through a few experts. In its Mixtral example, the sequence mixer was still full self-attention. During decode, processing another position extends those layers' key-value (KV) cache. The hybrid examples already hinted at a separate choice: change how a layer carries history.

A service receives a 32K-token API-documentation prompt. Full causal attention retains a key and value for every position, so a later query can attend directly to an earlier identifier. That access doesn't guarantee a correct answer. A state-space model (SSM) instead carries a fixed-size running summary forward. Memory becomes easier to budget, but exact lookup can get harder because a summary replaces addressable history.^{[1]Reference 1Mamba: Linear-Time Sequence Modeling with Selective State Spaceshttps://arxiv.org/abs/2312.00752}

Mamba makes that summary selective: the current input changes how the model writes, retains, and reads state. Mamba-1 keeps recurrent state and a short convolution buffer instead of a full token cache; later variants change the auxiliary buffers. Start with four scalar inputs to see what gets retained, then use the same recurrence to understand parallel execution and its limits.

A discrete recurrence you can compute by hand

The fixed state isn't a hidden database. It's an activation updated once per token. At each step, old state is carried forward, the current input writes into it, and the output reads the result:

$h_t = \bar{A}\, h_{t-1} + \bar{B}\, x_t$ $y_t = \bar{C}\, h_t + \bar{D}\, x_t$

Read the first term as memory and the second as a write. The readout turns the updated state into the layer output. During decode, state size is fixed by the model rather than by preceding context length.

Walkthrough with numbers

A one-dimensional state can track a running activation as four tokens arrive. Set:

• $\bar{A} = 0.9$ (the state keeps 90% of its previous value)
• $\bar{B} = 0.2$ (new input contributes 20%)
• $\bar{C} = 1.0$ (read the full state)
• $\bar{D} = 0$ (no skip)

Now process four token activations $x = [3, 1, 4, 2]$:

Before calculating, predict the direction. Because $\bar{A}=0.9$, old state should fade slowly. Because $\bar{B}=0.2$, input 4 should pull the state up more than input 1. The result should smooth the sequence, not copy it.

The state smooths the input sequence. Input 4 pulls it up, but the earlier 0.74 still contributes through the 0.9 carryover.

That compression is useful and lossy at the same time. The layer stores one summary instead of four addressable token values, so its memory stays bounded while exact recovery becomes a learned capability rather than a direct lookup.

inputs = [3, 1, 4, 2]
a_bar = 0.9
b_bar = 0.2
state = 0.0
states = []

for value in inputs:
    state = a_bar * state + b_bar * value
    states.append(round(state, 3))

print("states:", states)
print("final state:", states[-1])

states: [0.6, 0.74, 1.466, 1.719]
final state: 1.719

Now change only $\bar{A}$. If it were 0.1 instead of 0.9, the final state would be $0.4826$ with the same inputs and $\bar{B}=0.2$, rather than $1.7194$. The code rounds displayed values, not the state passed between steps.

Earlier events now contribute far less. Whether that helps or hurts depends on what the task needs to retain, which is why state dynamics matter as much as state size.

Where $\bar{A}$ and $\bar{B}$ come from

The hand calculation hid one design choice: where do $\bar{A}$ and $\bar{B}$ come from? State Space Models borrow a control-theory view of dynamical systems, where a continuous signal changes over time. The continuous variables are:

• $x(t) \in \mathbb{R}$: the input signal
• $h(t) \in \mathbb{R}^s$: the hidden state, a vector of dimension $s$ that acts as a compressed summary of all past inputs
• $y(t) \in \mathbb{R}$: the output
• $\mathbf{A} \in \mathbb{R}^{s \times s}$: the state transition matrix, shaping how much past information to retain
• $\mathbf{B} \in \mathbb{R}^{s \times 1}$: the input matrix, shaping how much new information to absorb
• $\mathbf{C} \in \mathbb{R}^{1 \times s}$: the output matrix, which reads out the relevant part of the state

The scalar input and output keep notation readable. Real neural layers run many channels in parallel, but each channel follows the same kind of state update.

From these, the continuous dynamics are:

$h'(t) = \mathbf{A}h(t) + \mathbf{B}x(t)$ $y(t) = \mathbf{C}h(t) + \mathbf{D}x(t)$

The derivative equation says how state changes between observations. $\mathbf{A}$ controls carryover, while $\mathbf{B}$ determines how the input enters. The output $y(t)$ reads the state through $\mathbf{C}$, and $\mathbf{D}$ provides a direct input-to-output path.

This is the continuous version of the table above. We still need one step that turns smooth time into token steps.

Discretization

Tokens arrive as samples, but the equations above describe smooth change. To connect them, sample every $\Delta$ time units and assume the current input stays constant across each interval. This assumption is zero-order hold (ZOH). It converts continuous dynamics into a discrete recurrence. It's one discretization choice, not a claim that every S4 or Mamba implementation uses this exact input term:^{[1]Reference 1Mamba: Linear-Time Sequence Modeling with Selective State Spaceshttps://arxiv.org/abs/2312.00752}

Here $e^{\Delta \mathbf{A}}$ is the matrix exponential. The displayed inverse expression requires invertible $\mathbf{A}$ and $\Delta>0$. The general input term is $\bar{\mathbf{B}}=\int_0^\Delta e^{u\mathbf{A}}\mathbf{B}\,du$, which also works when $\mathbf{A}$ is singular; at $\Delta=0$ it is zero. These matrices give an exact token-step update for the stated held-input dynamics:

$h_t = \bar{\mathbf{A}} h_{t-1} + \bar{\mathbf{B}} x_t$ $y_t = \mathbf{C} h_t + \mathbf{D} x_t$

After discretization, each token step does exactly what the hand calculation did: carry state through $\bar{\mathbf{A}}$, absorb input through $\bar{\mathbf{B}}$, and read out through $\mathbf{C}$. The $\mathbf{D}$ term still skips input directly to output. Continuous notation explains where the coefficients come from; it isn't a second model.

For scalar $A=-1$, $B=1$, and $\Delta=1$, exact ZOH gives $\bar A\approx0.368$ and $\bar B\approx0.632$. If $A=0$, the input simply accumulates: $\bar A=1$ and $\bar B=\Delta B$. This helper handles that limit and uses expm1(z) to avoid subtracting nearly equal numbers when $z$ is small. It rejects nonfinite or unrepresentable computations rather than quietly returning infinite coefficients. Checking finite inputs alone wouldn't catch an overflowing exponential or input write.

from math import exp, expm1, isfinite

def finite_real(value) -> float:
    if type(value) not in (int, float):
        raise ValueError("a real coefficient is required")
    try:
        result = float(value)
    except OverflowError as error:
        raise ValueError("coefficient doesn't fit a finite float") from error
    if not isfinite(result):
        raise ValueError("a finite coefficient is required")
    return result

def zoh(a: float, b: float, delta: float) -> tuple[float, float]:
    a, b, delta = map(finite_real, (a, b, delta))
    if delta < 0:
        raise ValueError("nonnegative delta required")
    step = finite_real(delta * a)
    try:
        a_bar = exp(step)
        b_bar = delta * b if step == 0 else expm1(step) / a * b
    except OverflowError as error:
        raise ValueError("discretization exceeds float range") from error
    finite_real(b_bar)
    return a_bar, b_bar

assert zoh(-1, 1, 0) == (1.0, 0.0)
assert zoh(-1e-320, 1, 1e-320)[1] == 1e-320
for args in ((-1, 1, True), (-1, 1, 10**1000), (1000, 1, 1),
             (0, 1e308, 1e308), (-1, float("nan"), 1), (-1, 1, -1)):
    try:
        zoh(*args)
    except ValueError:
        pass
    else:
        raise AssertionError("invalid discretization accepted")

for a in (-1.0, 0.0, -1e-12):
    print(f"A={a:g}: A_bar, B_bar =", tuple(round(v, 6) for v in zoh(a, 1, 1)))

A=-1: A_bar, B_bar = (0.367879, 0.632121)
A=0: A_bar, B_bar = (1.0, 1.0)
A=-1e-12: A_bar, B_bar = (1.0, 1.0)

With that bridge in place, we can ask how to choose dynamics that retain useful long-range signals and still run efficiently.

S4: structured state spaces

The Structured State Space for Sequence Modeling (S4)^{[2]Reference 2Efficiently Modeling Long Sequences with Structured State Spaceshttps://arxiv.org/abs/2111.00396} makes that bounded-state idea useful over long horizons. Its structured dynamics retain signals and expose two ways to execute the same recurrence. The original S4 paper uses bilinear discretization, not the exact ZOH rule just demonstrated. Its diagonal-plus-low-rank parameterization makes recurrent updates and kernel construction efficient; a generic dense transition wouldn't inherit those costs. The step-by-step form is the generation view; training can use a parallel view.

HiPPO initialization

In S4, the continuous-time $\mathbf{A}$ matrix starts from the HiPPO (High-order Polynomial Projection Operators) framework, which compresses input history into coefficients of orthogonal polynomials.^{[3]Reference 3HiPPO: Recurrent Memory with Optimal Polynomial Projections.https://arxiv.org/abs/2008.07669[2]Reference 2Efficiently Modeling Long Sequences with Structured State Spaceshttps://arxiv.org/abs/2111.00396}

HiPPO turns history into coefficients over orthogonal-polynomial bases, giving state different coordinates for different temporal patterns. Without useful structure, a long recurrence can wash out old signals. S4 starts from those dynamics instead of asking an arbitrary recurrent matrix to discover long-horizon retention from scratch.

Dual computation modes

S4 can be computed in two equivalent ways. At training time, the whole sequence is known; at generation time, only the next token is available. For notational simplicity, the costs below are written for one channel; a full layer runs many channels in parallel.

Recurrent mode (generation)

Efficient for generation: $h_t = \bar{\mathbf{A}} h_{t-1} + \bar{\mathbf{B}} x_t, \quad y_t = \mathbf{C} h_t + \mathbf{D} x_t$

With an efficient structured transition, processing $L$ positions takes $O(Ls)$ work and retains $O(s)$ recurrent state for one channel. A generic dense $s\times s$ transition would instead cost $O(s^2)$ per step; structure is doing real work here.

Convolutional mode (training)

Efficient for training: $y = x * \bar{K} + \mathbf{D}x, \quad \text{where } \bar{K}_j = \mathbf{C} \bar{\mathbf{A}}^j \bar{\mathbf{B}},\quad j\geq0$

In plain terms

For zero initial state and fixed coefficients, unrolling the recurrence gives one causal convolution. The kernel is fixed across positions within a forward pass, but changes as its learned parameters train. An FFT (Fast Fourier Transform) can apply the prepared kernel in $O(L\log L)$ work; kernel construction has its own cost.

For our scalar example the first four kernel entries are $[0.2,0.18,0.162,0.1458]$. The fourth output is $0.2(2)+0.18(4)+0.162(1)+0.1458(3)=1.7194$. Verify every position, not just the final one:

from math import isclose

inputs = [3, 1, 4, 2]
kernel = [0.2 * 0.9**lag for lag in range(len(inputs))]
convolved = [sum(kernel[t - j] * inputs[j] for j in range(t + 1))
             for t in range(len(inputs))]
state = 0.0
recurrent = []
for value in inputs:
    state = 0.9 * state + 0.2 * value
    recurrent.append(state)
assert all(isclose(a, b, abs_tol=1e-12) for a, b in zip(convolved, recurrent))
print("matching outputs:", [round(v, 4) for v in convolved])

matching outputs: [0.6, 0.74, 1.466, 1.7194]

This direct convolution is intentionally quadratic Python code, not an FFT implementation or an S4 benchmark. A nonzero initial state adds the separate contribution $\mathbf C\bar{\mathbf A}^{t}h_0$ at step $t\geq1$.

The convolution works because the response to an input depends on its lag, not its content. What if an identifier needs a different update from a filler token?

The LTI problem

Inside an LTI SSM, every position uses the same transition and input/output maps. Its input values can differ, and surrounding neural layers can be nonlinear, but the SSM coefficients themselves don't change with content.

That property is called Linear Time-Invariant (LTI). Matrices $\mathbf{A}$, $\mathbf{B}$, and $\mathbf{C}$ stay fixed regardless of the input. The recurrence can still decay with position, but it can't change its write or read policy when a particular token matters.^{[1]Reference 1Mamba: Linear-Time Sequence Modeling with Selective State Spaceshttps://arxiv.org/abs/2312.00752}

Attention takes a different route: query-key interactions provide content-dependent access. Mamba's paper uses selective-copy tasks to expose the limitations of fixed LTI dynamics, then makes part of the update input-dependent.^{[1]Reference 1Mamba: Linear-Time Sequence Modeling with Selective State Spaceshttps://arxiv.org/abs/2312.00752}

Mamba: selective state spaces

Mamba^{[1]Reference 1Mamba: Linear-Time Sequence Modeling with Selective State Spaceshttps://arxiv.org/abs/2312.00752} solves the LTI limitation with one core idea: make part of the SSM input-dependent.

The selective mechanism

A documentation-QA decoder moving through a long API spec makes the limitation concrete. A filler token such as the may need a small state update, while a rare identifier such as max_retries may need a stronger write. Predict what a useful memory system should do: preserve the identifier, then spend less state capacity on filler.

Mamba makes that prediction part of the update. Instead of fixed $\mathbf{B}$, $\mathbf{C}$, and $\Delta$, it computes them as functions of the input:

$\mathbf{B}_t = \text{Linear}_B(x_t), \quad \mathbf{C}_t = \text{Linear}_C(x_t), \quad \Delta_t = \text{softplus}(\text{Linear}_\Delta(x_t))$

Here $x_t$ means the activation entering the SSM after projection, causal convolution, and SiLU, not a raw token ID. Unlike S4's fixed maps, Mamba computes these parameters from that activation. The learned diagonal transition $\mathbf{A}$ and skip term $\mathbf{D}$ remain input-independent in Mamba-1. Selectivity enters through $\mathbf{B}_t$, $\mathbf{C}_t$, and $\Delta_t$.^{[1]Reference 1Mamba: Linear-Time Sequence Modeling with Selective State Spaceshttps://arxiv.org/abs/2312.00752}

The step size $\Delta_t$ controls how much of that change is applied. With negative diagonal entries in $\mathbf{A}$, it determines how quickly each state coordinate decays. Holding the input and $B_t$ fixed in Mamba-1's implemented update:

• Large $\Delta_t$: Decay older state faster and increase the magnitude of the input-write term $\Delta_t B_t x_t$
• Small $\Delta_t$: Keep more of the previous state and reduce that write term

This is a statement about two terms, not a monotonic guarantee for the final state or a measure of token importance. The write can be negative or zero, and the learned $B_t$ and $C_t$ change too. Its coefficient isn't bounded by one as a sigmoid gate would be.

from math import exp

a = -1.0
for delta in (0.1, 1.0, 2.0):
    retained_fraction = exp(delta * a)
    print(f"delta={delta:.1f} retains {retained_fraction:.3f} of previous state")

delta=0.1 retains 0.905 of previous state
delta=1.0 retains 0.368 of previous state
delta=2.0 retains 0.135 of previous state

This gives Mamba a learned, content-dependent recurrent update while retaining linear-length recurrent computation. It still doesn't give every output explicit access to every prior token as attention does.

Avoid a common inference: no KV cache doesn't mean Mamba remembers everything. A large $\Delta_t$ can make old state decay quickly, and a poorly learned gate can wash out max_retries when the arrives. Mamba learns how much to forget per token, not whether to open a full history for lookup.

Architecture

Selective recurrence is the core, not the entire block. Mamba-1 uses projection, causal convolution, a selective scan over input-dependent updates, and a separate output gate.^{[1]Reference 1Mamba: Linear-Time Sequence Modeling with Selective State Spaceshttps://arxiv.org/abs/2312.00752} SiLU means $z\,\mathrm{sigmoid}(z)$; unlike a sigmoid-only gate, it isn't restricted to $[0,1]$. Softplus, $\log(1+e^z)$, makes the step size positive in real arithmetic. Stable implementations avoid evaluating $e^z$ directly for large positive $z$.

The example below follows Mamba-1's scalar input update, $h_t=e^{\Delta_t A}h_{t-1}+\Delta_t B_t x_t$. The official implementation uses $\bar B_t=\Delta_t B_t$, not the exact ZOH input integral above.^{[4]Reference 4Mamba selective scan implementationhttps://github.com/state-spaces/mamba/blob/e9594ce1c732d97440f0332fdc43170a2294dbfa/mamba_ssm/ops/selective_scan_interface.py} These two discretizations agree approximately for small steps, but aren't interchangeable for finite steps. For $A=-1$, $B=1$, $\Delta=1$, their input coefficients are 1 and about 0.632 respectively.

Use the same four values as SSM inputs. We choose simple, untrained functions for $B_t$, $C_t$, and $\Delta_t$ so every term is inspectable. This leaves out convolution, projection, skip, and output gating. Input 4 should have the smallest carryover factor. The readout can differ from the stored state because $C_t$ is also selective.

from math import exp, log1p

def softplus(value: float) -> float:
    return max(value, 0.0) + log1p(exp(-abs(value)))

inputs = [3.0, 1.0, 4.0, 2.0]
a = -1.0
state = 0.0
deltas = []
retained = []
states = []
outputs = []
pairs = []

for token in inputs:
    write = 0.1 * token
    read = 1.0 + 0.05 * token
    delta = softplus(-1.0 + 0.3 * token)
    a_bar = exp(delta * a)
    additive = delta * write * token
    state = a_bar * state + additive
    pairs.append((a_bar, additive))
    deltas.append(round(delta, 3))
    retained.append(round(a_bar, 3))
    states.append(round(state, 3))
    outputs.append(round(read * state, 3))

print("delta:", deltas)
print("retained fraction:", retained)
print("states:", states)
print("readouts:", outputs)
print("final state:", states[-1])

delta: [0.644, 0.403, 0.798, 0.513]
retained fraction: [0.525, 0.668, 0.45, 0.599]
states: [0.58, 0.428, 1.47, 1.085]
readouts: [0.667, 0.449, 1.764, 1.194]
final state: 1.085

The recurrence code explains behavior, not language-model quality or serving capacity. Count stored elements separately. This hypothetical 32-layer comparison uses 32 KV heads of width 128 for attention, and 4,096 inner channels with state size 16 for Mamba-1. These aren't matched trained models. Both buffers use two bytes per value; actual runtimes may use different state dtypes.

For a convolution of width 4, three prior activations are mathematically sufficient before the next input arrives. The official Mamba-1 cache allocates four slots, including the current position during its rolling update, so the calculation below counts all four.^{[5]Reference 5Mamba-1 block implementationhttps://github.com/state-spaces/mamba/blob/e9594ce1c732d97440f0332fdc43170a2294dbfa/mamba_ssm/modules/mamba_simple.py}

def transformer_kv_bytes(layers: int, kv_heads: int, head_dim: int, tokens: int, bytes_per_value: int = 2) -> int:
    return 2 * layers * kv_heads * head_dim * tokens * bytes_per_value

def ssm_state_bytes(
    layers: int,
    inner_channels: int,
    state_size: int,
    conv_width: int,
    bytes_per_value: int = 2,
) -> int:
    recurrent_state = layers * inner_channels * state_size
    convolution_state = layers * inner_channels * conv_width
    return (recurrent_state + convolution_state) * bytes_per_value

layers = 32
heads = 32
head_dim = 128
inner_channels = 4096
state_size = 16
conv_width = 4
attention_layers = 4

for tokens in [128, 512, 2048, 8192, 65536]:
    kv_mb = transformer_kv_bytes(layers, heads, head_dim, tokens) / 1024**2
    ssm_mb = ssm_state_bytes(layers, inner_channels, state_size, conv_width) / 1024**2
    hybrid_mb = kv_mb * attention_layers / layers + ssm_mb * (layers - attention_layers) / layers
    print(
        f"context={tokens:5d} -> transformer_kv={kv_mb:7.1f} MiB, "
        f"hybrid={hybrid_mb:6.1f} MiB, ssm_state={ssm_mb:4.1f} MiB"
    )

context=  128 -> transformer_kv=   64.0 MiB, hybrid=  12.4 MiB, ssm_state= 5.0 MiB
context=  512 -> transformer_kv=  256.0 MiB, hybrid=  36.4 MiB, ssm_state= 5.0 MiB
context= 2048 -> transformer_kv= 1024.0 MiB, hybrid= 132.4 MiB, ssm_state= 5.0 MiB
context= 8192 -> transformer_kv= 4096.0 MiB, hybrid= 516.4 MiB, ssm_state= 5.0 MiB
context=65536 -> transformer_kv=32768.0 MiB, hybrid=4100.4 MiB, ssm_state= 5.0 MiB

Three debugging questions before writing a kernel

With zero input and $A<0$, homogeneous state should decay. If it grows, inspect the transition and step-size parameterization. With nonzero input, negative $A$ alone doesn't bound writes, projections, or the final state; the finite-coefficient exercise above already exposed that distinction.

If two activations produce equal read/write coefficients, try more inputs and inspect the parameter maps. Equality on one pair can reflect a learned insensitive direction. Replacing the complete $B_t$, $C_t$, and $\Delta_t$ functions with constants makes this core update LTI.

If training is slow, profile the executed path. A Python loop, missing fused kernel, extra copies, and an inefficient shape need different repairs. The teaching loop isn't a serving benchmark. Production Mamba-1 uses optimized scan kernels, including CUDA implementations; changing the recurrence to silence a timing symptom would change the model.

The hardware-aware parallel scan

The recurrence can be trained with ordinary backpropagation through a loop, but that loop underuses a GPU. Mamba's hardware-aware scan changes the execution strategy. Its parameters depend on the known input sequence, not on the preceding hidden state, so they can be prepared across positions before the recurrence is evaluated.

The bottleneck isn't only arithmetic. It's memory traffic too. If every token update spills intermediate state back to high-bandwidth memory (HBM), the recurrence becomes bandwidth-bound. Mamba fuses discretization and state updates so recurrent state stays in fast on-chip static random-access memory (SRAM) as much as possible, then uses an associative scan to recover parallelism across positions.^{[1]Reference 1Mamba: Linear-Time Sequence Modeling with Selective State Spaceshttps://arxiv.org/abs/2312.00752}

The key reformulation treats each update as an affine transform, $h \mapsto a h + b$. Two such transforms can be composed without evaluating every intermediate state, and composition is associative. Like a parallel prefix sum, that lets chunks be combined in a tree instead of waiting on one long chain.

Order still matters. Applying $(a_1,b_1)$ and then $(a_2,b_2)$ gives $(a_2a_1,a_2b_1+b_2)$. Regrouping these compositions is valid; swapping them generally isn't.

One tree reduction gives only the final transform. Training needs a prefix scan: the composition ending at every position, so every output can be read. The following CPU implementation first builds subtree totals, then distributes each subtree's incoming transform. Padding with the identity $(1,0)$ handles non-power-of-two lengths.

from math import isclose

def compose(right: tuple[float, float], left: tuple[float, float]) -> tuple[float, float]:
    """Compose h -> a*h + b transforms: apply left, then right."""
    a_right, b_right = right
    a_left, b_left = left
    return a_right * a_left, a_right * b_left + b_right

def prefix_scan(steps, initial=0.0):
    if not steps:
        return []
    size = 1
    while size < len(steps):
        size *= 2
    tree = [(1.0, 0.0)] * (2 * size)
    tree[size:size + len(steps)] = steps
    for node in range(size - 1, 0, -1):
        tree[node] = compose(tree[2 * node + 1], tree[2 * node])
    result = [0.0] * len(steps)

    def visit(node, before):
        if node >= size:
            index = node - size
            if index < len(steps):
                a, b = compose(tree[node], before)
                result[index] = a * initial + b
            return
        visit(2 * node, before)
        visit(2 * node + 1, compose(tree[2 * node], before))

    visit(1, (1.0, 0.0))
    return result

steps = [(0.9, 0.6), (0.9, 0.2), (0.9, 0.8), (0.9, 0.4)]
expected = [0.6, 0.74, 1.466, 1.7194]
scanned = prefix_scan(steps)
assert all(isclose(a, b, abs_tol=1e-12) for a, b in zip(scanned, expected))
assert prefix_scan([]) == []
assert isclose(prefix_scan(steps[:3], initial=2)[-1], 0.9**3 * 2 + 1.466)
print("all prefix states:", [round(v, 4) for v in scanned])
print("swapping two updates changes the transform:", compose(steps[1], steps[0]) != compose(steps[0], steps[1]))

all prefix states: [0.6, 0.74, 1.466, 1.7194]
swapping two updates changes the transform: True

The tree has $O(L)$ nodes and $O(\log L)$ depth. Independent nodes at each level can run in parallel; this Python traversal doesn't do so. Mamba's optimized scan applies this algebra coordinate-wise to structured state updates and avoids storing all expanded intermediate states in HBM.^{[1]Reference 1Mamba: Linear-Time Sequence Modeling with Selective State Spaceshttps://arxiv.org/abs/2312.00752}

Associativity holds in exact arithmetic. Floating-point cancellation can make regrouping errors large relative to the answer. Reuse the scan with carryover equal to one, so the recurrence becomes a sum. Predict whether either traversal retains all three small writes between the large opposite writes.

from math import fsum

writes = [1e16, 1.0, -1e16, 1.0, 1.0]
serial = 0.0
for write in writes:
    serial += write
tree_result = prefix_scan([(1.0, write) for write in writes])[-1]
reference = fsum(writes)
assert serial == 2.0 and tree_result == 1.0 and reference == 3.0
print("sequential / tree / accurate sum:", serial, tree_result, reference)

sequential / tree / accurate sum: 2.0 1.0 3.0

fsum provides an accurate reference for this sum, not a replacement for a general selective scan. Both traversals lose information in this deliberately extreme fixture. Check representative value ranges and precision against a suitable reference before choosing tolerances; a loose tolerance can conceal a broken state update.

Transformers can reuse FlashAttention and a mature PyTorch stack. SSMs need their own fused kernels, serving paths, and state management, so linear asymptotics alone don't settle adoption. That remaining kernel gap is what Mamba-2 addresses.

Mamba-2: state space duality

Mamba showed that selective SSMs could work well, but its parallel scan still followed a custom hardware path. That path underused the matrix-multiplication engines (Tensor Cores) that make GPUs fast.

Mamba-2^{[6]Reference 6Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Dualityhttps://arxiv.org/abs/2405.21060} attacks that gap through Structured State Space Duality (SSD). Its SSM uses a scalar-times-identity transition within each head, a more restricted structure than Mamba-1's diagonal transition. That structure relates the recurrence to masked kernel attention, not to arbitrary softmax attention.

Unroll the recurrence and each output becomes a weighted sum of earlier inputs. The weights factor into a write at the earlier position, decay between positions, and a read at the current position. The resulting matrix is semiseparable: its off-diagonal blocks admit low-rank factorizations. SSD uses those factors without allocating a full $L\times L$ matrix. Its chunked algorithm:

• intra-chunk outputs are computed in parallel
• chunk final states are computed with batched matrix multiplications
• inter-chunk propagation handles a much shorter sequence of boundary states, using a scan or structured matrix multiplication
• the initial-state contribution for each chunk is added back with more batched matrix multiplications

In SSD, more work becomes matrix multiplication, while a reduced inter-chunk recurrent core remains. The model keeps the bounded recurrent state, but changes the route used to compute it.

Mamba-1 vs. Mamba-2 comparison

In its dedicated implementation, the Mamba-2 paper reports 2-8x faster core-layer execution than the optimized selective-scan implementation of Mamba-1; its efficiency figure uses state expansion 64.^{[6]Reference 6Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Dualityhttps://arxiv.org/abs/2405.21060} The paper separately evaluates language-model quality. This isn't a guaranteed end-to-end server speedup. Check hardware, kernels, batch shape, sequence lengths, and quality on your own stack.

chunk_size = 256
for sequence_length in (16384, 16385):
    chunks = (sequence_length + chunk_size - 1) // chunk_size
    final_chunk_tokens = sequence_length - (chunks - 1) * chunk_size
    print(f"tokens={sequence_length}: chunks={chunks}, final chunk={final_chunk_tokens}")

tokens=16384: chunks=64, final chunk=256
tokens=16385: chunks=65, final chunk=1

Beyond Mamba-2 toward Mamba-3

SSD changes how the recurrence runs. Mamba-3^{[7]Reference 7Mamba-3: Improved Sequence Modeling using State Space Principleshttps://arxiv.org/abs/2603.15569} asks a more specific serving question: once state must be read every token, how can the update carry more useful arithmetic without moving more state bytes? It takes an explicitly inference-first view and changes the recurrent update along three axes:

• Exponential-trapezoidal discretization, which generalizes the exponential-Euler rule used in Mamba-1 and Mamba-2. The stated second-order result requires $\lambda_t=1/2+O(\Delta_t)$ and regularity/stability assumptions; the paper's ablations prefer not enforcing that constraint.
• A complex-valued state update, implemented through data-dependent rotations, evaluated on state-tracking tasks such as parity and modular arithmetic where the paper's Mamba-2 baseline struggles. This isn't an impossibility claim about every real-valued recurrence.
• A multi-input, multi-output (MIMO) formulation that raises arithmetic intensity during memory-bound decode by doing more work in the state update without growing the state.

SISO means single-input, single-output; MIMO uses multiple write/read directions per head. In the paper's 1.5B-parameter experiments, SISO improved average downstream accuracy by 0.6 percentage points over Gated DeltaNet (GDN), and MIMO improved it by 1.8 points. In a separate 440M-model sweep, MIMO with state size 64 matched Mamba-2 perplexity at state size 128.^{[7]Reference 7Mamba-3: Improved Sequence Modeling using State Space Principleshttps://arxiv.org/abs/2603.15569}

The paper also reports up to 4× more decode FLOPs for MIMO at similar wall-clock decode latency. That is a hardware-utilization result, not a 4× tokens/s claim.^{[7]Reference 7Mamba-3: Improved Sequence Modeling using State Space Principleshttps://arxiv.org/abs/2603.15569} Its single-step kernel comparison uses batch 128 on one H100, FP32/BF16, and state sizes 64/128. The 1.5B quality comparison uses 100B FineWeb-Edu training tokens and 2K contexts; the state-size/perplexity sweep uses 440M models. Neither experiment certifies 32K API-document retrieval. Treat these as separate measurements, then re-measure hardware, precision, batch, context, and correctness in the target serving stack.

The buffers change too. The paper's Mamba-3 block removes the separate short causal convolution, while its trapezoidal recurrence needs the previous state-input write. A Mamba-1 convolution-cache formula doesn't describe every generation of the architecture.^{[7]Reference 7Mamba-3: Improved Sequence Modeling using State Space Principleshttps://arxiv.org/abs/2603.15569}

For one real-valued coordinate, the new update is

$h_t=\alpha_t h_{t-1}+(1-\lambda_t)\Delta_t\alpha_t B_{t-1}x_{t-1}+\lambda_t\Delta_t B_t x_t,\qquad \alpha_t=e^{\Delta_t A_t}.$

Setting $\lambda_t=1$ recovers the exponential-Euler input term. At $\lambda_t=1/2$, both interval endpoints contribute. Compare those approximations with exact ZOH for constant $A=-1$, $B=1$, and unit input at both endpoints. This scalar exercise omits complex rotations, learned selection, and MIMO; it explains the discretization rather than running Mamba-3.

from math import exp, expm1

for delta in (1.0, 0.1):
    alpha = exp(-delta)
    euler_write = delta
    trapezoidal_write = 0.5 * delta * (alpha + 1.0)
    exact_held_write = -expm1(-delta)
    print(f"delta={delta:.1f}: Euler={euler_write:.6f}, "
          f"trapezoidal={trapezoidal_write:.6f}, ZOH={exact_held_write:.6f}")

delta=1.0: Euler=1.000000, trapezoidal=0.683940, ZOH=0.632121
delta=0.1: Euler=0.100000, trapezoidal=0.095242, ZOH=0.095163

The trapezoidal value is closer here, but still isn't the exact held-input integral. Smaller numerical integration error on this signal doesn't prove better learned-model answers.

Complexity comparison

Separate training from decode to see the serving payoff. The chart uses the same hypothetical shapes and two-byte state values as the byte-count script. The table isolates sequence mixing in one layer, not MLP or projection costs. Here $L$ is sequence length, $d$ is channel width, and $s$ is the SSM state dimension per channel.

These are fixed-width asymptotics. S4's row excludes kernel construction; its fast recurrent update relies on structure. The attention row assumes full-context multi-head attention. GQA reduces KV width, and sliding-window attention changes the retained length. Kernelized linear attention maintains a matrix summary; its $d^2$ shorthand assumes both feature and value dimensions scale as $d$.

SSMs retain roughly $sd$ recurrent elements per layer, plus any short convolution buffer, regardless of processed length. This bounds decode state, not training activations, input storage, model weights, or total server memory. A hybrid with even one full-attention layer still has state growing with $L$.

One live state isn't the whole cache policy

Each concurrent request needs its own mutable state. A shared-prefix cache can retain additional snapshots so new requests resume from an earlier boundary. That saves recomputation but adds memory; snapshot count can grow with processed length even though one live recurrent state stays fixed. A new branch must get independent mutable state rather than overwrite another request's checkpoint.

Implementation snapshot, October 10, 2026: the pinned vLLM cache configuration distinguishes none, all, and align Mamba strategies, permits different convolution/SSM dtypes, and includes state-checkpoint and replay options with backend/decoding restrictions.^{[8]Reference 8vLLM Mamba cache configuration (pinned source)https://github.com/vllm-project/vllm/blob/4edb55169f69c76bb87ab9f154585917cb01bb17/vllm/config/cache.py} Those are implementation policies, not properties guaranteed by every Mamba checkpoint. Verify the exact engine, model, cache mode, and speculative method together; this lesson doesn't execute those serving paths.

Suppose the illustrative 5 MiB stack above retains a separate current buffer plus a complete snapshot every 1,024 tokens. No sharing, padding, eviction, or allocator overhead is included:

state_mib = 5
checkpoint_interval = 1024
for tokens in (8192, 65536):
    snapshots = tokens // checkpoint_interval
    total_mib = state_mib * (1 + snapshots)
    print(f"tokens={tokens}: live={state_mib} MiB, "
          f"snapshots={snapshots}, total={total_mib} MiB")
print("eight independent live requests:", 8 * state_mib, "MiB")

tokens=8192: live=5 MiB, snapshots=8, total=45 MiB
tokens=65536: live=5 MiB, snapshots=64, total=325 MiB
eight independent live requests: 40 MiB

This is a declared policy worksheet, not vLLM's physical allocation formula. An engine can share retained prefixes, evict snapshots, replay intervening tokens, or reserve extra state slots. Budget those choices explicitly.

That is the core contrast: explicit history versus fixed summary. Sometimes a specific token is more useful than the summary, so memory savings need a retrieval test beside them.

The original Mamba paper reports 4–5× inference throughput versus its similar-sized Transformer baseline, helped by fitting larger batches without a KV cache.^{[1]Reference 1Mamba: Linear-Time Sequence Modeling with Selective State Spaceshttps://arxiv.org/abs/2312.00752} Its benchmark uses an A100 80GB PCIe, 2,048-token prompts, 128-token generations, and batches from 1 to 128. It compares Mamba 1.4B and an untrained 6.9B configuration against standard Hugging Face GPT-3-style Transformers at 1.3B and 6.7B. This isn't a matched-quality comparison against every modern serving stack. Neither this byte-count script nor the recurrence checks reproduce it. A serving comparison needs the same hardware, precision, prompt/output distribution, quality target, and stated batching policy.

Hybrid architectures: combining trade-offs

A hybrid decoder sends every token through both kinds of layers in a fixed stack. It doesn't turn attention on only when it detects a difficult lookup. Most sequence-mixing layers may use recurrent state, while a smaller number preserve token-addressable KV histories.

Fixed-state SSM designs can regress on precise retrieval, while attention-only stacks incur KV-cache growth as contexts expand. Hybrid architectures are candidates when a workload needs both properties.

Hybrid models: Jamba and Nemotron-H

The original Jamba-v0.1^{[9]Reference 9Jamba: A Hybrid Transformer-Mamba Language Modelhttps://arxiv.org/abs/2403.19887} interleaves Transformer and Mamba layers in a structured hybrid stack and adds MoE to some multilayer perceptron (MLP) layers. Its released configuration uses one attention layer for every seven Mamba layers in each 8-layer block, with MoE applied every other layer.

The 2025 Nemotron-H recipe^{[10]Reference 10Nemotron-H: A Family of Accurate and Efficient Hybrid Mamba-Transformer Modelshttps://arxiv.org/abs/2504.03624} uses attention in roughly 8% of layers: 4 self-attention layers out of 52 in its 8B model. The remainder is an even split of Mamba-2 and FFN layers. Its 8B stack starts with Mamba-2, ends with FFN, and places an FFN immediately after each attention layer.

These layer counts use different conventions: Jamba counts mixer-plus-MLP layers, while Nemotron-H counts Mamba-2, attention, and FFN sublayers separately. Their fractions aren't directly comparable compute budgets. Attention placement is chosen during architecture design and training, not routed per question. MoE routing in Jamba is a separate mechanism inside selected MLPs.

A hybrid needn't use Mamba at all. Qwen3.6-35B-A3B mixes three Gated DeltaNet layers followed by one gated-attention layer, repeated ten times; MoE is its feed-forward choice.^{[11]Reference 11Qwen3.6-35B-A3Bhttps://huggingface.co/Qwen/Qwen3.6-35B-A3B} Delta-rule linear attention and selective SSMs are distinct recurrence families. Inspect the actual update and buffer shapes rather than infer them from an engine's broad “Mamba cache” category.

total_layers = 32
attention_layers = 4
full_attention_kv_gib = 32.0
hybrid_kv_gib = full_attention_kv_gib * attention_layers / total_layers

print("attention layer fraction:", f"{attention_layers / total_layers:.1%}")
print("illustrative full-attention KV:", f"{full_attention_kv_gib:.1f} GiB")
print("illustrative hybrid KV only:", f"{hybrid_kv_gib:.1f} GiB")

attention layer fraction: 12.5%
illustrative full-attention KV: 32.0 GiB
illustrative hybrid KV only: 4.0 GiB

Why hybrids work

Along the depth axis, every position passes through recurrent and attention layers. Along the sequence axis, each recurrent layer carries its own state while each attention layer retains its own KV cache:

When to use SSMs vs. transformers vs. hybrids

There's no universal winner. Start with the failure you need to avoid: exact copying, growing decode memory, or both. Then benchmark the candidate that addresses it.

Asymptotic complexity doesn't choose a model. Prefill still processes the prompt, and short-context inference may favor optimized attention kernels. Evaluate the actual checkpoints and serving implementations, not architecture names alone.

Exact-retrieval release gate (NIAH / multi-needle)

Don't ship a pure SSM long-context path on recall SLOs without the same style of needle-in-a-haystack (NIAH) and multi-needle / variable-tracking suite taught in Long Context Window Management. Attention keeps explicit token access. A fixed recurrent state compresses history, so middle-depth and multi-hop retrieval can regress even when average language-modeling loss looks fine.

Place identifiers near the beginning, middle, and end of prompts at each supported length. Add distractors, multiple facts, and updates that supersede an earlier value. Compare exact-answer accuracy under the same decoding settings, including unsupported questions that should be declined. Then measure peak device memory, time to first token, and per-token latency under the intended concurrency. More attention layers don't certify recall, and a single successful needle doesn't establish document-QA quality.

Build a counterexample to exact recovery

Can two different histories leave exactly the same final state? For $h_t=0.5h_{t-1}+0.5x_t$, compare [3, 1] and [1, 2]. Both end at 1.25. A readout that sees only this scalar state can't tell which first value occurred, even with exact arithmetic.

def summarize(values, initial=0.0):
    state = initial
    for value in values:
        state = 0.5 * state + 0.5 * value
    return state

left, right = [3, 1], [1, 2]
assert left != right and summarize(left) == summarize(right) == 1.25
suffix = [4, 2]
assert summarize(suffix, summarize(left)) == summarize(suffix, summarize(right))
print("different first values:", left[0], right[0])
print("same final state:", summarize(left))
print("same continuation:", summarize(suffix, summarize(left)))

different first values: 3 1
same final state: 1.25
same continuation: 2.3125

This proves a limitation of this scalar recurrence, not a measured failure rate for trained Mamba. More state dimensions and selective updates can preserve distinctions that it loses. For the documentation service, that is a reason to test identifier recovery, not a reason to assume every SSM fails it.

Resume a stream, then deliberately break it

Adapt prefix_scan to the selective update above. Five inputs exercise padding; splitting after two tests continuation. Predict whether resetting state at the split can pass the same comparison.

from math import exp, log1p, isclose

def selective_pairs(values):
    pairs = []
    for value in values:
        z = -1.0 + 0.3 * value
        delta = max(z, 0.0) + log1p(exp(-abs(z)))
        pairs.append((exp(-delta), delta * (0.1 * value) * value))
    return pairs

updates = selective_pairs([3.0, 1.0, 4.0, 2.0, 0.5])
sequential, state = [], 0.0
for a, b in updates:
    state = a * state + b
    sequential.append(state)
scanned = prefix_scan(updates)
first = prefix_scan(updates[:2])
resumed = first + prefix_scan(updates[2:], initial=first[-1])
reset = first + prefix_scan(updates[2:], initial=0.0)

def matches(left, right):
    return len(left) == len(right) and all(
        isclose(a, b, abs_tol=1e-12) for a, b in zip(left, right)
    )

assert matches(scanned, sequential) and matches(resumed, sequential)
assert not matches(reset, sequential)
print("padded scan matches sequential:", matches(scanned, sequential))
print("resumed stream matches:", matches(resumed, sequential))
print("reset stream matches:", matches(reset, sequential))

padded scan matches sequential: True
resumed stream matches: True
reset stream matches: False

This checks the scalar selective core. In a full Mamba-1 block, the convolution buffer must survive the split too, and batched state rows must follow their request IDs when the scheduler reorders requests.^{[5]Reference 5Mamba-1 block implementationhttps://github.com/state-spaces/mamba/blob/e9594ce1c732d97440f0332fdc43170a2294dbfa/mamba_ssm/modules/mamba_simple.py} A fresh independent request starts from clean state. Conversely, resuming an old request with clean state discards its history. Test both boundaries.

Make the buffer dependency visible with a width-three causal convolution followed by a scalar recurrence. Its weights run from oldest to newest. This fixture omits Mamba's nonlinear activation, selective parameter maps, and output gate. After the prefix [3, 1], compare restoring both buffers with restoring only the recurrent state.

from math import isclose

def run_stream(values, state=0.0, buffer=(0.0, 0.0, 0.0)):
    outputs = []
    for value in values:
        buffer = (*buffer[1:], value)
        filtered = sum(w * x for w, x in zip((0.2, 0.3, 0.5), buffer))
        state = 0.5 * state + filtered
        outputs.append(state)
    return outputs, state, buffer

full, _, _ = run_stream([3.0, 1.0, 4.0, 2.0])
prefix, saved_state, saved_buffer = run_stream([3.0, 1.0])
suffix, _, _ = run_stream([4.0, 2.0], saved_state, saved_buffer)
lost_buffer, _, _ = run_stream([4.0, 2.0], saved_state)
assert all(isclose(a, b) for a, b in zip(prefix + suffix, full))
assert not isclose(lost_buffer[0], full[2])
print("reference suffix:", [round(x, 4) for x in full[2:]])
print("restored both:", [round(x, 4) for x in suffix])
print("restored recurrence only:", [round(x, 4) for x in lost_buffer])

reference suffix: [3.975, 4.3875]
restored both: [3.975, 4.3875]
restored recurrence only: [3.075, 3.7375]

The first suffix input still depends on the preceding two inputs through the convolution. Restoring only saved_state changes the write before the recurrence even runs. A checkpoint must retain every live dependency of its block, and each resumed request must receive its own buffers.

Rejection requires a state boundary

In Speculative Decoding, the valid attention KV prefix can be kept while unaccepted entries are discarded. One overwritten recurrent state has no such token-indexed suffix. Keep a checkpoint, retained intermediate states, or a correct replay path.

Use the scalar recurrence from the collision exercise: the processed prefix is [3, 1], provisional inputs are [99, 77], only 99 is accepted, and the correction is 4. Predict what happens if the correction starts from state after rejected 77.

from math import exp

prefix_state = summarize([3, 1])
snapshots = [prefix_state]
for provisional in (99, 77):
    snapshots.append(summarize([provisional], snapshots[-1]))
accepted_count = 1
restored = snapshots[accepted_count]
corrected = summarize([4], restored)
contaminated = summarize([4], snapshots[-1])
reference = summarize([99, 4], prefix_state)
assert corrected == reference == 27.0625
assert contaminated == 33.78125 and contaminated != reference

# An exp(-1000) carryover underflows to zero in this float computation.
carryover = exp(-1000.0)
assert carryover == 0.0
assert carryover * 3.0 + 7.0 == carryover * 9.0 + 7.0
print("restored then corrected:", corrected)
print("corrected from rejected state:", contaminated)
print("distinct prior states become identical:", carryover == 0.0)

restored then corrected: 27.0625
corrected from rejected state: 33.78125
distinct prior states become identical: True

Sampling a correction supplies its ID; processing it computes the next state. Restore every affected layer's valid buffers before continuation, including convolution or previous-write buffers where applicable. Algebraic inversion is unsafe as a generic rollback plan: carryover can underflow to zero, and reversing tiny factors can amplify rounding. The checkpoint/replay design adds memory or work. This fixture demonstrates the dependency, not an engine's speculative support or measured speedup.

Finally, separate startup from continuation latency. In this explicitly hypothetical timing model, the prompt pass costs 120 ms for one candidate and 150 ms for another; subsequent decode steps cost 3.2 and 1.8 ms. The prompt pass already produces the first token, so $N$ generated tokens require $N-1$ decode steps. Predict which path wins at 16 and 256 tokens.

transformer = {"prefill_ms": 120, "decode_ms_per_token": 3.2}
ssm = {"prefill_ms": 150, "decode_ms_per_token": 1.8}

for generated_tokens in (16, 256):
    decode_steps = generated_tokens - 1
    transformer_total = transformer["prefill_ms"] + decode_steps * transformer["decode_ms_per_token"]
    ssm_total = ssm["prefill_ms"] + decode_steps * ssm["decode_ms_per_token"]
    winner = "ssm" if ssm_total < transformer_total else "transformer"
    print(f"generated={generated_tokens}: winner={winner}, delta_ms={abs(ssm_total - transformer_total):.1f}")

generated=16: winner=transformer, delta_ms=9.0
generated=256: winner=ssm, delta_ms=327.0

Where fixed-state models lead

The four-input calculation connects the main ideas: a recurrence retains a summary, selection changes its updates, and a prefix scan regroups those updates without changing their order. Hybrids retain explicit history in only part of the stack. None makes extra generated tokens free. Reasoning & Test-Time Compute asks when spending those tokens improves the answer enough to justify their cost.