Generated by Codex with GPT 5.6 Sol High

State-space models promise a compact alternative to the transformer’s ever-growing key-value cache. Yet compact does not automatically mean fast. At production batch sizes, repeatedly reading and rewriting an SSM’s fixed-size recurrent state can consume so much memory bandwidth that the supposedly efficient layers become the decoding bottleneck.

The NVIDIA Technical Blog published “ReplaySSM: Rethinking State Storage for Faster Hybrid-Model Decoding”, describing joint work with Princeton’s Dao AI Lab that changes what an inference server stores. Instead of materializing the latest recurrent state after every token, ReplaySSM keeps an older checkpoint plus a short buffer of the inputs needed to reconstruct subsequent state updates. That one representation separates producing an output from committing the next state, unlocking improvements in ordinary decoding, speculative decoding, and quantization.

Why a fixed-size state can still be expensive

Hybrid language models combine many SSM layers with a smaller number of attention layers. The attention cache grows with context length, while each SSM layer carries a fixed-size matrix that summarizes all previous tokens. On paper, this makes SSMs attractive for long contexts. On a GPU, however, every decoding step traditionally loads that entire matrix from high-bandwidth memory, applies a small update, uses the result to produce an output, and writes the full matrix back.

This work performs little arithmetic relative to the amount of data moved, so it is limited by memory bandwidth. In NVIDIA’s example, one Mamba-2 layer in Nemotron 3 uses a 4 MiB FP32 state per sequence. The model also has five times as many SSM layers as attention layers. At batch size 256 on a Blackwell Ultra B300, those SSM layers remain the largest part of decode latency through contexts of 100,000 tokens.

The conventional representation also frustrates two common inference optimizations. Speculative decoding asks a smaller mechanism to draft several tokens, then accepts only a prefix after the target model verifies them. An attention server can roll back by moving a pointer in its append-only KV cache. An SSM has already folded every draft into one recurrent state, so serving engines preserve a separate state snapshot after each draft token. Those copies add traffic and consume memory that could otherwise support more concurrent requests.

Low-precision state has a different failure mode. Quantization error does not merely affect one read: the rounded state becomes the input to the next recurrent update, allowing error to compound across a long generation. This is why models whose weights and KV caches use four- or eight-bit formats may still retain their SSM state in FP32.

Cache the ingredients, not every result

ReplaySSM observes that the recurrence can be unrolled over a short window. The current state can be represented exactly, aside from normal floating-point differences, as an older checkpoint plus the sequence of recent decay, key, and value inputs that would update it. Most decoding steps therefore append a small input record to a buffer and compute the token output from the checkpoint and buffered records. Only when the buffer fills does a flush fold all records into a new checkpoint and write the large state back to memory.

With a 16-token window, the checkpoint is still read whenever an output needs it, but it is written once per window rather than once per token. Small input vectors replace most full-matrix writes. Across four Mamba-2 and Gated DeltaNet models, from dense 4B models on an H100 to 120B-class mixture-of-experts models on a B300, this cut enough memory traffic to make the SSM kernel 1.46x to 1.80x faster at batch size 256. Full decode steps improved by 1.20x to 1.47x. The technique did not help at batch size one, where bandwidth pressure was not the dominant cost—an important boundary on the result.

The design is more than a write-combining trick. It makes state commitment an occasional operation rather than an inseparable part of producing every output. That distinction creates independent optimization choices for reads, updates, rollback, and numerical precision.

Speculation becomes pointer movement again

During speculative decoding, ReplaySSM appends the draft tokens’ input records to the buffer before verification. Once the target model decides how many drafts to accept, the server advances the buffer’s committed pointer to that position. Rejected records are ignored and later overwritten. No per-draft state snapshots are created, and a flush incorporates only committed records.

ReplaySSM can also verify a chunk without constructing each intermediate state. Algebraically rearranging the recurrence turns verification into small matrix products over the checkpoint, cached inputs, and draft queries. This “output-only” evaluation avoids unnecessary state materialization and maps well to Tensor Cores.

On GSM8K prompts with four draft tokens, the two evaluated 120B-class models delivered 1.85x and 2.10x the throughput of vLLM’s speculative implementation at batch size 512. Conventional snapshots caused verification latency to rise sharply with draft length and capped concurrency through memory use. ReplaySSM verified eight drafts in roughly the time of one standard decoding step and continued scaling at larger batches.

Separate the state used for reading from the state used for recurrence

Buffering also changes the quantization problem. ReplaySSM maintains a high-precision checkpoint for recurrent updates and a low-precision copy for ordinary output reads. At a flush, it derives both versions from the same updated high-precision state. The low-precision copy is never fed back into the recurrence, so its rounding error affects outputs inside the current window but does not accumulate from token to token.

The long-generation measurements show why this separation matters. On Qwen3.5-4B after 128,000 generated tokens, using one NVFP4 checkpoint for both jobs produced a KL divergence of 0.136 from FP32 decoding. Keeping the recurrence in FP32 while reading an NVFP4 copy reduced that divergence to 0.0038. An FP16 recurrent checkpoint with stochastic rounding tracked FP32 closely as well.

Downstream accuracy exposed the remaining limits. A shared NVFP4 checkpoint lost between 3 and 31 percentage points on GPQA-Diamond across the four models. Separate NVFP4 reads with an FP16 or FP8 recurrent checkpoint kept three models within roughly two points of their FP32 result, but the smallest Nemotron model still lost more than five points. ReplaySSM makes aggressive state compression much safer; it does not make quantization quality automatic.

State representation is part of the serving architecture

ReplaySSM’s broader lesson is that inference performance often depends less on changing the model’s mathematics than on choosing the right materialized state. Conventional decoding stores the latest result because that seems like the natural dynamic-programming representation. Storing a checkpoint and the recent inputs instead exposes delayed commitment, append-only rollback, output-only evaluation, and different precision policies for reads and updates.

The idea is already integrated into TensorRT-LLM, vLLM, SGLang, and FlashInfer, and related forms run in production for GLM-5.3-Flash and Kimi-K3. That adoption matters because the benefits span layers that are usually optimized separately: model architecture, GPU kernels, memory management, speculative serving, and numerical formats. ReplaySSM shows that a carefully chosen state representation can give all of them a better interface at once.