2.11 Beyond Attention: SSMs, Mamba, RWKV & Linear Attention¶

The transformer has dominated language modeling since Vaswani et al. introduced it in 2017. Yet it carries a fundamental cost: the standard attention mechanism scales quadratically in both compute and memory with sequence length. For a 128 K-token context, a naive attention over a single head requires computing an \(N \times N\) score matrix with \(N = 131072\) — about 128 billion floating-point multiplications before the value aggregation step. At that scale, alternatives to attention stop being academic curiosities and start being engineering necessities.

This chapter is about those alternatives. We will study what makes attention expensive, then tour the major architectural families that attempt to fix the scaling problem: sparse and sliding-window attention, linear attention, structured state space models (SSMs) like S4 and Mamba, recurrent language models like RWKV, and the Retention mechanism behind RetNet. We end by looking at hybrid architectures that mix attention and these alternatives — a pragmatic approach that is finding real adoption in production LLMs.

Before diving in, note that the upstream machinery these models plug into — tokenization, embeddings, residual-stream design, layer normalization — is shared with standard transformers, covered in Tokenization: BPE, WordPiece, Unigram & Byte-Level, Embeddings & The Input Pipeline, and The Transformer Block: Norms, Residuals, MLPs & Activations. The KV-cache implications of inference are discussed in depth in The Anatomy of LLM Inference: Prefill, Decode & The KV Cache.

The Quadratic Bottleneck¶

Let us be precise about what “quadratic attention” actually means and where the cost comes from.

For a sequence of \(N\) tokens with model dimension \(d\) and head dimension \(d_k\), the attention computation is:

The \(QK^\top\) product produces an \(N \times N\) matrix. This single step costs \(O(N^2 d_k)\) FLOPs and requires \(O(N^2)\) memory to store the attention weights. For \(H\) heads:

• FLOPs per layer: \(O(H \cdot N^2 d_k) = O(N^2 d)\)
• Memory: \(O(N^2)\) attention matrix, which dominates for large \(N\)

At inference, the autoregressive KV cache grows as \(O(N \cdot d)\) per layer, which is linear — but the prefill step still reads \(O(N^2)\) attention weights. FlashAttention (covered in FlashAttention I: IO-Awareness & The Online Softmax) tames the memory cost by not materializing the full matrix, but does not reduce the FLOP count.

Sliding-Window and Sparse Attention¶

Before abandoning attention entirely, it is worth asking a cheaper question: does every query really need every key? Most linguistic dependencies are local, and the ones that are not tend to be sparse. Two families exploit this — sliding-window (local) attention and sparse attention — keeping the softmax (and its exact-retrieval sharpness) while cutting the \(O(N^2)\) cost to \(O(N \cdot W)\) or \(O(N \cdot k)\).

Sliding-Window (Local) Attention¶

A sliding-window layer restricts each query at position \(t\) to attend only to the last \(W\) keys, positions \([t - W + 1,\ t]\). The score matrix becomes a banded lower-triangular matrix of bandwidth \(W\) instead of a full triangle:

• FLOPs / memory per layer: \(O(N \cdot W \cdot d)\) instead of \(O(N^2 d)\) — linear in \(N\) for fixed \(W\).
• KV cache: bounded at \(W\) tokens (a rolling buffer), so decode memory is \(O(W)\) regardless of how long the sequence gets.

Mistral 7B uses \(W = 4096\); Gemma ⅔ and GPT-oss interleave sliding-window layers with occasional full-attention layers (below); Longformer’s local branch is the same banded mask.

import torch

def sliding_window_causal_mask(seq_len: int, window: int, device="cpu") -> torch.Tensor:
    """
    Additive attention mask for causal sliding-window attention.
    Query i may attend key j  iff  j <= i (causal)  AND  i - j < window (local).
    Returns (seq_len, seq_len) with 0 where allowed and -inf where masked;
    add it to QK^T / sqrt(d_k) before the softmax.
    """
    i = torch.arange(seq_len, device=device)[:, None]   # (N, 1)
    j = torch.arange(seq_len, device=device)[None, :]    # (1, N)
    allowed = (j <= i) & (i - j < window)                # banded lower-triangular
    mask = torch.zeros(seq_len, seq_len, device=device)
    mask.masked_fill_(~allowed, float("-inf"))
    return mask

The key trick for real speedups is block sparsity: skip any \(128 \times 128\) tile of the score matrix that is entirely masked. PyTorch’s FlexAttention does this for you from a mask_mod:

# PyTorch 2.5+: FlexAttention skips fully-masked 128x128 blocks for a real speedup.
from torch.nn.attention.flex_attention import flex_attention, create_block_mask

def sliding_window_mod(window: int):
    def mask_mod(b, h, q_idx, kv_idx):
        causal = kv_idx <= q_idx
        local  = q_idx - kv_idx < window
        return causal & local
    return mask_mod

# Build the block mask once, reuse across the forward pass:
block_mask = create_block_mask(
    sliding_window_mod(window=1024),
    B=None, H=None, Q_LEN=8192, KV_LEN=8192,
)
# out = flex_attention(q, k, v, block_mask=block_mask)   # q,k,v: (B, H, N, d_k)

Rolling-buffer KV cache. Because a query never looks back further than \(W\), the cache only needs to hold the last \(W\) keys/values. On each decode step the oldest entry is evicted:

class RollingKVCache:
    """
    Fixed-size KV cache for sliding-window attention. Holds at most `window`
    tokens, so decode memory is O(window) per layer regardless of sequence length.
    (Production kernels use an in-place ring buffer; this slice version is clearer.)
    """
    def __init__(self, window: int):
        self.W = window
        self.k = None   # (B, n_kv_heads, <=W, head_dim)
        self.v = None

    def update(self, k_t: torch.Tensor, v_t: torch.Tensor):
        # k_t, v_t: (B, n_kv_heads, 1, head_dim) -- RoPE already applied
        self.k = k_t if self.k is None else torch.cat([self.k, k_t], dim=2)[:, :, -self.W:]
        self.v = v_t if self.v is None else torch.cat([self.v, v_t], dim=2)[:, :, -self.W:]
        return self.k, self.v   # attend the current query over just these

Receptive field grows with depth. A single sliding-window layer sees \(W\) tokens, but stacking them compounds: layer 2 attends to positions that themselves summarized a window, so after \(L\) layers a token’s effective receptive field is \(\sim L \cdot (W - 1)\). Mistral’s 32 layers \(\times\) 4096 window reach an effective context of \(\sim\) 131 K tokens even though no single layer materializes more than a 4096-wide band — the same “information hops one window per layer” argument as a deep stack of convolutions.

Interleaving global and local layers. A pure sliding-window stack still cannot do exact long-range retrieval in one hop, so modern models sprinkle in a few global (full-attention) layers: Gemma 2 alternates local/global 1:1 (\(W = 4096\)); Gemma 3 uses 5 local : 1 global (\(W = 1024\)); GPT-oss alternates banded and dense layers. The global layers restore precise long-range lookup at a small fraction of total attention cost — the same insight that motivates the attention/SSM hybrids later in this chapter.

Sparse Attention Patterns¶

Sliding windows are one point in a larger design space of sparse attention, where the mask is any fixed or learned subset of key positions:

• Strided / block-sparse (Sparse Transformer, Child et al. 2019): combine a local band with a strided pattern (every \(s\)-th position) so two stacked layers cover the whole sequence; blocks are sized to GPU tiles.
• Global + local (+ random) (Longformer, Beltagy et al. 2020; BigBird, Zaheer et al. 2020): a local window, a handful of global tokens that every query attends to (and that attend to everything), plus — in BigBird — a few random keys for expander-graph connectivity. All are \(O(N)\).
• Learned / trainable sparsity (2025): DeepSeek NSA (Native Sparse Attention) runs three branches per query — a compressed branch over coarse block summaries, a selected branch over the top-\(k\) most relevant fine blocks (chosen by the compressed scores), and a sliding-window branch for locality — all trainable end-to-end. MoBA (Mixture of Block Attention) routes each query to a small set of key-blocks with a top-\(k\) gate, importing the MoE idea into attention.

Why blocks? GPU tensor cores compute dense \(128 \times 128\) (or larger) tiles efficiently; token-level sparsity wastes them, whereas block granularity keeps every computed tile dense. This is why production sparse attention is almost always block-sparse.

Legend: X = attention computed,  . = masked/skipped,  G = global token

  Full (causal)      Sliding window W=3     Block-sparse         Global + local
  k:0 1 2 3 4 5 6    k:0 1 2 3 4 5 6        k:0 1 2 3 4 5 6      k:0 1 2 3 4 5 6
q0  X . . . . . .    q0  X . . . . . .      q0  X . . . . . .    q0  G . . . . . .
q1  X X . . . . .    q1  X X . . . . .      q1  X X . . . . .    q1  G X . . . . .
q2  X X X . . . .    q2  X X X . . . .      q2  X X X . . . .    q2  G X X . . . .
q3  X X X X . . .    q3  . X X X . . .      q3  X X X X . . .    q3  G X X X . . .
q4  X X X X X . .    q4  . . X X X . .      q4  X . . X X . .    q4  G . . . X . .
q5  X X X X X X .    q5  . . . X X X .      q5  X X . X X X .    q5  G . . . X X .
q6  X X X X X X X    q6  . . . . X X X      q6  X . . X . X X    q6  G . . . X X X

