NVIDIA and Princeton’s ReplaySSM speeds up hybrid model decoding by caching inputs, not state

What is ReplaySSM?

ReplaySSM is a new way to run the state-space model (SSM) layers inside hybrid language models during decoding. It is described in the paper “ReplaySSM: Cache SSM Inputs, Not State”, submitted to arXiv on 9 October 2026 by Ze-Wei Liou and Tri Dao of Princeton and five NVIDIA engineers, and in an NVIDIA technical blog post published on 12 October.

Hybrid models mix a few attention layers with many SSM layers, such as Mamba-2, Gated DeltaNet (GDN) or Kimi Delta Attention. The authors name NVIDIA’s Nemotron 3, Alibaba’s Qwen3.5 to Qwen3.8, Moonshot’s Kimi-K3 and Z.ai’s GLM-5.3-Flash as production models built this way. Their appeal is that an SSM layer keeps a fixed-size state instead of a key-value cache that grows with every token.

The paper argues that this fixed state is not as cheap as it looks, and that a small change in what gets stored fixes three problems at once.

Why are SSM layers slow at serving batch sizes?

Conventional decoding loads an SSM layer’s state from GPU memory, applies a decay and a small update, computes the token’s output and writes the whole state back, every step. That is a lot of memory movement for very little arithmetic, so the kernel is limited by memory bandwidth.

The numbers are large. In NVIDIA’s example, one Mamba-2 layer’s FP32 state is 4 MiB per sequence, about the size of an attention layer’s FP8 key-value cache for 8,000 tokens, and Nemotron-3-Super has five times as many SSM layers as attention layers. At batch size 256 on one Blackwell Ultra B300 GPU, the SSM layers are the largest part of each decode step up to roughly 100,000 tokens of context.

The same design causes two further problems. Speculative decoding needs to undo rejected draft tokens, but an SSM state that has absorbed them has no cheap inverse, so servers such as vLLM store a snapshot per draft token. And if the state is stored in low precision, its rounding error is fed back into the next update and accumulates over a long generation.

How does caching inputs instead of state work?

The recurrence can be unrolled: the current state equals an earlier checkpoint, decayed, plus a short sum over the inputs received since. ReplaySSM stores exactly that, a checkpoint and a bounded buffer of recent input records.

On most steps, which the authors call non-flush steps, the layer appends the new inputs to the buffer and computes the output from the checkpoint and the buffer, without writing a state. When the buffer of L records fills, a flush step folds the inputs into the checkpoint, writes it once and clears the buffer. The outputs match conventional decoding up to floating-point differences.

Over a window of L tokens, conventional decoding reads and writes the state L times each. ReplaySSM still reads the checkpoint every step but writes it once, plus a few small input vectors per step, which roughly halves the dominant traffic. The authors use L = 16, the best window for all four models they tested.

How much faster is it?

The team implemented ReplaySSM in vLLM and compared it with vLLM’s own kernels at their fastest launch settings, on four models: Nemotron-3-Nano-4B and Qwen3.5-4B on one H100, and Nemotron-3-Super-120B-A12B and Qwen3.5-122B-A10B with NVFP4 weights on one B300.

At batch size 256 the SSM kernel ran 1.46x to 1.80x faster and the whole decode step 1.20x to 1.47x faster. The end-to-end gain is smaller because attention and the rest of the model are unchanged, and it grows with batch size as the SSM share of the step grows. At batch size 1, where the kernel is limited by latency rather than bandwidth, ReplaySSM finished within a few percent of the baseline, so it is a serving optimisation, not a single-user one.

Why does it help speculative decoding?

Because recent inputs are kept, rolling back rejected drafts is just moving a pointer in the buffer, like rewinding a key-value cache. There is no per-draft snapshot to save or restore. For background on the technique, see our explainer on speculative decoding.

ReplaySSM also verifies drafts with what the authors call output-only evaluation: a few small matrix products over the cached keys, values and draft queries, without building intermediate states. On Nemotron-3-Super at batch 256, vLLM’s verification step with eight drafts took 4.66 times as long as a standard decode step; ReplaySSM verified eight drafts in 0.89 to 1.21 times a standard step, flushes included.

Removing the snapshots also frees GPU memory for more concurrent requests. On GSM8K prompts with four drafts, vLLM’s speculative decoding stopped improving beyond batch 128, and on Qwen3.5-122B-A10B it fell below plain decoding at batch 512. There, ReplaySSM delivered 1.85x and 2.10x the throughput of vLLM’s speculation on the two large models, with matched acceptance lengths.

Can the state be stored in 4-bit precision?

A low-precision state is attractive because it means fewer bytes to move: on Mamba-2, reading an NVFP4 state made the non-flush kernel 3.8x faster than reading FP32. But a single shared 4-bit checkpoint keeps feeding its rounding error back into the recurrence.

ReplaySSM separates the two jobs. A high-precision checkpoint carries the recurrence, and a low-precision copy, written at each flush, serves output reads. The low-precision copy is never read back into the recurrence, so its error affects only the current window.

On Qwen3.5-4B at 128,000 generated tokens, NVFP4 reads with an FP32 recurrence cut the output KL divergence from 0.136 to 0.0038, about 36 times lower. On GPQA-Diamond, a shared NVFP4 checkpoint lost 13.05 points on average across the four models; NVFP4 reads with an FP16 recurrence lost 2.36 points, and three of the four models stayed within about 2 points of FP32. Nemotron-3-Nano-4B still lost 6.7 points.

Where can developers use it?

The code is in a public GitHub repository under the Apache 2.0 licence, as a fork of vLLM with kernels for Mamba-2 and GDN. Its README says upstream vLLM has merged the Mamba-2 standard decoding kernel, while the other decoding paths and the low-precision kernels are so far only in the fork.

NVIDIA’s post lists pull requests for TensorRT-LLM, vLLM and FlashInfer and says ReplaySSM is integrated with SGLang too. It also says GLM-5.3-Flash serves with ReplaySSM, and that Kimi-K3 uses a related, concurrently developed approach for draft verification. Both models already appear in our coverage, from Goodfire’s cyber monitors on Kimi K3 and GLM 5.3 to Z.ai’s GLM-Legal, built on GLM-5.3-Flash.

What are the caveats?

The results come from the authors, several of whom work at NVIDIA, on NVIDIA GPUs, and the paper is a preprint that has not yet been peer reviewed. The gains also depend on workload: they appear at large batch sizes, and the end-to-end improvement is capped by how much of a model’s time its SSM layers take.

The low-precision results are promising but not free. Keeping two copies of the checkpoint adds memory and write traffic at each flush, and the smallest Nemotron model still lost several points on long GPQA-Diamond answers. The production claims for GLM-5.3-Flash and Kimi-K3 come from NVIDIA’s post and the model makers’ reports, not from independent measurement.