This section is about computing a linear recurrence fast on a GPU, during training, forward and backward. Every layer in this chapter keeps a state and updates it the same way. At each step it multiplies the old state by a decay, a number that says how much of it survives, and adds something new. Run as a loop over $T$ tokens, that is $T$ steps in a row, each waiting for the one before, while thousands of GPU cores wait with it. A layer that cannot train fast on long sequences is never shipped, however cheap it is at decoding, so this is where recurrent layers win or lose against attention.

§21.1 showed the way out in principle. Because each update is linear in the state, any stretch of steps can be summarized by one update of the same form. A parallel scan combines those summaries in about $\log_2T$ rounds instead of $T$. This section turns that idea into what production kernels actually run, for the vector and matrix states that real layers carry.

One hardware fact shapes everything that follows. A GPU has two kinds of arithmetic unit. Its tensor cores do nothing but multiply small matrices, and on an H100 they do it at about 989 TFLOP/s in bf16. Its ordinary float32 units do everything else at about 67 TFLOP/s, nearly 15 times slower (§11.1). A scan is elementwise multiplies and adds, so it runs on the slow units however cleverly it is arranged. An algorithm that does twice the arithmetic as matrix products can therefore win easily. That is what happened in 2024. A chunked algorithm that does more arithmetic than Mamba's carefully fused scan (§21.4) ran 2 to 8 times faster than it at a state size of 64 (Dao & Gu 2024).

We get there in steps. First we extend the scan from one number of state to vectors and matrices, and count what each step of it costs. Then we derive the backward pass, which is another linear recurrence, run backward in time. Then we see why long products of decays break in float32 and how kernels avoid it. Then we derive, prove, and size the chunkwise algorithm, which splits the sequence into short chunks and does most of its work as matrix products. The section ends with when a scan is the wrong tool, and how the same idea splits one sequence across several GPUs.

The scan for one number of state

Start with the smallest case, a state that is one number. At each step the state is the old state times a decay, plus a new input,

Here $a_t$ is the decay, the fraction of the old state that survives step $t$, and $b_t$ is the input written at step $t$. The state starts at $h_0 = 0$ unless we say otherwise. This is the scalar linear recurrence on which §21.1 builds the chapter.

Each step is a rule of the form "multiply by $a_t$, then add $b_t$." A function of this form, $h\mapsto a_th + b_t$, is called an affine map, and the pair $(a_t, b_t)$ describes it completely. The fact that makes everything work is that two such steps in a row are again one such step. Take "multiply by 0.5 and add 1" followed by "multiply by 0.5 and add 0." Starting from any $h$, that gives $0.5(0.5h + 1) + 0 = 0.25h + 0.5$, which is the single step $(0.25, 0.5)$.

In general, applying $(a_1, b_1)$ and then $(a_2, b_2)$ gives $a_2(a_1h + b_1) + b_2 = a_1a_2h + (a_2b_1 + b_2)$. §11.4 wrote this as an operator on pairs,

whose identity is $(1, 0)$, the step that changes nothing.

In plain words, combining two steps multiplies their decays and carries the first step's input through the second step's decay.

The operator is associative, which means $(p_1\bullet p_2)\bullet p_3 = p_1\bullet(p_2\bullet p_3)$ for any three steps. So a long chain of steps can be combined in any grouping, and different processors can combine different parts of it at the same time. That freedom is what allows a parallel computation.

Call one application of $\bullet$ to two pairs a combine. A parallel algorithm is measured by two numbers. Its work is the total number of combines it does. Its depth is the number of rounds that must run one after another. A plain loop has work $T - 1$ and depth $T - 1$. Two parallel scans matter here, and both come back later.

• In the first, every element combines with the element 1 place to its left, then with the one 2 places to its left, then 4, and so on, each round using the values from the round before. After $\lceil\log_2T\rceil$ rounds every position holds its full prefix. The work is about $T\log_2T$ combines. This is the Hillis-Steele scan (Hillis & Steele 1986).
• The second sweeps up a binary tree, combining neighbors in pairs until the root holds the whole sequence. Then it sweeps back down, handing each node the combination of everything to its left. Its depth is $2\log_2T$ and its work is $2(T-1)$ combines, about twice the loop's. This is the Blelloch scan (Blelloch 1990).

Brent's theorem, also from §11.4, turns these counts into time. On $p$ processors, a computation with work $W$ and depth $D$ finishes in at most $W/p + D$ steps. So low depth helps only when there are enough processors to use it, a point that matters at the end of this section.

§11.4 stopped at a scalar state, $h_0 = 0$, the forward pass, and float64 in NumPy. A real layer differs in five ways. It has a vector or matrix state. It starts from a nonzero state whenever a long sequence is processed in pieces. It needs a backward pass. It multiplies thousands of decays in float32. And it runs on a GPU whose tensor cores a scan never touches. Each of these changes the algorithm, and the rest of the section takes them one at a time.

Scans over vector and matrix states