Real implementations: FlashAttention ships a block_mask/varlen path, PyTorch FlexAttention compiles arbitrary mask_mods into block-sparse kernels, and state-spaces/DeepSeek release fused NSA kernels. Citations: Child et al., Generating Long Sequences with Sparse Transformers (2019); Beltagy et al., Longformer (2020); Zaheer et al., Big Bird (2020); Jiang et al., Mistral 7B (2023); Yuan et al., Native Sparse Attention (2025); Lu et al., MoBA (2025).

Linear Attention: The Kernel Trick¶

The core insight of linear attention, introduced by Katharopoulos et al. (2020), is that the softmax in standard attention is not strictly necessary — it just ensures the output is a weighted sum of values with non-negative, normalized weights. If we can approximate the softmax with a kernel function, we can rewrite the computation to change the order of operations and get \(O(N)\) cost.

The Reformulation¶

Define a feature map \(\phi: \mathbb{R}^{d_k} \rightarrow \mathbb{R}^r\) such that

Then:

The key: compute \(\phi(K)^\top V\) first. This is an \(r \times d_v\) matrix that requires \(O(N r d_v)\) operations and can be computed once. Then multiplying by \(\phi(Q)\) (shape \(N \times r\)) costs another \(O(N r d_v)\). Total is \(O(N)\) rather than \(O(N^2)\).

The price: vanilla linear attention drops the softmax normalization, losing the “focusing” ability that makes attention heads specialize. Empirically, models trained with simple linear attention tend to underperform transformers on tasks requiring selective copying of distant tokens.

The Recurrent Form¶

Linear attention has an equivalent recurrent interpretation that is crucial for efficient inference. Define the outer-product state matrix:

Then the output at time step \(t\) is:

where \(z_t = \sum_{i=1}^t \phi(k_i)\) is a normalizer vector.

This recurrent form processes one token at a time in \(O(r d_v)\) operations with \(O(r d_v)\) state — analogous to an RNN. During training, the parallel form is faster; during autoregressive inference, the recurrent form is constant cost per step regardless of sequence length. This training-parallel / inference-recurrent duality is a recurring theme in this entire chapter.

import torch
import torch.nn.functional as F

def elu_feature_map(x: torch.Tensor) -> torch.Tensor:
    """ELU+1 feature map: keeps positivity, cheap to compute."""
    return F.elu(x) + 1.0  # shape preserved; always > 0

def linear_attention_parallel(
    Q: torch.Tensor,  # (B, N, H, d_k)
    K: torch.Tensor,  # (B, N, H, d_k)
    V: torch.Tensor,  # (B, N, H, d_v)
    causal: bool = True,
) -> torch.Tensor:
    """
    Parallel (training) form of linear attention.
    For causal (autoregressive) models we cannot simply do K^T V first —
    we need the cumulative version. Here we show the non-causal form for clarity.
    The causal form requires a cumulative outer-product scan.
    """
    B, N, H, d_k = Q.shape
    d_v = V.shape[-1]

    phi_Q = elu_feature_map(Q)  # (B, N, H, d_k)
    phi_K = elu_feature_map(K)  # (B, N, H, d_k)

    # Compute K^T V: (B, H, d_k, d_v)
    # phi_K: (B, N, H, d_k) -> reshape to (B, H, N, d_k)
    phi_K_t = phi_K.permute(0, 2, 1, 3)  # (B, H, N, d_k)
    V_t     = V.permute(0, 2, 1, 3)       # (B, H, N, d_v)
    KV      = torch.einsum('bhnk,bhnv->bhkv', phi_K_t, V_t)  # (B, H, d_k, d_v)

    # Normalizer: sum of phi_K over N
    z = phi_K_t.sum(dim=2)  # (B, H, d_k)

    # Compute output
    phi_Q_t = phi_Q.permute(0, 2, 1, 3)  # (B, H, N, d_k)
    # Numerator: phi(Q) @ KV -> (B, H, N, d_v)
    num = torch.einsum('bhnk,bhkv->bhnv', phi_Q_t, KV)
    # Denominator: phi(Q) @ z -> (B, H, N)
    den = torch.einsum('bhnk,bhk->bhn', phi_Q_t, z).unsqueeze(-1)
    # Clamp for numerical stability
    den = den.clamp(min=1e-6)

    out = (num / den).permute(0, 2, 1, 3)  # back to (B, N, H, d_v)
    return out


def linear_attention_recurrent_step(
    q_t: torch.Tensor,  # (B, H, d_k) — current query
    k_t: torch.Tensor,  # (B, H, d_k) — current key
    v_t: torch.Tensor,  # (B, H, d_v) — current value
    S:   torch.Tensor,  # (B, H, d_k, d_v) — running state
    z:   torch.Tensor,  # (B, H, d_k)       — running normalizer
):
    """Single-step recurrent update — O(d_k * d_v) per step, O(d_k * d_v) memory."""
    phi_q = elu_feature_map(q_t)  # (B, H, d_k)
    phi_k = elu_feature_map(k_t)  # (B, H, d_k)

    # Update state: S += phi_k outer v
    S = S + torch.einsum('bhk,bhv->bhkv', phi_k, v_t)
    z = z + phi_k  # (B, H, d_k)

    # Output
    num = torch.einsum('bhk,bhkv->bhv', phi_q, S)   # (B, H, d_v)
    den = torch.einsum('bhk,bhk->bh', phi_q, z).unsqueeze(-1).clamp(min=1e-6)
    return num / den, S, z

The Chunked (Causal) Form¶

The parallel form above computes \(\phi(K)^\top V\) once and shares it across all queries — but that is only valid non-causally, because a causal model must not let query \(t\) see keys with index \(> t\). Running the recurrent step token-by-token restores causality but is \(O(N)\) sequential — a throughput disaster on a GPU. The standard fix, and the algorithmic core reused by GLA, RetNet, and Mamba-2 chunkwise kernels, is to split the sequence into chunks of length \(C\): do exact quadratic masked attention inside each chunk (a parallel matmul), and carry a running outer-product state \(S = \sum \phi(k)\, v^\top\) between chunks.

def linear_attention_chunked(
    Q: torch.Tensor,  # (B, H, N, d_k)  -- head-major here for clean einsums
    K: torch.Tensor,  # (B, H, N, d_k)
    V: torch.Tensor,  # (B, H, N, d_v)
    chunk_size: int = 64,
) -> torch.Tensor:
    """
    Chunked *causal* linear attention (unnormalized numerator).
    Assumes N is divisible by chunk_size. Cost: O(N * C * d) intra-chunk
    matmuls + O((N/C) * d_k * d_v) state carries -- fully parallel within a chunk.
    Divide by the same scan run with V replaced by ones to get the normalizer.
    """
    B, H, N, d_k = Q.shape
    d_v = V.shape[-1]
    phi_Q = elu_feature_map(Q)
    phi_K = elu_feature_map(K)
    n_chunks = N // chunk_size
    phi_Q = phi_Q.view(B, H, n_chunks, chunk_size, d_k)
    phi_K = phi_K.view(B, H, n_chunks, chunk_size, d_k)
    Vc    = V.view(B, H, n_chunks, chunk_size, d_v)

    mask = torch.tril(torch.ones(chunk_size, chunk_size, device=Q.device))  # intra-chunk causal
    S = torch.zeros(B, H, d_k, d_v, device=Q.device, dtype=Q.dtype)         # inter-chunk state
    outs = []
    for c in range(n_chunks):
        q, k, v = phi_Q[:, :, c], phi_K[:, :, c], Vc[:, :, c]   # (B, H, C, *)
        inter = torch.einsum('bhck,bhkv->bhcv', q, S)           # from all previous chunks
        A     = torch.einsum('bhck,bhdk->bhcd', q, k) * mask    # masked within-chunk scores
        intra = torch.einsum('bhcd,bhdv->bhcv', A, v)           # within-chunk contribution
        outs.append(inter + intra)
        S = S + torch.einsum('bhck,bhcv->bhkv', k, v)           # fold chunk into running state
    return torch.cat(outs, dim=2)                                # (B, H, N, d_v)

Tuning \(C\) trades intra-chunk quadratic work (\(O(N C d)\)) against the number of sequential state carries (\(N / C\)); \(C \in [64, 256]\) is typical. Use the following checks to convince yourself the three forms agree (all pass):

