The Fixed State Budget at Inference

A Llama-3-8B server holding a 32,000-token conversation is carrying about 4 GiB of KV cache for that one request. The arithmetic: 32 layers, 8 key/value heads after grouped-query attention, head dimension 128, two tensors (K and V), two bytes each in fp16, which is \(32 \times 8 \times 128 \times 2 \times 2 = 131{,}072\) bytes per token, or 128 KiB. Multiply by 32,000. Now double the context and the number doubles with it. A state space model in the same position carries the same number of bytes at token 32,000 as it did at token 1. That single difference is what linear-time sequence models are actually selling, and it is worth doing the arithmetic before believing the marketing.

The arithmetic on the other side

A selective SSM layer keeps a state matrix \(h \in \mathbb{R}^{D \times N}\), where \(D\) is the number of channels the layer operates on and \(N\) is the state dimension per channel. Take a layer with model width 4096, an expansion factor of 2 (so \(D = 8192\)) and \(N = 16\). That is \(8192 \times 16 = 131{,}072\) floats. In fp32, which is what most selective-scan kernels use for the recurrence because the state accumulates, that is 512 KiB per layer. Across 32 layers: 16 MiB per sequence, plus a short causal convolution buffer of three or four tokens per channel.

Set the two side by side. The transformer crosses 16 MiB of KV cache at around 128 tokens. Everything after that is pure growth, and the SSM's 16 MiB is flat out to a million tokens. The crossover is early enough that for any long-context workload the comparison is not close.

The consequence that matters commercially is batch size. Serving throughput on a memory-bound decoder is roughly the number of sequences you can hold resident times the rate you can stream weights past them. If each sequence costs 4 GiB, an 80 GiB accelerator holding a 16 GiB model fits sixteen of them. If each costs 16 MiB, it fits thousands, and the decode step becomes weight-bound rather than cache-bound. Mamba reported roughly 5x higher inference throughput than a Transformer of comparable size on this basis (Gu and Dao, 2023, Mamba: Linear-Time Sequence Modeling with Selective State Spaces, arXiv:2312.00752).

Constant memory is not constant work

Prefill does not become free. Reading a 32,000-token prompt still means running 32,000 recurrence steps, or the chunked parallel form of them, and the FLOPs are linear in sequence length rather than quadratic but they are not zero. What changes is that prefill produces a fixed-size artefact instead of a growing one.

That artefact is an unusually convenient object. A conversation's entire history compresses to one 16 MiB blob that can be written to a cache, shipped to another host, or reloaded next week, and the cost does not depend on how long the conversation was. Prefix caching for a transformer requires storing and matching a growing KV prefix; for an SSM it is a fixed-size checkpoint keyed on the prefix hash.

The mirror image is the problem. Transformer KV caches are append-only, so you can truncate them, fork them, or roll back a few tokens by dropping entries. An SSM state is overwritten at every step. Rolling back \(k\) tokens means re-running from the last snapshot, which makes tree search, beam search over long branches, and rejection sampling in speculative decoding more awkward than they look. The usual fix is to snapshot the state at branch points, which is cheap, but it has to be designed in rather than assumed.

When it breaks

The state has to fit in fast memory during the scan, not just in HBM. Hardware-aware kernels load the state into SRAM and keep it there across the steps of a chunk, precisely to avoid materialising \(h\) for every position in HBM. Growing \(N\) to buy recall pushes against SRAM capacity, and past that point the kernel spills and the linear-time advantage narrows sharply. The fixed budget is a real budget, enforced by the memory hierarchy.

Per-sequence state multiplies by batch size, and fp32 doubles it. The 16 MiB above is not 16 MiB for the server; at batch 512 it is 8 GiB, which is real but still an order of magnitude below the transformer equivalent. Teams that quote the per-sequence figure and then plan capacity from it get a pleasant surprise rather than a nasty one, but the habit is still wrong.

Constant memory does not imply constant information. The state is a lossy summary, and what it drops is not chosen by the user. What that costs, and on which tasks, is the subject of State Capacity and Associative Recall.