Real layers keep many numbers of state, and in general a step may mix them. The old state vector is multiplied by a matrix, and then a new vector is added. The general linear recurrence has a state $\mathbf{h}_t\in\R^N$ of $N$ numbers, a transition matrix $\mathbf{A}_t\in\R^{N\times N}$ that decays and mixes the old state, and an input $\mathbf{b}_t\in\R^N$,

Both $\mathbf{A}_t$ and $\mathbf{b}_t$ may be computed from the layer's input, but never from the state. That is what keeps the recurrence linear in its state.

Two layers of this chapter show the range. In Mamba (§21.4), a recurrent layer whose decays are computed from each token, every channel keeps $N$ state numbers, each with its own decay. So $\mathbf{A}_t$ is a diagonal matrix with the decays on its diagonal, and $\mathbf{b}_t$ is the token's input written into the state.

In linear attention (§21.2), attention without the softmax rewritten as a recurrence, the state is a $d_k\times d_v$ matrix. Each token adds the outer product $\mathbf{k}_t\mathbf{v}_t^\top$ of its key and value, the matrix whose entry $(i, j)$ is $k_{t,i}v_{t,j}$. With a decay, the update is $\mathbf{S}_t = a_t\mathbf{S}_{t-1} + \mathbf{k}_t\mathbf{v}_t^\top$. Each of the $d_v$ columns of $\mathbf{S}_t$ then follows the vector recurrence above, with the transition $a_t\mathbf{I}$. So one definition covers every layer in the chapter, once the state is allowed to be a matrix.

The formula is the expansion $\mathbf{A}_2(\mathbf{A}_1\mathbf{H} + \mathbf{b}_1) + \mathbf{b}_2 = (\mathbf{A}_2\mathbf{A}_1)\mathbf{H} + (\mathbf{A}_2\mathbf{b}_1 + \mathbf{b}_2)$. The later step's matrix stands on the left in both components. For scalars the order of $a_1a_2$ did not matter, which is why §11.4 could write it either way. For matrices it does.

In plain words, steps can be grouped in any way but never reordered. A parallel scan only regroups, which is why it is allowed. But every implementation must keep the earlier step on the correct side. In the Blelloch down-sweep, for example, the prefix arriving from the left must stay on the left of each combine, which is how §11.4 wrote it.

A nonzero initial state $\mathbf{h}_0$ costs nothing extra, and there are two ways to handle it. Write $(\mathbf{A}_{t:0}, \boldsymbol\beta_t)$ for the scan's output at position $t$, the composed map of steps 1 to $t$. Applied to $\mathbf{h}_0$ it gives $\mathbf{h}_t = \mathbf{A}_{t:0}\mathbf{h}_0 + \boldsymbol\beta_t$, so one option is to apply the prefixes to $\mathbf{h}_0$ at the end.

The other option is to fold $\mathbf{h}_0$ into the first step before scanning. Replace $\mathbf{b}_1$ by $\mathbf{b}_1 + \mathbf{A}_1\mathbf{h}_0$, then scan as if $\mathbf{h}_0 = \mathbf{0}$. For a scalar with $h_0 = 2$ and first step $(0.5, 1)$, the folded input is $1 + 0.5\cdot2 = 2$, which is indeed $h_1$. Kernels use this second way. It is what lets a long sequence be processed in pieces, each piece starting from the state where the previous one ended.

What one combine costs

The work of a scan is the number of combines times the cost of one combine. The cost of one depends entirely on the structure of $\mathbf{A}_t$.

This is why no practical model uses a dense transition computed from the input. With $N = 16$, a dense combine costs about 8,700 FLOPs against 48 for a diagonal one, 180 times more. The layer would also have to produce $N^2$ numbers per token and channel just to build $\mathbf{A}_t$. Little is lost by avoiding it. §21.1 showed that a dense $\mathbf{A}$ that stays the same at every step can almost always be diagonalized over the complex numbers, which turns it into independent scalar recurrences that compute the same outputs.

Layers whose transitions change with the input use one of three structures instead.

• A diagonal matrix, one decay per state entry. Mamba uses it, and so do gated linear attention (GLA), linear attention with one decay per key dimension (Yang et al. 2023), and the recurrent layer of Griffin (De et al. 2024).
• A scalar times the identity, one decay for a whole head. Mamba-2 uses it (Dao & Gu 2024). RetNet uses the same structure with a fixed decay for each head, one that does not depend on the input (Sun et al. 2023).
• The identity minus a rank-one matrix, used by the delta-rule layers of §21.6, which overwrite what the state stores under a key instead of only adding to it. These are harder to compose, because a product of $k$ such matrices is the identity minus a matrix of rank up to $k$. §21.6 keeps those products cheap by storing them as two thin factors, the WY representation.

The last sentence of the box, one stored state per token, is the real obstacle for matrix states.

Keep that number in mind, because the chunkwise algorithm later in this section exists to avoid it. An interviewer who asks how to parallelize a recurrence over time wants both halves of the answer. The combine is cheap only for structured transitions, and a scan stores one state per token.