if __name__ == "__main__":
    torch.manual_seed(0)

    # (a) Non-causal parallel form matches an explicit softmax-free reference.
    B, N, H, d_k, d_v = 2, 6, 2, 4, 5
    Q = torch.randn(B, N, H, d_k); K = torch.randn(B, N, H, d_k); V = torch.randn(B, N, H, d_v)
    pQ, pK = elu_feature_map(Q), elu_feature_map(K)
    scores = torch.einsum('bnhk,bmhk->bhnm', pQ, pK)
    scores = scores / scores.sum(-1, keepdim=True).clamp(min=1e-6)
    ref = torch.einsum('bhnm,bmhv->bnhv', scores, V)
    assert torch.allclose(linear_attention_parallel(Q, K, V, causal=False), ref, atol=1e-5)
    print("non-causal parallel form matches reference: OK")

    # (b) Recurrent step, rolled over t, reproduces the *causal* output.
    Sr = torch.zeros(B, H, d_k, d_v); zr = torch.zeros(B, H, d_k); rec = []
    for t in range(N):
        y, Sr, zr = linear_attention_recurrent_step(Q[:, t], K[:, t], V[:, t], Sr, zr)
        rec.append(y)
    rec = torch.stack(rec, dim=1)   # (B, N, H, d_v)

    def causal_reference(Q, K, V):
        B, N, H, d_k = Q.shape; d_v = V.shape[-1]
        pQ, pK = elu_feature_map(Q), elu_feature_map(K)
        out = torch.zeros(B, N, H, d_v)
        for t in range(N):
            s = torch.einsum('bhk,bihk->bih', pQ[:, t], pK[:, :t + 1])   # (B, t+1, H)
            num = torch.einsum('bih,bihv->bhv', s, V[:, :t + 1])
            den = s.sum(1).clamp(min=1e-6)                               # (B, H)
            out[:, t] = num / den.unsqueeze(-1)
        return out
    assert torch.allclose(rec, causal_reference(Q, K, V), atol=1e-5)
    print("recurrent rollout reproduces causal output: OK")

    # (c) Chunked numerator == recurrent (unnormalized) numerator.
    B, N, H, d_k, d_v = 2, 32, 2, 4, 5
    Q = torch.randn(B, N, H, d_k); K = torch.randn(B, N, H, d_k); V = torch.randn(B, N, H, d_v)
    Qh, Kh, Vh = [t.permute(0, 2, 1, 3).contiguous() for t in (Q, K, V)]   # (B, H, N, *)
    chk = linear_attention_chunked(Qh, Kh, Vh, chunk_size=8)
    S = torch.zeros(B, H, d_k, d_v); num = []
    for t in range(N):
        pk = elu_feature_map(Kh[:, :, t]); pq = elu_feature_map(Qh[:, :, t])
        S = S + torch.einsum('bhk,bhv->bhkv', pk, Vh[:, :, t])
        num.append(torch.einsum('bhk,bhkv->bhv', pq, S))
    assert torch.allclose(chk, torch.stack(num, dim=2), atol=1e-4)
    print("chunked numerator matches recurrent numerator: OK")

The fla-org/flash-linear-attention library implements exactly this chunked recurrence as fused Triton kernels (chunk_linear_attn, chunk_gla), which is what makes causal linear-attention LMs train at transformer-like throughput.

Structured State Space Models: S4¶

State space models (SSMs) originate in control theory. A continuous-time linear SSM maps an input signal \(u(t)\) to an output signal \(y(t)\) through a hidden state \(x(t)\):

where \(A \in \mathbb{R}^{N \times N}\), \(B \in \mathbb{R}^{N \times 1}\), \(C \in \mathbb{R}^{1 \times N}\), and \(D\) is a skip connection scalar. The Structured State Space Sequence model (S4), introduced by Gu et al. (2021), discretizes this system and makes \(A\) efficiently structured: the HiPPO matrix that preserves long-range history is normal plus low-rank, and S4 conjugates it into a diagonal-plus-low-rank (DPLR) form whose convolution kernel can be evaluated in the frequency domain with a Cauchy kernel rather than by materializing \(\bar A^k\) for every \(k\). With that structure, the entire sequence can be computed as a convolution in \(O(N \log N)\) time during training, or as a recurrence in \(O(1)\) per step during inference. (Later variants — DSS, S4D, and Mamba — showed a purely diagonal \(A\) with a good initialization recovers most of the quality at a fraction of the implementation complexity, which is why every model after S4 in this chapter uses diagonal state matrices.)

Discretization¶

The zero-order hold discretization with step size \(\Delta\) gives:

The recurrence becomes:

And the convolutional kernel is \(\bar{K} = (\bar{C}\bar{B},\, \bar{C}\bar{A}\bar{B},\, \bar{C}\bar{A}^2\bar{B},\, \ldots)\), enabling \(y = u * \bar{K}\) to be computed with an FFT in \(O(N \log N)\).

The HiPPO (High-order Polynomial Projection Operator) initialization of \(A\) — specifically the LegS HiPPO matrix — is what makes S4 effective at remembering distant tokens. The eigenvalue structure of the HiPPO-LegS matrix means the state \(x_t\) projects the history \(u_{\leq t}\) onto Legendre polynomials, giving near-optimal compression of the signal over time.

Convolutional Language Models: H3, Hyena & FlashFFTConv¶

S4’s convolutional view invites a different question than Mamba’s: instead of learning a recurrence, what if we learn the long convolution kernel directly? A branch of the field — largely from Dan Fu, Chris Re, and collaborators — pursued exactly this, and its ideas fed straight into Mamba’s design.

• H3 (Hungry Hungry Hippos; Fu et al., 2023): stacks two SSMs — a shift SSM that acts like a local memory and a diagonal SSM — with multiplicative gating between them, emulating the compare-then-recall motion that attention performs. H3 closed most of the associative-recall gap that plagued plain S4 and directly motivated Mamba’s gated, selective block.
• Hyena (Poli et al., 2023): replaces attention with a data-controlled long convolution — an implicit convolution filter parameterized by a small MLP over positional encodings, interleaved with elementwise gating. Because the filter spans the whole sequence, it is evaluated as an FFT convolution in \(O(N \log N)\), and matches attention quality at sub-billion-parameter scale.
• FlashFFTConv (Fu et al., 2023): the hardware-aware FFT convolution — the long-convolution counterpart of FlashAttention. It recasts the FFT as a Monarch (butterfly) matrix decomposition so the \(O(N \log N)\) convolution runs as a short sequence of dense matmuls on tensor cores, kernel-fused to keep intermediates in SRAM. This turns Hyena/H3-style long convolutions from bandwidth-bound into compute-bound and yields large speedups for long filters.
• Monarch / butterfly matrices (Dao et al., 2022): structured sub-quadratic matrices that factor a dense linear map into a few sparse block-diagonal factors. The DFT, Hadamard transform, and many structured mixers are Monarch-expressible, which is precisely what lets FlashFFTConv map the FFT onto tensor cores.
• Based (Arora et al., 2024): a hybrid token mixer that combines a Taylor-feature-map linear attention (a 2^nd-order Taylor approximation of the softmax, giving a large recurrent state for global recall) with a short sliding-window attention branch (for precise local shifting). Based recovers much of the associative-recall quality that pure linear attention loses, at a fraction of the KV-cache memory — an explicit bridge between the linear-attention and sliding-window sections of this chapter.

Citations: Fu et al., Hungry Hungry Hippos (H3), arXiv:2212.14052; Poli et al., Hyena Hierarchy, arXiv:2302.10866; Fu et al., FlashFFTConv, arXiv:2311.05908; Dao et al., Monarch: Expressive Structured Matrices, arXiv:2204.00595; Arora et al., Based, arXiv:2402.18668.

Mamba: Selective State Space Models¶

S4 and its relatives (S5, DSS, GSS) proved that SSMs can model long sequences efficiently, but they struggled on tasks requiring selective recall — like looking up a specific value from a long context. The state matrix \(A\), \(B\), \(C\) in S4 are fixed for all time steps (content-independent), so the model cannot decide which inputs to focus on.

Mamba (Gu & Dao, 2023) introduces input-dependent (selective) parameters: \(B_t\), \(C_t\), and \(\Delta_t\) are now functions of the input \(u_t\), computed through a linear projection. This selectivity mechanism is the central innovation.

The Selective Scan¶

The Mamba recurrence for a single channel is:

where \(\bar{A}_t = \exp(\Delta_t A)\) and \(\bar{B}_t = (\exp(\Delta_t A) - I) A^{-1} B_t\). Now \(\Delta_t\), \(B_t\), \(C_t\) are all computed from the input, giving the model the ability to selectively gate what information flows into the state.

Critically, \(A\) is kept diagonal (not the full HiPPO matrix) for efficiency, and initialized with a specific negative-real diagonal structure that encourages stable long-range memory.

Mamba Block (per layer):

Input x: (B, L, D)
         |
    Linear proj → z: (B, L, D)   [skip branch, gated by SiLU]
         |
    Linear proj → u: (B, L, expand*D)   [SSM branch]
         |
    Depthwise conv (width 4)
         |
    SiLU activation
         |
    Selective SSM:
      Linear proj → delta: (B, L, expand*D)
      Linear proj → B:     (B, L, N)      [state dim N ~16]
      Linear proj → C:     (B, L, N)
      Discretize (A fixed diagonal, per-channel)
      Selective scan recurrence
         |
    Multiply by z (gating)
         |
    Linear proj back to D
         |
Output: (B, L, D)

The selective scan is the computationally hot path. Naively it is \(O(L \cdot D \cdot N)\) with \(O(D \cdot N)\) state, but running it sequentially on a GPU is extremely memory-bandwidth-bound. Mamba solves this with a custom parallel scan algorithm (prefix sum / work-efficient parallel scan) and a hardware-aware kernel that fuses the scan, keeps intermediate activations in SRAM, and avoids materializing the full state sequence in HBM.

