Parallel Scan and Hardware-Aware SSM Kernels
How an associative scan recovers training parallelism after selectivity destroys the convolution, and why the arithmetic-intensity argument means the kernel is the architecture.
The selective state space recurrence looks hopelessly sequential: each state depends on the previous one. It is not, and the reason has been known in parallel computing since the 1980s. Recognising the recurrence as an associative scan is what makes selective models trainable, and implementing that scan without touching high-bandwidth memory is what makes them fast.
The scan
An associative scan computes all prefix results of an associative binary operator in parallel. Prefix sum is the familiar case. The state space recurrence
fits, because composing two consecutive steps yields another step of the same form: applying \((\bar{A}_1, b_1)\) then \((\bar{A}_2, b_2)\) is equivalent to the single operation \((\bar{A}_2\bar{A}_1,\ \bar{A}_2 b_1 + b_2)\). That composition rule is associative, so a Blelloch-style scan computes every \(h_k\) with \(O(L)\) total work and \(O(\log L)\) sequential depth.
For \(L = 4096\) this turns 4,096 dependent steps into about 12 parallel rounds. The sequential bottleneck is gone without any approximation.
Why the kernel matters more than the algorithm
Count the tensors. For batch \(B\), length \(L\), model dimension \(D\) and state dimension \(N\), the hidden state across the sequence is \(B \times L \times D \times N\). With \(B=8\), \(L=4096\), \(D=2048\), \(N=16\) in bf16, that is roughly 4 GB for one layer's states. Writing that to high-bandwidth memory and reading it back for the backward pass is far more expensive than the arithmetic performed on it, so a straightforward implementation is memory-bound and slower than attention.
The hardware-aware kernel avoids materialising it. Parameters are loaded from HBM into SRAM, the discretisation and the scan run inside SRAM, and only the \(B \times L \times D\) output is written back. For the backward pass the states are recomputed from the saved inputs rather than stored, trading extra FLOPs for a large reduction in memory traffic. This is precisely the FlashAttention argument applied to a different operator, and it is why published throughput numbers for these models are numbers about a kernel, not about an architecture.
Chunked formulations
Later work reorganised the computation into a chunked form: run the scan sequentially across chunks of, say, 256 tokens while handling everything inside a chunk with dense matrix multiplication. Within a chunk the operation becomes a matmul, which uses tensor cores; across chunks the sequential dependency is short. This raises arithmetic intensity substantially, because pure scans are elementwise operations that leave tensor cores idle, and it is the basis of the state space duality framing that connects these models back to a form of linear attention.
When it breaks
The advantage is kernel-dependent, not architecture-dependent. Run a selective state space model in a framework without the fused kernel and it will typically be slower than an equivalent transformer at moderate sequence lengths. Reported speedups should always be read as speedups of a particular implementation on a particular device.
Recomputation is not free. Recomputing states in the backward pass costs real FLOPs, and on a compute-bound configuration, meaning large state dimension and short sequence, that trade can go the wrong way. The optimum depends on \(N\), \(L\) and the device's ratio of compute to bandwidth, which is exactly the roofline calculation.
Numerical care is required inside the scan. Products of many \(\bar{A}_k\) terms underflow or overflow easily. Implementations work in log space, or in complex form with a magnitude-phase decomposition, or accumulate in fp32 regardless of the storage dtype. A scan that is correct in fp32 and silently wrong in bf16 is a common and difficult bug, because the model still trains and merely trains worse.
Variable-length batching is awkward. The scan assumes a contiguous sequence. Packing several documents into one batch row, standard practice for transformers, requires resetting the state at document boundaries, which means either masking \(\bar{A}\) to zero at those positions or running separate scans. Implementations that ignore this leak information across documents in a way that is easy to miss and inflates evaluation numbers.
12 flashcards for this concept
Click a card to reveal the answer.