The cell checks associativity and non-commutativity for dense $3\times3$ steps. Then it runs the Hillis-Steele scan with matrix-affine elements and compares it with a loop. The only change from the scalar scan is the order of the matrix product.

import numpy as np

def compose(p1, p2):
    raise NotImplementedError

def compose_all(maps):
    raise NotImplementedError

A vectorized scan for diagonal states

The common case is a diagonal transition, one decay per state entry. Write $\mathbf{a}_t\in\R^N$ for the diagonal and $\odot$ for the elementwise product. The recurrence is $\mathbf{h}_t = \mathbf{a}_t\odot\mathbf{h}_{t-1} + \mathbf{b}_t$, and the combine is

That is $N$ independent scalar scans side by side. So one Hillis-Steele round is three whole-array operations on arrays of shape $(T, N)$, and NumPy (or a GPU) does every channel and every position of the round at once.

The Blelloch scan works for these pairs too. It does $2(T-1)$ combines on a length padded to a power of two, instead of about $T\log_2T$. The cell implements it vectorized over channels and counts the combines. In the down-sweep the prefix is always the first argument of combine, so the non-commutative operator is applied in the right order.

Blelloch did 3,046 combines here. That is 1,023 on the way up, 1,023 on the way down, and 1,000 more to turn the exclusive scan into the inclusive one. Hillis-Steele did 8,977 and the loop 999. So the work-efficient scan costs about three times the loop's work, the price of depth 20 instead of 999.

import numpy as np

def parallel_linear_scan(a, b, h0=None):
    a, b = np.asarray(a, dtype=float), np.asarray(b, dtype=float)
    return np.cumsum(b, axis=0), 0      # placeholder: ignores every decay, and is not a scan of the pairs

The backward pass is a scan in reverse

Training needs the gradient of the loss with respect to every decay $a_t$, every input $b_t$, and the initial state. The loss depends on all $T$ states, and each parameter affects every later state. Tracing each parameter's effect on each later state separately would cost $O(T^2)$. Reverse-mode differentiation (§0.3) carries one running gradient backward through the computation instead, and does it in $O(T)$. Here the running gradient obeys a recurrence of exactly the kind we already know how to scan.

Work with one scalar channel, since a diagonal state is $N$ of them side by side. Let the loss $\mathcal{L}$ depend on the states $h_1, \dots, h_T$. Two kinds of derivative matter.

• The direct gradient $g_t = \partial\mathcal{L}/\partial h_t$ holds the other states fixed. It is what the layer's output reports back. For an output $y_t = c_th_t$, it is $g_t = c_t\,\partial\mathcal{L}/\partial y_t$.
• The adjoint $\bar h_t = d\mathcal{L}/dh_t$ is the total derivative. It counts every path from $h_t$ to the loss, including the paths through later states. It is the running gradient of §0.3.

In this section an overbar always means an adjoint, never a discretized parameter as in §21.3.

In plain words, the gradient flows backward through the same chain of decays, picking up each step's direct gradient on the way. Reverse time and shift the decays by one place, since the gradient passes from step $t+1$ back to step $t$ through $a_{t+1}$. Then the adjoint recursion is a linear recurrence of exactly the forward kind, with inputs $g_t$, and the same scan kernel computes it on reversed arrays. The decay gradient then needs one elementwise product with the previous states.

The theorem keeps its shape for larger states. For a diagonal state every quantity is a vector and every product is elementwise. For a dense transition, $\bar{\mathbf{h}}_t = \mathbf{g}_t + \mathbf{A}_{t+1}^\top\bar{\mathbf{h}}_{t+1}$ and $\partial\mathcal{L}/\partial\mathbf{A}_t = \bar{\mathbf{h}}_t\mathbf{h}_{t-1}^\top$, by the matrix calculus of §0.2.

The cell runs the four-step example, then checks all three gradients against finite differences on a random recurrence of 50 steps.

That leaves memory. The decay gradient $\bar h_th_{t-1}$ needs the previous state at every step. Storing all of them costs $TN$ numbers for every channel, the very tensor that Mamba's fused kernel (§21.4) refuses to write to HBM. So Mamba recomputes the states from the inputs during the backward pass. That is the same trade FlashAttention makes when it recomputes its attention matrix instead of storing it (§12.3). Chunked kernels store only the state at each chunk boundary and recompute the states inside a chunk from it. An interviewer who asks how recurrent layers train on long sequences without storing a state per token wants this answer.

import numpy as np

def linear_recurrence_backward(a, b, h0, g):
    g = np.asarray(g, dtype=float)
    return np.zeros_like(g), g.copy(), np.zeros(g.shape[1])     # placeholder: db = g ignores the path through later states

The backward pass of linear attention