import torch
import torch.nn as nn
import math

class SelectiveSSM(nn.Module):
    """
    Simplified Mamba-style selective SSM for one feature dimension.
    Full Mamba also uses depthwise conv before this and expand*D channels;
    this illustrates the core selective scan logic.

    d_model: input/output dimension
    d_state: SSM state dimension N (Mamba uses 16)
    """
    def __init__(self, d_model: int, d_state: int = 16, dt_rank: int = None):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        dt_rank = dt_rank or math.ceil(d_model / 16)

        # Log of A diagonal: initialized to evenly spaced negative reals
        # A = -exp(log_A), kept negative for stability
        A_log = torch.log(torch.arange(1, d_state + 1, dtype=torch.float32)
                          .unsqueeze(0).repeat(d_model, 1))
        self.A_log = nn.Parameter(A_log)  # (d_model, d_state)

        # D skip connection (one per channel)
        self.D = nn.Parameter(torch.ones(d_model))

        # Projections for input-dependent parameters
        self.x_proj = nn.Linear(d_model, dt_rank + 2 * d_state, bias=False)
        self.dt_proj = nn.Linear(dt_rank, d_model, bias=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x: (B, L, d_model)
        returns: (B, L, d_model)
        """
        B, L, D = x.shape
        N = self.d_state

        # A: (D, N), kept negative for stability
        A = -torch.exp(self.A_log)  # (D, N)

        # Compute input-dependent B, C, delta
        x_dbl = self.x_proj(x)  # (B, L, dt_rank + 2*N)
        dt, B_ssm, C = x_dbl.split([self.dt_proj.in_features, N, N], dim=-1)
        dt = torch.nn.functional.softplus(self.dt_proj(dt))  # (B, L, D), positive

        # Discretize: zero-order hold
        # dA[t] = exp(dt[t] * A)  — shape (B, L, D, N)
        dA = torch.exp(dt.unsqueeze(-1) * A.unsqueeze(0).unsqueeze(0))  # (B,L,D,N)
        # dB[t] = dt[t] * B[t] (simplified ZOH for diagonal A)
        dB = dt.unsqueeze(-1) * B_ssm.unsqueeze(2)  # (B, L, D, N)

        # Sequential scan (simple version — production uses parallel scan)
        h = torch.zeros(B, D, N, device=x.device, dtype=x.dtype)
        ys = []
        for t in range(L):
            # h: (B, D, N), dA[:,t]: (B, D, N), dB[:,t]: (B, D, N)
            # x[:,t]: (B, D)
            h = dA[:, t] * h + dB[:, t] * x[:, t].unsqueeze(-1)
            # y: (B, D) via C projection
            y = (h * C[:, t].unsqueeze(1)).sum(-1)  # (B, D)
            ys.append(y)

        # Stack and add skip connection
        y_seq = torch.stack(ys, dim=1)  # (B, L, D)
        y_seq = y_seq + x * self.D.unsqueeze(0).unsqueeze(0)
        return y_seq


class MambaBlock(nn.Module):
    """
    Full Mamba block with expansion, depthwise conv, gating, and SSM.
    expand: expansion factor (Mamba uses 2)
    d_conv: depthwise conv width (Mamba uses 4)
    """
    def __init__(self, d_model: int, d_state: int = 16,
                 expand: int = 2, d_conv: int = 4):
        super().__init__()
        self.d_model = d_model
        self.d_inner = d_model * expand

        # Input projection: split into SSM branch (d_inner) and gate (d_inner)
        self.in_proj = nn.Linear(d_model, 2 * self.d_inner, bias=False)

        # Causal depthwise convolution over sequence dimension
        self.conv1d = nn.Conv1d(
            in_channels=self.d_inner,
            out_channels=self.d_inner,
            kernel_size=d_conv,
            padding=d_conv - 1,          # left-pad for causality
            groups=self.d_inner,         # depthwise
            bias=True,
        )
        self.act = nn.SiLU()

        # SSM
        self.ssm = SelectiveSSM(self.d_inner, d_state=d_state)

        # Output projection
        self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """x: (B, L, d_model) -> (B, L, d_model)"""
        B, L, _ = x.shape

        # Split into SSM branch u and gate z
        xz = self.in_proj(x)                   # (B, L, 2*d_inner)
        u, z = xz.chunk(2, dim=-1)             # each (B, L, d_inner)

        # Depthwise conv (operates on L dimension)
        u = u.transpose(1, 2)                  # (B, d_inner, L)
        u = self.conv1d(u)[..., :L]            # causal: trim right
        u = u.transpose(1, 2)                  # (B, L, d_inner)
        u = self.act(u)

        # SSM
        y = self.ssm(u)                        # (B, L, d_inner)

        # Gate: multiply by SiLU(z)
        y = y * self.act(z)

        # Output projection
        return self.out_proj(y)                # (B, L, d_model)

# Smoke test: run SelectiveSSM and the full MambaBlock end-to-end, then verify
# the sequential selective scan against a cumulative-product closed form.
if __name__ == "__main__":
    torch.manual_seed(0)
    B, L, d_model = 2, 16, 32
    x = torch.randn(B, L, d_model)

    ssm = SelectiveSSM(d_model, d_state=8)       # note d_state != d_model on purpose
    y = ssm(x)
    print("SelectiveSSM:", tuple(x.shape), "->", tuple(y.shape))   # (2,16,32) -> (2,16,32)
    assert y.shape == x.shape

    block = MambaBlock(d_model, d_state=8)
    y = block(x)
    print("MambaBlock:  ", tuple(x.shape), "->", tuple(y.shape))    # (2,16,32) -> (2,16,32)
    assert y.shape == x.shape

    # Selective scan vs closed form:
    #   h_t = sum_{i<=t} (prod_{j=i+1}^{t} dA_j) * dB_i * x_i,   y_t = sum_n C[t,n] * h[t,n]
    torch.manual_seed(1)
    B, L, D, N = 1, 5, 2, 3
    dA = torch.rand(B, L, D, N) * 0.5 + 0.4      # per-step decays in (0.4, 0.9)
    dB = torch.randn(B, L, D, N)
    xin = torch.randn(B, L, D)
    C = torch.randn(B, L, N)

    # (i) sequential scan -- the exact loop used inside SelectiveSSM.forward
    h = torch.zeros(B, D, N); seq = []
    for t in range(L):
        h = dA[:, t] * h + dB[:, t] * xin[:, t].unsqueeze(-1)
        seq.append((h * C[:, t].unsqueeze(1)).sum(-1))
    seq = torch.stack(seq, dim=1)

    # (ii) cumulative-product closed form
    clo = []
    for t in range(L):
        h = torch.zeros(B, D, N)
        for i in range(t + 1):
            prod = torch.ones(B, D, N)
            for j in range(i + 1, t + 1):
                prod = prod * dA[:, j]
            h = h + prod * dB[:, i] * xin[:, i].unsqueeze(-1)
        clo.append((h * C[:, t].unsqueeze(1)).sum(-1))
    clo = torch.stack(clo, dim=1)

    assert torch.allclose(seq, clo, atol=1e-5)
    print("selective scan matches cumulative-product closed form: OK")

Training at Speed: The Parallel (Chunked) Scan¶

The sequential for t in range(L) loop above is correct, but it serializes the sequence and is memory-bandwidth-bound — it is to SSMs what naive attention is to transformers. Production kernels replace it with a hardware-aware scan, the SSM counterpart of the FlashAttention story. There are two routes.

Route 1: the associative (prefix) scan. The linear recurrence \(h_t = a_t\, h_{t-1} + b_t\) (with \(a_t = \bar A_t\), \(b_t = \bar B_t x_t\)) is an affine map \(h \mapsto a_t h + b_t\), and composition of affine maps is associative. Composing an earlier map \((a_e, b_e)\) then a later map \((a_l, b_l)\) gives

Because the operator is associative, a work-efficient parallel prefix scan computes all prefix states \(h_0, \ldots, h_{L-1}\) in \(O(\log L)\) sequential steps instead of \(O(L)\) — this is exactly what torch.associative_scan and Mamba’s CUDA selective_scan do under the hood.

import torch

def associative_scan(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """
    Parallel (Hillis-Steele) inclusive scan of the affine recurrence
        h_t = a_t * h_{t-1} + b_t,   h_{-1} = 0
    a, b: (L, ...) with matching trailing shape, e.g. (L, D, N).
    Runs in O(log L) doubling steps; returns h stacked along dim 0.
    Combine (earlier=(a_e,b_e) then later=(a_l,b_l)) -> (a_l*a_e, a_l*b_e + b_l).
    """
    L = a.shape[0]
    a = a.clone(); b = b.clone()
    d = 1
    while d < L:
        a_prev, b_prev = a[:-d], b[:-d]        # partials ending d steps earlier
        a_new, b_new = a.clone(), b.clone()
        a_new[d:] = a[d:] * a_prev
        b_new[d:] = a[d:] * b_prev + b[d:]
        a, b = a_new, b_new
        d *= 2
    return b   # h_t == accumulated b_t because h_{-1} = 0

if __name__ == "__main__":
    torch.manual_seed(0)
    L, D, N = 128, 4, 3
    a = torch.rand(L, D, N) * 0.9              # per-step decays in (0, 0.9)
    b = torch.randn(L, D, N)

    par = associative_scan(a, b)

    h = torch.zeros(D, N); ref = []            # ground-truth sequential loop
    for t in range(L):
        h = a[t] * h + b[t]
        ref.append(h.clone())
    ref = torch.stack(ref, dim=0)

    assert torch.allclose(par, ref, atol=1e-4)
    print("parallel associative scan matches sequential loop: OK")

To drive the selective SSM with it, set a = dA and b = dB * x[..., None] (both shaped (L, B, D, N)), call associative_scan over the time axis to get every h_t at once, then read out y_t = (h_t * C_t).sum(-1).

Route 2: the chunkwise (intra-chunk parallel + inter-chunk recurrent) algorithm. This is what Mamba-2 (SSD) and GLA actually ship, because it maps onto tensor-core matmuls rather than a scalar scan. Split the sequence into chunks of length \(C\). Within a chunk, materialize the \(C \times C\) decay-weighted score matrix and multiply by the chunk’s values — a dense masked matmul, identical in shape to linear_attention_chunked above but with the data-dependent decay mask \(L_{ij} = \prod_{k=j+1}^{i} a_k\) in place of the plain causal mask. Between chunks, carry the low-rank state \(S \in \mathbb{R}^{d \times N}\) and pass it forward. Intra-chunk work is \(O((N/C) \cdot C^2 \cdot d) = O(N C d)\) of tensor-core matmul; inter-chunk work is \(O((N/C) \cdot d N)\) of state passing. See state-spaces/mamba (mamba_chunk_scan_combined) and fla-org/flash-linear-attention (chunk_gla, chunk_simple_gla) for the fused Triton/CUDA implementations.

Mamba-2 and the State Space Duality¶

Mamba-2 (Dao & Gu, 2024) reframes the selective SSM as a special case of structured matrix multiplication, revealing a duality between SSMs and linear attention. Specifically, with a scalar \(A_t\) (instead of a diagonal matrix), the Mamba-2 selective scan is mathematically equivalent to linear attention with a specific data-dependent decay mask. This insight leads to a cleaner algorithm, better hardware utilization (tiled matrix multiplications instead of a sequential scan), and strong theoretical grounding.

The key operation in Mamba-2 is the SSD (State Space Duality) layer, which computes:

Here \(C B^\top\) is the linear-attention score matrix (entry \(C_i^\top B_j\)), and \(L\) is the 1-semiseparable causal decay mask built from the scalar decays \(a_k = \exp(\Delta_k A)\). In words: the SSD layer is exactly linear attention with a data-dependent multiplicative causal mask — swap the softmax’s normalization for the cumulative-product decays \(L_{ij}\) and you recover the selective SSM. In practice \(Y\) is computed by a block decomposition of \(L\): diagonal blocks are handled by masked intra-chunk matmuls, off-diagonal blocks by low-rank inter-chunk state passing, which is what lets Mamba-2 run on tensor cores.

RWKV: RNNs Strike Back¶

RWKV (Peng et al., 2023) takes a different philosophical approach: it revisits the recurrent neural network (RNN) architecture but designs it to be trainable in parallel (like a transformer) by expressing the recurrence as a time-weighted attention.

The WKV Attention¶

The central operation in RWKV is WKV (time-mixing through weighted key-value aggregation):

Here \(w \in \mathbb{R}^d\) is a learned decay vector (one per channel), \(u \in \mathbb{R}^d\) is a bonus for the current token, and \(k_t, v_t \in \mathbb{R}^d\) are per-token key and value vectors. This is essentially exponentially-decayed attention, where older tokens get exponentially down-weighted by a channel-wise learned rate.

The recurrent form has a scalar state per channel and can be written as:

This recurrence runs in \(O(1)\) per step with \(O(d)\) state — vastly smaller state than a transformer’s KV cache.

RWKV Versions¶

RWKV-6 and RWKV-7 introduce data-dependent decays, similar to Mamba’s selectivity, and represent a convergence of the two design philosophies.

import torch
import torch.nn as nn

class RWKVTimeMixing(nn.Module):
    """
    RWKV-4 style time mixing block.
    n_embd: model dimension
    layer_id: used to initialize w differently per layer
    """
    def __init__(self, n_embd: int, layer_id: int = 0):
        super().__init__()
        self.n_embd = n_embd

        # Learned time decay (w): initialized to negative values (decay)
        # Stored as log(-w) for stability; w should be < 0
        self.w = nn.Parameter(torch.zeros(n_embd))
        nn.init.uniform_(self.w, -8.0, -0.5)

        # Time-first (u): bonus for current token
        self.u = nn.Parameter(torch.zeros(n_embd))
        nn.init.uniform_(self.u, -0.5, 0.5)

        # Token shift mixing ratios
        self.time_mix_k = nn.Parameter(torch.ones(1, 1, n_embd))
        self.time_mix_v = nn.Parameter(torch.ones(1, 1, n_embd))
        self.time_mix_r = nn.Parameter(torch.ones(1, 1, n_embd))

        # Projections
        self.key    = nn.Linear(n_embd, n_embd, bias=False)
        self.value  = nn.Linear(n_embd, n_embd, bias=False)
        self.receptance = nn.Linear(n_embd, n_embd, bias=False)
        self.output = nn.Linear(n_embd, n_embd, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x: (B, T, n_embd)
        Token shift: mix current and previous token for k, v, r inputs.
        Then run the WKV scan.
        """
        B, T, C = x.shape

        # Token shift: concat zero-padded version of x (shifted right by 1)
        x_prev = torch.cat([torch.zeros(B, 1, C, device=x.device), x[:, :-1]], dim=1)

        # Compute K, V, R with time-mixed inputs
        xk = x * self.time_mix_k + x_prev * (1 - self.time_mix_k)
        xv = x * self.time_mix_v + x_prev * (1 - self.time_mix_v)
        xr = x * self.time_mix_r + x_prev * (1 - self.time_mix_r)

        k = self.key(xk)        # (B, T, C)
        v = self.value(xv)      # (B, T, C)
        r = torch.sigmoid(self.receptance(xr))  # gating

        # WKV computation (simplified sequential loop over T)
        # In practice this is a custom CUDA kernel for speed
        w = -torch.exp(self.w)   # (C,), negative → exponential decay
        u = self.u               # (C,)

        # Running accumulators
        aa = torch.zeros(B, C, device=x.device)  # numerator state
        bb = torch.zeros(B, C, device=x.device)  # denominator state
        pp = torch.full((B, C), -1e38, device=x.device)  # log-sum for numerical stability

        wkv_outputs = []
        for t in range(T):
            kt = k[:, t]  # (B, C)
            vt = v[:, t]  # (B, C)

            # Numerically stable WKV update using log-space
            # Output uses the *undecayed* running state a_{t-1}, b_{t-1};
            # the decay w is applied only in the state update below.
            # (Matches the displayed WKV recurrence and the official RWKV-4 kernel.)
            qq = torch.maximum(pp, kt + u)         # (B, C)
            a  = torch.exp(pp - qq) * aa + torch.exp(kt + u - qq) * vt
            b  = torch.exp(pp - qq) * bb + torch.exp(kt + u - qq)
            wkv_outputs.append(a / b.clamp(min=1e-8))

            # Update state
            qq2 = torch.maximum(pp + w, kt)
            aa  = torch.exp(pp + w - qq2) * aa + torch.exp(kt - qq2) * vt
            bb  = torch.exp(pp + w - qq2) * bb + torch.exp(kt - qq2)
            pp  = qq2

        wkv = torch.stack(wkv_outputs, dim=1)  # (B, T, C)
        out = r * wkv
        return self.output(out)

RetNet and the Retention Mechanism¶

RetNet (Sun et al., 2023) proposes Retention — a recurrence-free, attention-free sequence model that explicitly targets the “training parallelism vs. inference efficiency” tradeoff with what the authors call the “impossible triangle” — the claim that retention achieves all three: training parallelism, \(O(1)\) inference, and good performance.

The retention score between query \(q_i\) and key \(k_j\) is:

where \(\gamma \in (0, 1)\) is a per-head decay constant (not learned — fixed at training time based on head index). This can equivalently be written in a parallel (training) form as a lower-triangular masked matrix or in a recurrent (inference) form as an \(O(d_k \times d_v)\) state matrix updated at each step.

Multi-scale retention uses different \(\gamma\) values per head, allowing different heads to specialize in short-range versus long-range dependencies. Chunk-wise retention computes retention over fixed-size chunks, balancing the efficiency of both forms.

RetNet replaces the softmax normalization with a group normalization applied to the retention output — a practical fix for the scale instability of un-normalized attention.

The FLOP / Memory Trade-off Table¶

Before we look at hybrids, it is worth summarizing the algorithmic tradeoffs:

\(N_s\) denotes SSM state dimension (e.g., 16 in Mamba). Note that Mamba/SSM state size is far smaller than the KV cache of a transformer at long sequences.

Hybrid Architectures: Jamba, Zamba & Beyond¶

Pure SSM models have one well-documented weakness: they struggle at tasks requiring precise in-context retrieval — looking up an exact value from the context, copying specific tokens, or doing multi-hop reasoning across distant facts. This is because the fixed-size SSM state is a lossy compression of the history, while softmax attention can attend to exact past tokens. The limit is quantitative, not just qualitative: a recurrent layer carries a fixed budget of roughly \(d \cdot N_s\) scalars, so once the number of key–value associations the context asks it to hold exceeds what that budget can encode, recall must degrade. The Zoology/Based line of work (Arora et al., 2024) measured exactly this recall–memory tradeoff, finding associative-recall accuracy improving smoothly as the recurrent state is enlarged. A transformer escapes the bound only by letting its state — the KV cache — grow with \(N\), which is precisely why it is expensive.

The natural solution is to combine attention and SSM layers in the same model.

Jamba (AI21 Labs, 2024)¶

Jamba interleaves Mamba SSM blocks and transformer (attention + MLP) blocks, with a ratio heavily weighted toward Mamba (roughly 1 attention layer per 8 Mamba layers). It also incorporates MoE (Mixture-of-Experts; see Mixture-of-Experts (MoE) Architectures) in the MLP blocks to increase model capacity without proportionally increasing compute.

The key design insight: a small number of attention layers provides the exact-retrieval capability that SSMs lack, while the majority of layers being Mamba keeps memory and compute efficient. Jamba demonstrated that this hybrid could match or exceed pure transformer models on standard benchmarks while using far less KV-cache memory at long contexts.

Zamba and Other Hybrids¶

Zamba (Zyphra, 2024) uses a 1-attention-per-6-Mamba ratio with a shared global attention block that is applied at multiple depths in the network (parameter sharing). This reduces the number of attention-induced KV cache entries further.

The general design space for hybrids can be parameterized by: 1. Attention-to-SSM ratio: how many attention layers vs. SSM layers 2. Placement: interleaved evenly, attention at top/bottom, or clustered 3. Attention type: full attention, sliding window attention, or grouped-query attention (see Multi-Head Attention, MQA, GQA & MLA) 4. SSM variant: Mamba-1, Mamba-2, GLA, gated DeltaNet, or RWKV blocks

Production Hybrids (2025-2026)¶

What was a research demo with Jamba has become the default recipe for a wave of shipped, general-purpose 2025 models. NVIDIA’s Nemotron-H (2025) replaces the majority of self-attention layers in an 8B/56B transformer with Mamba-2 layers, reporting up to ~3x faster inference at accuracy on par with similarly-sized Qwen-2.5 and Llama-3.1 models. IBM’s Granite 4.0 (October 2025) ships a Mamba-2 + attention + MoE stack (roughly 4 attention to 36 Mamba-2 layers) as a mainstream enterprise model line. On the token-mixer side, the gated delta rule (Gated DeltaNet, Yang et al., 2025 — derived below) augments Mamba-2’s gating with the delta update rule, and its kernels now underpin several of these production hybrids: Alibaba’s Qwen3-Next (2025) interleaves gated-DeltaNet layers with a minority of gated full-attention layers (roughly three recurrent layers per attention layer) on top of a sparse MoE. The through-line is unchanged from Jamba: a few attention layers for exact retrieval, many recurrent/SSM layers for cheap long-context memory.

A Minimal Hybrid Model¶

import torch
import torch.nn as nn
from typing import Literal

# We assume MambaBlock and a standard TransformerBlock are already defined.
# Below shows how to compose them in a hybrid model.

class HybridBlock(nn.Module):
    """
    A single hybrid layer that is either an SSM block or an Attention block.
    block_type: 'mamba' | 'attention'
    """
    def __init__(
        self,
        d_model: int,
        block_type: Literal["mamba", "attention"],
        n_heads: int = 8,
        d_state: int = 16,
    ):
        super().__init__()
        self.block_type = block_type
        self.norm = nn.RMSNorm(d_model)

        if block_type == "mamba":
            self.block = MambaBlock(d_model, d_state=d_state)
        elif block_type == "attention":
            # Minimal multi-head self-attention
            self.block = nn.MultiheadAttention(
                d_model, n_heads, batch_first=True
            )

    def forward(self, x: torch.Tensor, attn_mask=None) -> torch.Tensor:
        residual = x
        x = self.norm(x)
        if self.block_type == "mamba":
            x = self.block(x)
        else:
            x, _ = self.block(x, x, x, attn_mask=attn_mask)
        return residual + x


class HybridLanguageModel(nn.Module):
    """
    A hybrid SSM-Attention language model.
    attn_every: place an attention layer every N layers (rest are Mamba).
    E.g. attn_every=8 → 1 attention layer per 7 Mamba layers.
    """
    def __init__(
        self,
        vocab_size: int,
        d_model: int,
        n_layers: int,
        attn_every: int = 8,
        n_heads: int = 8,
        d_state: int = 16,
    ):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)

        layers = []
        for i in range(n_layers):
            # Place attention layers at fixed intervals
            if (i + 1) % attn_every == 0:
                block_type = "attention"
            else:
                block_type = "mamba"
            layers.append(
                HybridBlock(d_model, block_type, n_heads=n_heads, d_state=d_state)
            )

        self.layers = nn.ModuleList(layers)
        self.norm_out = nn.RMSNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, tokens: torch.Tensor) -> torch.Tensor:
        """tokens: (B, L) -> logits (B, L, vocab_size)"""
        x = self.embed(tokens)  # (B, L, d_model)
        for layer in self.layers:
            x = layer(x)
        x = self.norm_out(x)
        return self.lm_head(x)


# Quick sanity check
if __name__ == "__main__":
    model = HybridLanguageModel(
        vocab_size=32000, d_model=512, n_layers=12, attn_every=4
    )
    # Count attention vs Mamba layers
    attn_count  = sum(1 for l in model.layers if l.block_type == "attention")
    mamba_count = sum(1 for l in model.layers if l.block_type == "mamba")
    print(f"Attention layers: {attn_count}, Mamba layers: {mamba_count}")
    # → Attention layers: 3, Mamba layers: 9

    tokens = torch.randint(0, 32000, (2, 128))
    logits = model(tokens)
    print(f"Output shape: {logits.shape}")  # (2, 128, 32000)

From Scratch to Production: The Real Libraries¶

Everything above is written out so the mechanism is visible; nobody trains with a Python for t in range(L) loop. Four rungs of the real stack — reference blocks, portable kernels, pretrained checkpoints, serving — and the package that owns each:

• state-spaces/mamba — the reference Mamba/Mamba-2 blocks with their fused CUDA kernels: selective_scan_fn (Mamba-1’s hardware-aware scan) and mamba_chunk_scan_combined (Mamba-2’s chunkwise SSD), plus the separate causal-conv1d package for the fused depthwise conv. Mamba2(d_model=...) is a drop-in replacement for the MambaBlock we wrote, with the same (B, L, D) -> (B, L, D) signature. Requires a CUDA GPU.
• fla-org/flash-linear-attention (imported as fla) — Triton kernels for the whole gated-linear-recurrence family (GLA, RetNet, Mamba-2, DeltaNet, gated DeltaNet, RWKV-6/7, Based), exposed both as functional ops (fla.ops.gla.chunk_gla) and as nn.Module layers (fla.layers.GatedLinearAttention, fla.layers.GatedDeltaNet). This is what you reach for to swap the token mixer in an existing PyTorch model, and it runs on NVIDIA, AMD, and Intel backends.
• HF transformers — MambaForCausalLM, Mamba2ForCausalLM, JambaForCausalLM and their siblings give published checkpoints (e.g. state-spaces/mamba-130m-hf) the same from_pretrained / generate interface as any transformer, so your training, eval, and tokenizer code is unchanged. Note that generation carries a MambaCache (or a hybrid cache) rather than a KV cache — a fixed-size tensor of conv states and SSM states.
• vLLM / SGLang — serving. vLLM allocates a fixed-size state cache for the recurrent layers alongside the paged KV cache for the (few) attention layers, which is exactly what makes hybrids such as Jamba, Nemotron-H, and Granite 4.0 cheap to serve at long context; see vLLM: Architecture, PagedAttention & Internals and PagedAttention & KV-Cache Memory Management.