The matrix state of linear attention is worth doing in full, because the chunked kernels later in the section train exactly this layer. The forward pass starts from $\mathbf{S}_0 = \mathbf{0}$ and updates $\mathbf{S}_t = a_t\mathbf{S}_{t-1} + \mathbf{k}_t\mathbf{v}_t^\top$. It then reads the output $\mathbf{o}_t = \mathbf{S}_t^\top\mathbf{q}_t$ with the query $\mathbf{q}_t\in\R^{d_k}$. Write $\bar{\mathbf{o}}_t = \partial\mathcal{L}/\partial\mathbf{o}_t$ for the gradient arriving at each output.

The read $\mathbf{o}_t = \mathbf{S}_t^\top\mathbf{q}_t$ gives the state the direct gradient $\mathbf{q}_t\bar{\mathbf{o}}_t^\top$. So the theorem gives the reverse recurrence

and each step has four gradients,

where $\langle\cdot,\cdot\rangle$ is the sum of elementwise products.

In plain words, the query gradient is the odd one out. A query is used once, in the read at its own step, so its gradient needs the forward state $\mathbf{S}_t$ and is computed in a forward sweep. A key, a value, or a decay is written into the state and reaches every read from its own step on. So its gradient goes through the adjoint state $\bar{\mathbf{S}}_t$, in a backward sweep. Each sweep holds one $d_k\times d_v$ state at a time.

The decay gradient $\bar a_t = \langle\bar{\mathbf{S}}_t, \mathbf{S}_{t-1}\rangle$ seems to need every forward state $\mathbf{S}_{t-1}$ during the reverse sweep. Those are the $T$ matrices that no kernel wants to write to HBM. In fact it needs none of them.

In plain words, the loss is a sum of contributions, one for each pair of a write at step $s$ and a later read at step $t\ge s$. Raising one log decay $\ell_u$ scales up exactly the contributions whose path crosses step $u$, those with $s < u\le t$. The number $\mathbf{q}_t^\top\bar{\mathbf{q}}_t$ is the sum of all contributions to the read at step $t$. The number $\mathbf{k}_s^\top\bar{\mathbf{k}}_s$ is the sum of all contributions from the write at step $s$. Summing the reads from $u$ on and subtracting the writes from $u$ on leaves exactly the paths that cross step $u$. A token's read of its own write appears in both sums and cancels, as it must, since that path is never decayed.

So the decay gradient costs about $4Td_k$ more operations on numbers we already have, and no states at all. Section 4.3 of Yang et al. (2023) gives this closed form for gated linear attention, which has a vector of decays per step, and notes that earlier work had assumed the states must be stored to compute it. Exercise 21.5f builds the whole backward pass this way.

import numpy as np

def linear_attention_backward(Q, K, V, log_a, dO):
    Q, K, V, dO = (np.asarray(x, dtype=float) for x in (Q, K, V, dO))
    dV = np.stack([np.outer(Q[t], dO[t]).T @ K[t] for t in range(len(Q))])
    return np.zeros_like(Q), np.zeros_like(K), dV, np.zeros(len(Q))   # placeholder: dV sees only the read at its own step

Long products without underflow

Every algorithm in this chapter needs segment products $a_{t:s} = a_t\cdots a_{s+1}$. They fill the mixing matrix of §21.1, and below they fill every chunk. There are two tempting formulas. The first divides cumulative products, $a_{t:s} = P_t/P_s$ with $P_t = a_1\cdots a_t$ from np.cumprod. The second is its log-space twin, $\exp(S_t - S_s)$ with the cumulative log $S_t = \sum_{r\le t}\log a_r$. §11.4 warned against the ratio, and §21.1 against both. Here is what actually happens in float32.

The quotient is 1.0, a confident claim that a hundred decays of 0.9 left the state untouched. Here is why. The smallest normal float32 is $1.18\times10^{-38}$. Below it lie the subnormal numbers, which keep fewer and fewer significant bits as they shrink. The product passes into them at step 829. From step 964 on, multiplying by 0.9 rounds back to the same value, $5.6\times10^{-45}$, so $P_{1100} = P_{1000}$ exactly.

Hardware makes it worse. Many GPU settings flush subnormals to zero for speed, and then the same quotient is $0/0$, a NaN. A stronger decay of $e^{-0.5}$ per step makes the product stick after only 206 steps. And a reset, $a_t = 0$, makes every later $P_t$ zero at once.

Log space removes the underflow but not the cancellation, the loss of accuracy when two nearly equal large numbers are subtracted. The spacing of float32 at a number is the gap to its nearest neighbor, and it grows with the number's size. After 100,000 steps the cumulative log is about $-10{,}600$. Neighboring float32 numbers there are $2^{-10}\approx0.001$ apart, so each of the 20 additions between $S_s$ and $S_t$ rounds to a multiple of about 0.001. The difference $S_t - S_s$ is off by a few thousandths.

Since $e^{x+\epsilon}\approx e^x(1+\epsilon)$, the segment product is off by the same amount in relative terms, $3\times10^{-3}$ in this run. Summing only the segment's own 20 logs is accurate to float32's precision, $10^{-7}$ or better. This matters most for decays near 1, which are where models keep their long memories, and §21.1 showed how sensitive those are.