# CUDA GPU required for the fused kernels; causal-conv1d is a separate build.
pip install mamba-ssm causal-conv1d
pip install flash-linear-attention   # imported as `fla`

import torch

# --- Real Mamba-2 block (state-spaces/mamba) ---------------------------------
try:
    from mamba_ssm import Mamba2
except ImportError:
    Mamba2 = None

# --- Real gated-linear-attention / gated-DeltaNet layers (fla) ---------------
try:
    from fla.layers import GatedLinearAttention
except ImportError:
    GatedLinearAttention = None

if __name__ == "__main__":
    B, L, D = 2, 64, 256
    if Mamba2 is not None and torch.cuda.is_available():
        x = torch.randn(B, L, D, device="cuda", dtype=torch.bfloat16)
        # d_state=64 and headdim=64 are Mamba-2's defaults; d_inner = expand*D = 512.
        block = Mamba2(d_model=D, d_state=64, d_conv=4, expand=2, headdim=64).cuda().to(torch.bfloat16)
        y = block(x)                       # (B, L, D) -> (B, L, D), fused chunkwise SSD
        assert y.shape == x.shape
        print("mamba_ssm.Mamba2 OK:", tuple(y.shape))

    if GatedLinearAttention is not None and torch.cuda.is_available():
        x = torch.randn(B, L, D, device="cuda", dtype=torch.bfloat16)
        gla = GatedLinearAttention(hidden_size=D, num_heads=4, mode="chunk").cuda().to(torch.bfloat16)
        out = gla(x)                       # recent fla returns (output, attn_weights, cache)
        y = out[0] if isinstance(out, tuple) else out
        assert y.shape == x.shape
        print("fla GatedLinearAttention OK:", tuple(y.shape))