There are three safe ways to get segment products, and all three appear in real kernels.

Let the scan carry them. The scan's first component is the product of the decays it has composed, built by multiplications only. It may underflow to zero, but then the true value is also below $10^{-38}$, so nothing is lost. No division ever happens.

Sum logs within a short window. Inside a chunk of $T_c$ tokens, a cumulative sum of logs never exceeds $T_c\max_t|\log a_t|$ in magnitude, so its spacing stays small. A difference of two such sums is accurate as long as no decay in the chunk is exactly 0. With $T_c = 64$ and decays above 0.5, the sums stay above $-45$, where the float32 spacing is $4\times10^{-6}$. This is how the chunked kernels of flash-linear-attention, the open-source library of linear-attention kernels that grew out of the GLA paper (fla-org/flash-linear-attention), compute their decays.

Build each segment from its own terms with additions only. Put $\log a_t$ in row $t$ below the diagonal of a $T_c\times T_c$ matrix, with zeros elsewhere, and take cumulative sums down each column. Entry $(t, s)$ is then $\sum_{r=s+1}^{t}\log a_r$, summed from its own terms with no subtraction anywhere. A reset ($\log a_t = -\infty$) gives $-\infty$, hence a product of 0, instead of the NaN of $-\infty - (-\infty)$. This is the masked cumulative product of Exercise 21.1b, done with sums of logs, and it is the segsum function of Mamba-2 (Dao & Gu 2024). Tri Dao's blog post on the algorithm reports that without it the basic algorithm produced NaNs immediately in training, even in float32.

An interviewer who asks how you would compute the decay mask of a gated layer is checking for exactly this. The right answer never divides, never subtracts long cumulative sums, and keeps the decays and their logs in float32.

The chunkwise algorithm

§21.1 gave three ways to compute a linear recurrence. None of them is good enough for training a layer with a matrix state.

• The recurrent mode runs the loop. Its depth is $T$, and its work is elementwise, so it never touches the tensor cores.
• The matrix mode builds the $T\times T$ mixing matrix and multiplies by it. It is all matrix products, but it costs $O(T^2)$.
• The scan mode has depth $\log T$, but it stores one $d_k\times d_v$ state per token, 137 GB in the example above.

The chunkwise algorithm takes the best of the first two. Split the sequence into short chunks. Inside each chunk, use the matrix mode, a small masked attention computed with matrix products. Between chunks, use the recurrent mode, carrying one state from each chunk to the next. Only that short chain of chunk states is sequential. RetNet (Sun et al. 2023), gated linear attention and its kernel library flash-linear-attention (Yang et al. 2023), Lightning Attention-2 (Qin et al. 2024), and Mamba-2 (Dao & Gu 2024) all train this way.

We derive it for a matrix state with one scalar decay per step,

Here the queries and keys are $\mathbf{q}_t, \mathbf{k}_t\in\R^{d_k}$, the values are $\mathbf{v}_t\in\R^{d_v}$, the decays are $a_t\in[0, 1]$, and $\mathbf{S}_0$ is the initial state. The state $\mathbf{S}_t\in\R^{d_k\times d_v}$ has rows indexed by keys and columns by values, the convention of §21.2. Stack the tokens as rows, $\mathbf{Q}, \mathbf{K}\in\R^{T\times d_k}$ and $\mathbf{V}, \mathbf{O}\in\R^{T\times d_v}$.

This one recurrence is several layers at once.

• With $a_t = 1$ it is the unnormalized causal linear attention of §21.2.
• With a fixed decay for each head it is RetNet's retention.
• With $\mathbf{q}\to\mathbf{C}_t$, $\mathbf{k}\to\mathbf{B}_t$, $\mathbf{v}\to\Delta_t\mathbf{x}_t$, and $a_t = e^{\Delta_tA}$ it is one head of Mamba-2, as §21.6 will show. Here $\mathbf{B}_t$ and $\mathbf{C}_t$ are the input-dependent write and read vectors of §21.4, $\Delta_t$ is the step size, $A$ is the head's negative scalar rate, and $\mathbf{x}_t$ is the head's input.

The GLA paper writes tokens as row vectors, $\mathbf{S}_t = \mathbf{G}_t\odot\mathbf{S}_{t-1} + \mathbf{k}_t^\top\mathbf{v}_t$ with $\mathbf{o}_t = \mathbf{q}_t\mathbf{S}_t$, where $\mathbf{G}_t$ holds its decays. That is the same layout as ours, so its formulas can be read directly.

Before the general formulas, here is the whole algorithm on four tokens.

Every number in that example came from one of three decays, all segment products. The mask holds the decays between tokens of the same chunk. The decay from the chunk start says how much of the incoming state each token sees. The decay to the chunk end says how much of each write survives into the outgoing state. The definition below names them.

Split the sequence into $n_c = \lceil T/T_c\rceil$ chunks of length $T_c$, the last possibly shorter. Chunk $j$ covers the steps $t_0 + 1, \dots, t_0 + T_c$ with $t_0 = (j-1)T_c$. Its rows of $\mathbf{Q}$, $\mathbf{K}$, $\mathbf{V}$, $\mathbf{O}$ are $\mathbf{Q}_{[j]}$, $\mathbf{K}_{[j]}$, $\mathbf{V}_{[j]}$, $\mathbf{O}_{[j]}$. The state after it is $\mathbf{S}_{[j]} = \mathbf{S}_{t_0 + T_c}$, with $\mathbf{S}_{[0]} = \mathbf{S}_0$. Inside a chunk we index by the local position $r = 1, \dots, T_c$.

Read the three decays carefully, because each one hides a classic off-by-one error.

• The mask's diagonal is $a_{t:t} = 1$. A token's own write is not decayed before it is read, because this chapter updates the state first and then reads it.
• The decay from the chunk start includes the first token's decay, $(\boldsymbol\lambda)_1 = a_{t_0+1}$. Step $t_0 + 1$ decays the state it inherits before adding its write.
• The decay to the chunk end excludes the writing token's own decay, so $(\boldsymbol\gamma)_{T_c} = 1$.

In plain words, each output is what its own chunk wrote, read through a small masked attention matrix, plus what all earlier chunks wrote, read from the one state that summarizes them. The state is updated once per chunk by a single matrix product, and only that update is sequential.

Two limits make the trade-off visible. With $T_c = 1$ every chunk is one token, the mask is the scalar 1, and the algorithm is the recurrence. With $T_c = T$ there is one chunk, the inter-chunk term vanishes when $\mathbf{S}_0 = \mathbf{0}$, and the algorithm is the quadratic matrix mode. In between, the chunk length decides how much work is done as dense matrix products and how many sequential steps remain.

The picture behind this is the sequence-mixing matrix of §21.1, the $T\times T$ lower-triangular matrix $\mathbf{M}$ with $\mathbf{O} = \mathbf{M}\mathbf{V}$. Here its entries are $\mathbf{M}_{ts} = a_{t:s}\,\mathbf{q}_t^\top\mathbf{k}_s$ for $s\le t$. The algorithm cuts $\mathbf{M}$ into $T_c\times T_c$ blocks. The blocks on the diagonal are computed densely. Every block below them has rank at most $d_k$, and it is never formed, only applied through the states.

The widget draws this decomposition on a mixing matrix with $T = 64$ and the same decay $a = 0.93$ at every step, shaded by the decay part $a^{t-s}$ of each entry. Each shaded triangle on the diagonal is one masked matrix product, and everything below the band is reached only through states. Click an entry to see how its decay factors into the three pieces of the definition. Press Play to run the sequential pass one chunk at a time.

import numpy as np

def chunked_gla(Q, K, V, log_a, T_c, S0=None):
    # placeholder: plain causal linear attention, every decay treated as 1
    T = Q.shape[0]
    O = ((Q @ K.T) * np.tril(np.ones((T, T)))) @ V
    return O, K.T @ V

How big should a chunk be

Longer chunks put more of the work into matrix products and leave fewer sequential steps. But the work inside a chunk grows with the square of its length. So there is a balance, and it can be computed. Count the matrix-product FLOPs of one chunk, where a product of an $m\times k$ and a $k\times n$ matrix costs $2mkn$.

• The intra-chunk scores $\mathbf{Q}_{[j]}\mathbf{K}_{[j]}^\top$ cost $2T_c^2d_k$, and their product with $\mathbf{V}_{[j]}$ costs $2T_c^2d_v$.
• The inter-chunk read $\mathbf{Q}_{[j]}\mathbf{S}_{[j-1]}$ costs $2T_cd_kd_v$, and the state update $\mathbf{K}_{[j]}^\top\operatorname{diag}(\boldsymbol\gamma)\mathbf{V}_{[j]}$ costs another $2T_cd_kd_v$.

In plain words, lengthening the chunks moves work from the cheap linear part into the quadratic part and removes sequential steps. Below $T_c^\ast$ the quadratic part is the smaller one, so the total stays under twice the unavoidable linear part.

The balance point is the same as the crossover of §21.2, the length at which the two orders of multiplying out linear attention cost the same, $2d_\phi d_v/(d_\phi + d_v)$, with $d_\phi$ the width of the queries and keys (our $d_k$). That is no coincidence. A chunk is a short sequence, and for short sequences the quadratic order is the cheaper one.

Three hardware facts pin the choice down.