# Pretrained SSM checkpoints behave like any HF causal LM (downloads weights):
#   from transformers import AutoTokenizer, AutoModelForCausalLM
#   tok = AutoTokenizer.from_pretrained("state-spaces/mamba-130m-hf")
#   model = AutoModelForCausalLM.from_pretrained("state-spaces/mamba-130m-hf")
#   ids = tok("State space models are", return_tensors="pt").input_ids
#   print(tok.decode(model.generate(ids, max_new_tokens=20)[0]))

Building a ~100M Hybrid Yourself¶

If you are swapping some of your own model’s attention blocks for recurrent ones — the ablation offered in The Stack-100M Architecture: SOTA Components, Cited and Assembled, whose gated short-convolution mixer is the cheapest member of this family — four things change, and all four are easy to get wrong:

1. Parameter accounting. A Mamba block at expand=2 costs roughly \(6d^2\) parameters (in-projection \(4d^2\), out-projection \(2d^2\), plus negligible conv/\(\Delta\)/\(B\)/\(C\) terms), while a transformer layer costs \(4d^2\) (attention) \(+\ 8d^2\) (SwiGLU MLP) \(\approx 12d^2\). So a Mamba stack needs about twice the depth of a transformer to match parameters — the substitution is not one block for one block, and forgetting this silently halves your model.
2. Positional encoding. SSM, conv, and decayed-linear-attention mixers are implicitly positional: order enters through the recurrence itself. They take no RoPE. In a hybrid, apply RoPE only inside the attention layers — or, as Jamba does, drop explicit positional encodings entirely and let the recurrent layers carry position (Positional Encodings: Sinusoidal, Learned, RoPE & ALiBi).
3. Initialization of \(\Delta\). The dt_proj bias must be initialized so that \(\operatorname{softplus}(\Delta)\) starts in roughly \([0.001, 0.1]\) (code below); leave PyTorch’s default nn.Linear bias init in place and \(\Delta\) starts near \(0.7\), so \(\bar A = \exp(\Delta A)\) decays hard at every step, the state is overwritten each token, and the model never learns long-range memory. The reference recipe samples \(\Delta\) log-uniformly in that range and stores its inverse-softplus in the bias.
4. Optimizer groups. Exclude A_log, D, and the conv/dt biases from weight decay — they are not “weights” in the usual sense, and decaying A_log drags every state toward instant forgetting. This is the same parameter-group discipline as norms and biases in a transformer (Optimizers: SGD, Adam, Adafactor, Lion, Muon & Shampoo), and it slots into the capstone’s Muon/AdamW split in Optimizer & Schedule: Muon + MuonClip and Warmup-Stable-Decay.

import math
import torch
import torch.nn as nn

def init_dt_bias(dt_proj: nn.Linear, dt_min: float = 1e-3, dt_max: float = 1e-1) -> nn.Linear:
    """
    Mamba's Delta initialization: softplus(dt_proj.bias) is log-uniform in [dt_min, dt_max].
    Store the *inverse* softplus in the bias so the forward pass's softplus recovers dt.
    """
    d = dt_proj.out_features
    dt = torch.exp(torch.rand(d) * (math.log(dt_max) - math.log(dt_min)) + math.log(dt_min))
    inv_dt = dt + torch.log(-torch.expm1(-dt))      # inverse of softplus, numerically stable
    with torch.no_grad():
        dt_proj.bias.copy_(inv_dt)
    return dt_proj

def ssm_param_groups(model: nn.Module, weight_decay: float = 0.1):
    """No weight decay on A_log, D, biases, or norm/conv-1d parameters."""
    no_decay = [p for n, p in model.named_parameters()
                if p.ndim <= 1 or n.endswith(("A_log", "D"))]
    decay    = [p for n, p in model.named_parameters()
                if not (p.ndim <= 1 or n.endswith(("A_log", "D")))]
    return [{"params": decay, "weight_decay": weight_decay},
            {"params": no_decay, "weight_decay": 0.0}]

if __name__ == "__main__":
    torch.manual_seed(0)
    proj = init_dt_bias(nn.Linear(16, 128))
    dt = torch.nn.functional.softplus(proj.bias)
    assert dt.min() > 5e-4 and dt.max() < 0.2
    print(f"dt range after init: [{dt.min():.4f}, {dt.max():.4f}]")   # ~[0.001, 0.1]

    groups = ssm_param_groups(SelectiveSSM(64, d_state=8))
    print("decayed tensors:", len(groups[0]["params"]),
          "| undecayed (A_log, D, biases):", len(groups[1]["params"]))

At 100M parameters and a 2048–8192-token context, note honestly that a hybrid buys you very little: the KV cache of a 100M transformer at 8K context is only tens of megabytes, so the memory argument that motivates Nemotron-H does not bite. Build one to measure the tradeoff — a real ablation you can run in an afternoon — not because it is the better model at this scale.

Gated Linear Attention and the Convergence¶

The theoretical connections between all these architectures have become clearer over 2024-2025. The Gated Linear Attention (GLA) formulation (Yang et al., 2023) provides a unified view: all of these models can be seen as linear attention variants with different forms of gating and decay.

The general GLA recurrence is:

where \(G_t \in (0,1)^{d_k \times d_v}\) is a data-dependent gate (or decay) matrix. Different choices of \(G_t\) recover different architectures:

• \(G_t = \text{diag}(\gamma)\) (constant) → RetNet
• \(G_t = \text{diag}(\exp(-\exp(w)))\) (input-dependent scalar per channel) → RWKV-style
• \(G_t = \text{diag}(\exp(\Delta_t A))\) (Mamba-style discretization) → Mamba / selective SSM
• \(G_t = I\) (no decay) → linear attention

This convergence has practical implications: hardware-efficient kernels developed for one variant (e.g., GLA’s chunked parallel scan) can be adapted for others, and model designers can tune the gate expressiveness/cost tradeoff without changing the high-level architecture.

The Delta Rule: Overwriting Instead of Accumulating¶

One important family does not fit the elementwise-gate template, and it is worth stating explicitly because it is the token mixer behind the strongest 2025–2026 linear models. Every recurrence above adds \(k_t \otimes v_t\) into the state; nothing ever removes a stale association, so a key written early keeps polluting later reads. DeltaNet (Schlag et al., 2021; made hardware-efficient by Yang et al., 2024) replaces the additive write with an error-correcting one:

Read it as erase then write: the projector \(I - \beta_t k_t k_t^\top\) removes whatever the state currently associates with \(k_t\), and \(\beta_t k_t v_t^\top\) writes the new value in its place. Equivalently, it is exactly one step of online gradient descent with step size \(\beta_t\) on the regression loss \(\tfrac{1}{2}\lVert S^\top k_t - v_t \rVert^2\) — the state is a fast-weight memory being trained at test time. Because the update is rank-one rather than diagonal, it can overwrite a key–value binding instead of merely accumulating, which is precisely what repairs linear attention’s associative-recall failures.

Gated DeltaNet (Yang et al., 2025) composes the two ideas, multiplying in Mamba-2’s scalar decay \(\alpha_t\) as well:

— fast global forgetting (\(\alpha_t\)) and targeted edits (\(\beta_t\)). RWKV-7’s “dynamic state evolution” is a generalization of the same rule with a vector-valued gate. All of these admit the same chunkwise algorithm as the rest of the chapter (the rank-one products accumulate into a WY-style representation within a chunk), and fla ships them as chunk_delta_rule / chunk_gated_delta_rule.

HGRN and Hawk/Griffin¶

Other notable entries in this space include:

• HGRN / HGRN-2 (Qin et al.): Hierarchical Gated Recurrent Network, using input-dependent forget gates similar to LSTM but designed for parallelism.
• Hawk / Griffin (De et al., Google DeepMind, 2024): both are built from the RG-LRU (Real-Gated Linear Recurrent Unit) — a real-valued gated linear recurrence with an input gate and a recurrence gate, deliberately avoiding complex numbers and any discretization step. Hawk is the pure-recurrent stack; Griffin interleaves RG-LRU blocks with local (sliding-window) attention, making it the recurrent-side twin of the SSM/attention hybrids above. Both were trained to multi-billion-parameter scale with quality competitive with transformer baselines and lower decode latency.

Where the Frontier Is Heading¶

As of mid-2026, the field has reached some pragmatic conclusions:

Hybrids win in practice. Pure SSM models have not displaced transformers in production LLMs — but by 2026 hybrids have, at least at the efficiency frontier. Shipped model families like NVIDIA’s Nemotron-H and IBM’s Granite 4.0 are Mamba-2/attention stacks, not pure transformers. The combination of a small fraction of attention layers with SSM/linear-attention layers offers the best tradeoff: the exact-retrieval capability of attention for a small fraction of total compute, with the memory efficiency of SSMs for the bulk of the sequence processing.

The long-context use case is where alternatives shine. For standard (up to 8K token) language modeling, modern transformers with FlashAttention are hard to beat. The advantage of SSM and linear-attention models grows dramatically as context exceeds 32K tokens.

Training dynamics differ. SSMs and RWKV models have different sensitivity to hyperparameters than transformers. The learning rate schedule, initialization, and gradient clipping values that work for transformers often need significant adjustment for recurrent models — the concrete list (depth doubling to match parameters, \(\Delta\) initialization, no weight decay on A_log/D, no RoPE on recurrent layers) is in Building a ~100M Hybrid Yourself above.

Scaling laws are being remeasured. Early (2022-2023) Chinchilla-style scaling laws were derived for transformers. Whether SSMs follow the same laws is an open research question. Preliminary results suggest competitive scaling, but the optimal model size / data tradeoff may differ.

Hardware matters. Modern GPUs are optimized for dense matrix multiplication — the operation that attention and dense linear layers perform. SSM-specific operations (sequential scans, depthwise convolutions) are less naturally mapped to tensor cores. This gives transformers an implementation advantage that partially offsets the algorithmic advantage of SSMs at long sequences. Future ASICs or NPUs designed for recurrent inference could shift this calculus.

Key Takeaways¶

Further Reading¶

• Gu, A., Goel, K., & Ré, C. (2021). Efficiently Modeling Long Sequences with Structured State Spaces (S4). ICLR 2022.
• Gu, A., & Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752.
• Dao, T., & Gu, A. (2024). Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality (Mamba-2). ICML 2024.
• Peng, B., et al. (2023). RWKV: Reinventing RNNs for the Transformer Era. EMNLP 2023 Findings.
• Sun, Y., et al. (2023). Retentive Network: A Successor to Transformer for Large Language Models (RetNet). arXiv:2307.08621.
• Katharopoulos, A., et al. (2020). Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. ICML 2020.
• Yang, S., et al. (2023). Gated Linear Attention Transformers with Hardware-Efficient Training. arXiv:2312.06635.
• Schlag, I., Irie, K., & Schmidhuber, J. (2021). Linear Transformers Are Secretly Fast Weight Programmers. ICML 2021 — the delta rule for linear attention.
• Yang, S., et al. (2024). Parallelizing Linear Transformers with the Delta Rule over Sequence Length (DeltaNet). NeurIPS 2024.
• GitHub: fla-org/flash-linear-attention — Triton kernels and nn.Module layers for GLA, RetNet, DeltaNet, gated DeltaNet, and RWKV-6/7.
• De, S., et al. (Google DeepMind, 2024). Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models. arXiv:2402.19427.
• AI21 Labs (2024). Jamba: A Hybrid Transformer-Mamba Language Model. arXiv:2403.19887.
• GitHub: state-spaces/mamba — official Mamba implementation and selective scan CUDA kernels.
• GitHub: BlinkDL/RWKV-LM — official RWKV training code.

Exercises¶

1. The reassociation trick and why causality breaks it. (a) Linear attention replaces \(\operatorname{softmax}(QK^\top)V\) with \(\phi(Q)\big(\phi(K)^\top V\big)\). Explain in words why computing \(\phi(K)^\top V\) first changes the asymptotic cost, and why this reordering is only valid non-causally — what has to change for a causal (autoregressive) model? (b) Take \(N = 4096\), feature dimension \(r = d_k = 64\), and \(d_v = 64\). Count the scalar multiplications for the quadratic route (materialize \(QK^\top\), then multiply by \(V\)) versus the linear route (\(\phi(K)^\top V\) first, then \(\phi(Q)\) times that). What is the ratio, and how does it depend on \(N\) and \(r\)?

2. Inference memory: KV cache vs. recurrent state at 128K context. Use the chapter’s 7B configuration: 32 layers, \(d = 4096\), \(d_k = d_v = 128\), fp16 (2 bytes/element). For a context of \(N = 131072\) (128K) tokens compute (a) the transformer KV-cache size across all layers, (b) the Mamba state size (state dim \(N_s = 16\)), and © the RWKV state size (scalar state per channel). Then give the transformer-to-Mamba ratio, and state which of the three grows with \(N\).

3. The semiseparable decay mask by hand. The Mamba-2 SSD layer computes \(Y = \big(L \circ (C B^\top)\big) V\) with the 1-semiseparable causal decay mask \(L_{ij} = \prod_{k=j+1}^{i} a_k\) for \(i \ge j\) (and \(0\) for \(i < j\)), where \(a_k = \exp(\Delta_k A)\) are scalar per-step decays. Take a length-4 sequence with \(a_1 = 0.9,\ a_2 = 0.5,\ a_3 = 0.8,\ a_4 = 0.5\) (1-indexed). (a) Write out the full \(4\times4\) matrix \(L\). (b) Show that if every \(a_k\) equals a constant \(\gamma\), then \(L\) reduces to RetNet’s decay mask \(D_{mn} = \gamma^{\,m-n}\).

4. Implement the recurrent form of RetNet retention. The chapter gives RetNet’s parallel retention \(Y = \big(QK^\top \odot D\big)V\) with \(D_{mn}=\gamma^{\,m-n}\) (for \(m \ge n\)). Implement its \(O(1)\)-per-step recurrent form using a running state matrix \(S_t \in \mathbb{R}^{d_k \times d_v}\), and verify numerically that it reproduces the parallel form. State the state-update and readout equations you are implementing.

import torch

def retention_parallel(Q, K, V, gamma: float):
    # Q, K: (T, d_k)   V: (T, d_v)   gamma in (0, 1)
    T = Q.shape[0]
    m = torch.arange(T)[:, None]
    n = torch.arange(T)[None, :]
    exp = (m - n).clamp(min=0).float()               # exponent m-n, masked below
    D = torch.where(m >= n, gamma ** exp, torch.zeros(()))
    return (Q @ K.T * D) @ V                          # (T, d_v)

def retention_recurrent(Q, K, V, gamma: float):
    T, d_k = Q.shape
    d_v = V.shape[-1]
    S = torch.zeros(d_k, d_v)                         # running state
    outs = []
    for t in range(T):
        S = gamma * S + torch.outer(K[t], V[t])       # S_t = gamma S_{t-1} + k_t v_t^T
        outs.append(Q[t] @ S)                         # y_t = q_t^T S_t
    return torch.stack(outs, dim=0)                   # (T, d_v)

if __name__ == "__main__":
    torch.manual_seed(0)
    T, d_k, d_v, gamma = 12, 8, 5, 0.9
    Q, K = torch.randn(T, d_k), torch.randn(T, d_k)
    V = torch.randn(T, d_v)
    par = retention_parallel(Q, K, V, gamma)
    rec = retention_recurrent(Q, K, V, gamma)
    assert torch.allclose(par, rec, atol=1e-5)
    print("recurrent retention matches parallel form: OK")

5. Verify the State Space Duality: SSD layer == selective scan. For a single channel, the selective SSM recurrence is \(h_t = a_t\,h_{t-1} + B_t\,x_t\) with readout \(y_t = C_t^\top h_t\), where \(a_t\) is a scalar decay, \(B_t, C_t \in \mathbb{R}^{N}\), and \(x_t \in \mathbb{R}\). The Mamba-2 SSD claim is that this equals \(Y_i = \sum_{j \le i} L_{ij}\,(C_i^\top B_j)\,x_j\) with \(L_{ij}=\prod_{k=j+1}^{i} a_k\). Implement both and confirm they agree. Why is the matmul (SSD) form preferred on GPUs?

import torch

def selective_scan(a, B, C, x):
    # a: (T,)   B, C: (T, N)   x: (T,)   -- single channel
    T, N = B.shape
    h = torch.zeros(N)
    ys = []
    for t in range(T):
        h = a[t] * h + B[t] * x[t]            # h_t = a_t h_{t-1} + B_t x_t
        ys.append((C[t] * h).sum())           # y_t = C_t . h_t
    return torch.stack(ys)                    # (T,)

def ssd_matmul(a, B, C, x):
    T, N = B.shape
    Acum = torch.cumsum(torch.log(a), dim=0)              # Acum[i] = sum_{k<=i} log a_k
    i = torch.arange(T)[:, None]; j = torch.arange(T)[None, :]
    # L[i,j] = prod_{k=j+1}^{i} a_k = exp(Acum[i] - Acum[j]) for i >= j, else 0
    L = torch.where(i >= j, torch.exp(Acum[:, None] - Acum[None, :]), torch.zeros(()))
    Score = C @ B.T                                      # (T, T), entry C_i . B_j
    return ((L * Score) * x[None, :]).sum(dim=-1)        # sum_j L_ij Score_ij x_j

if __name__ == "__main__":
    torch.manual_seed(0)
    T, N = 16, 4
    a = torch.rand(T) * 0.5 + 0.4          # decays in (0.4, 0.9), so log a is finite
    B, C = torch.randn(T, N), torch.randn(T, N)
    x = torch.randn(T)
    assert torch.allclose(selective_scan(a, B, C, x), ssd_matmul(a, B, C, x), atol=1e-5)
    print("SSD matmul matches selective scan: OK")