• Tile size. Tensor cores multiply fixed small tiles, $16\times8\times16$ for the bf16 instruction of §11.5. Chunks shorter than 16 leave the tiles mostly empty or put the products back on the slow units. The GLA paper sets its chunk length to a multiple of 16 for this reason.
• Shared memory. A chunk's $\mathbf{Q}$, $\mathbf{K}$, $\mathbf{V}$, its $T_c\times T_c$ scores, and the float32 state must fit in the fast on-chip memory of one streaming multiprocessor (SM), one of the GPU's 132 independent processors, about 228 KB on an H100. With $T_c = 64$ and $d = 128$, that is 48 KB of bf16 inputs, 16 KB of float32 scores, and 64 KB of state, which fits. At $T_c = 256$ the scores alone would take 256 KB.
• Latency. Each of the $n_c$ sequential steps costs a fixed delay, so very short chunks lose to synchronization even when their FLOPs are low.

In practice, flash-linear-attention's chunked kernels default to $T_c = 64$. Mamba-2's default is 256, with a state of $128\times64$ per head, above its balance point of 85. It trades extra matrix-product FLOPs for fewer sequential steps.

The cell sweeps $T_c$ on a small head in NumPy. Every chunk length gives the recurrence's outputs to rounding, and the FLOPs follow the theorem.

At $T_c = 1024$ the algorithm does 29 times the FLOPs of $T_c = 4$. In NumPy the fastest chunk is usually 16 or 64. There, Python's overhead per chunk, paid once per sequential step, balances the quadratic FLOPs, near the theorem's $T_c^\ast = 32$. On a GPU the overhead per step is smaller and the matrix products far faster, but the curve has the same shape.

import numpy as np

def chunk_cost(T, T_c, d_k, d_v):
    n_c = T // T_c
    return n_c * T_c * T_c * (d_k + d_v), n_c, n_c * d_k * d_v     # placeholder: misses factors and the ceiling

def step_time(T, T_c, d_k, d_v, heads, peak_matmul, peak_scalar, latency):
    raise NotImplementedError

At real scale

Now one real layer. Qwen3.5-397B-A17B, released in February 2026, has 60 layers, and 45 of them are Gated DeltaNet layers. Each has 64 value heads whose states are $128\times128$ (model card).

Gated DeltaNet is a linear-attention layer with a scalar decay per head and a write that corrects what the state stores under a key instead of only adding to it, the delta rule of §21.6. The delta rule needs an extra step inside each chunk to build the WY factors of that section, which adds work rather than replacing any. So here we count only the gated-linear-attention part of the work. We also treat each value head as its own head. The real layer shares queries and keys across groups of four value heads, which saves part of the intra-chunk work.

Dao and Gu (2024) measured the outcome of this arithmetic. Their chunked kernel for Mamba-2 is 2 to 8 times faster than Mamba's fused scan at state size 64. It allows states 8 times larger with little slowdown. And it overtakes FlashAttention-2 from 2K tokens on, reaching 6 times faster at 16K.

When a scan is the wrong tool

A parallel scan buys depth. Depth matters only when there is not enough other parallel work to fill the machine, and on a GPU there usually is. Brent's theorem counts all the work, including the independent channels and the batch.

Take one Mamba layer at the size of Mamba-2.8B, with $d_{\text{in}} = 5120$ channels of $N = 16$ state entries, on a batch of 8 sequences. That is $8\cdot5120\cdot16 = 655{,}360$ independent scalar recurrences. An H100 has $132\cdot128 = 16{,}896$ float32 lanes, the arithmetic units that each run one thread's instruction. Give each lane a few dozen recurrences to run one after another, and the GPU is already full, with depth $T$ and no extra work at all. A scan over time pays off when the batch is small, or inside one thread block, where a few hundred threads share one chunk of the sequence.

So what limits a recurrent layer is memory traffic, not depth, and the winning kernels are designed around bytes. Three designs show it.

Mamba keeps the scan in fast memory. Mamba's kernel loads $\Delta$, $\mathbf{A}$, $\mathbf{B}$, $\mathbf{C}$, and $x$ for a block of the sequence into the on-chip shared memory (SRAM). It discretizes there, runs a work-efficient parallel scan there, and writes only the outputs, so the state tensor of shape (batch, $T$, $d_{\text{in}}$, $N$) never reaches HBM. The backward pass recomputes the states instead of storing them. The authors report the fused scan 20 to 40 times faster than a standard scan in PyTorch (Sections 3.3 and 4.5 of Gu & Dao 2023).

Griffin keeps a sequential loop in fast memory. De et al. (2024) note that the associative scan does nothing about memory traffic, their real bottleneck, and they measured it significantly slower on TPU-v3 than JAX's native sequential scan. They wrote a sequential scan in Pallas, JAX's language for custom kernels, that keeps the state in on-chip memory throughout. It runs almost 3 times faster than JAX's native sequential scan.

Mamba-2 and linear attention use chunks. When the state is a matrix and the transition is a scalar, or a diagonal over the keys, the chunkwise algorithm turns most of the work into tensor-core matrix products and writes one state per chunk. Mamba-1's decay for channel $d$ and state entry $n$ is $e^{\Delta_{t,d}A_{d,n}}$, with a step size $\Delta_{t,d}$ per channel and a learned rate $A_{d,n}$ per channel and state entry. As a matrix over channels and state entries it is full rank, and such a decay has no matrix-product form (Yang et al. 2023 discuss this). Making the decay one scalar per head is exactly the restriction Mamba-2 accepted to get one (Dao & Gu 2024).

The widget races these strategies under a simple roofline model. Each strategy's time is the larger of its compute time and its memory time, plus a fixed latency for every dependent step, 0.1 µs on chip and 5 µs for each round of a scan. The matrix state is the Gated DeltaNet layer of the example above, and the diagonal state one Mamba-2.8B layer, both at 32,768 tokens on an H100. The scans write every state to HBM. Switch between a matrix state and a diagonal one, and watch the winner change. Then raise the number of sequences in the batch, and see that the winner stays the same on both.

Here $N$ is the number of state entries per recurrence ($d_kd_v$ for a matrix state), and a warp is the group of 32 GPU threads that run in lockstep. S5 is a diagonal state-space layer computed with JAX's associative scan (Smith et al. 2022).

Splitting one sequence across devices

The chunkwise form has one more consequence, and it matters for training on a million tokens. Context parallelism splits one sequence across several GPUs, so that no device holds the whole thing. For a recurrent layer, make each chunk a device.

Device $p$ of $P$ holds a contiguous slice of the sequence. It runs the chunkwise algorithm on its slice from a zero initial state. That gives its local outputs, its local final state $\mathbf{S}^{\text{loc}}_p$, and its total decay $\alpha_p$. The pair $(\alpha_p, \mathbf{S}^{\text{loc}}_p)$ is the composed affine map of the device's whole slice. The true state entering device $p$ is the exclusive scan of these pairs under $\bullet$, a scan over only $P$ elements. It is computed by passing one state from device to device, or in $\log_2P$ rounds.

Each device then adds the correction $\operatorname{diag}(\boldsymbol\lambda)\mathbf{Q}\mathbf{S}_{\text{in}}$ to its outputs, the inter-chunk term of the definition. The state-passing test of Exercise 21.5d runs the sequential version of this on two pieces.

Dao and Gu (2024) describe exactly this context parallelism for Mamba-2, with communication linear in the number of devices. The catch is the sequential dependence. Device $p$ cannot finish until the maps of devices $1$ to $p - 1$ have been composed. But the local work comes first and is fully parallel, and only the cheap correction waits. In a hybrid model the attention layers still need ring attention or another context-parallel scheme, and they dominate the communication.

In PyTorch

Core PyTorch has no linear-recurrence scan in its stable API. Production layers call fused Triton or CUDA kernels, from flash-linear-attention or the mamba_ssm package (state-spaces/mamba), that implement the chunked algorithm with these exact steps. Serving stacks call the same kernels (§20.7). The block below writes those steps with plain tensor operations. It builds the segment sums with Mamba-2's addition-only construction, runs a batched, multi-head chunked layer with einsums, and checks it against a loop. Then it checks the reverse-recurrence gradients of this section against autograd.

The decay to the chunk end reuses segsum. Appending a zero log to each chunk and reading the last row gives $\sum_{u=s+1}^{T_c}\log a_u$ for every $s$, again with additions only.

In JAX the scan itself is a library call, jax.lax.associative_scan(fn, elems). It takes any associative fn(earlier, later) over a tuple of arrays. It also runs in reverse (reverse=True), which is the backward pass of this section. S5 is built on it.

Read more

• Blelloch (1990). Prefix Sums and Their Applications. The work-efficient scan, with a classic treatment of first-order linear recurrences as scans over pairs.
• Martin & Cundy (2018). Parallelizing Linear Recurrent Neural Nets Over Sequence Length. The scan applied to linear RNNs, with a CUDA kernel and the observation that gates computed from the input keep training parallel.
• Yang et al. (2023). Gated Linear Attention Transformers with Hardware-Efficient Training. The chunkwise form with data-dependent decays, why a parallel scan must materialize matrix states, the closed-form decay gradient, and the FlashLinearAttention kernels with and without materialized chunk states.
• Qin et al. (2024). Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models. Intra-block and inter-block tiling of linear attention with decay, forward and backward, the kernel behind MiniMax-01.
• Dao & Gu (2024). Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. Section 6 gives the chunked SSD algorithm as a block decomposition, Section 8 the state-passing context parallelism, and the companion blog post on the algorithm explains the stable segment sum.
• De et al. (2024). Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models. Section 4 explains why a memory-bound diagonal recurrence runs faster as a sequential scan in on-chip memory than as an associative scan.
• Heinsen (2023). Efficient Parallelization of a Ubiquitous Sequential Computation. The log-space closed form with two prefix sums, short and worth testing against the numerical limits of this section.
• Smith, Warrington & Linderman (2022). Simplified State Space Layers for Sequence Modeling. S5, a multi-input multi-output diagonal SSM computed with a parallel scan instead of a convolution, the first deep SSM built on the scan.

Watch next

▶

▶

▶