NVIDIA and the Dao AI Lab at Princeton University released ReplaySSM, a technique that accelerates decoding in hybrid language models by caching recent inputs rather than storing the full state-space model (SSM) state at every step, according to a paper published October 12, 2026 on arXiv and a blog post by Michelle Horton on the NVIDIA Developer site. The method is already integrated into NVIDIA TensorRT-LLM, vLLM, SGLang and FlashInfer, and runs in production in GLM-5.3-Flash and Kimi-K3 models.
The SSM Decoding Bottleneck
Hybrid models—including NVIDIA Nemotron 3, Qwen3.5 through Qwen3.8, Kimi-K3 and GLM-5.3-Flash—interleave SSM layers such as Mamba-2, Gated DeltaNet or Kimi Delta Attention with a few Transformer attention layers. SSM layers use a fixed-size state instead of a growing key-value cache, but at serving batch sizes the SSM layers can become the latency bottleneck up to 100,000 tokens of context, according to the October 12 report.
On Nemotron-3-Super-NVFP4 at batch size 256 on one NVIDIA Blackwell Ultra B300 GPU, SSM layers accounted for the largest component of decode-step latency up to 100,000 tokens. Two factors drive this: the model has five times as many SSM layers as attention layers, and one layer’s FP32 state is 4 MiB per sequence, equivalent to an FP8 key-value cache for 8,000 tokens per attention layer, the report stated.
Conventional SSM decoding loads the full state from high-bandwidth memory, applies a data-dependent decay and rank-1 update, computes the output, and stores the updated state. Each step performs little arithmetic relative to the bytes moved, making the kernel memory-bandwidth bound rather than compute bound.
How ReplaySSM Works
ReplaySSM represents the SSM state as a checkpoint plus a buffer of recent input records—the decay parameters, keys and values for Mamba-2 layers. On most steps, it appends the current inputs to the buffer and computes the output from the checkpoint and buffer without storing a full state matrix. When the buffer of L records fills, a flush step folds the buffered inputs into the checkpoint, writes the new checkpoint to memory once, and clears the buffer. The outputs match conventional decoding within floating-point rounding differences, according to the paper.
The technique separates output computation, which happens every step, from state updates, which occur once per window. Over a window of L tokens, conventional decoding reads the state L times and writes it L times. ReplaySSM still reads the checkpoint every step to compute outputs but writes it only during the flush, roughly halving the dominant state traffic, the October 12 report stated.
Performance Gains
The authors evaluated four hybrid models from two families—Nemotron 3 with Mamba-2 layers and Qwen3.5 with Gated DeltaNet layers—ranging from 4B dense models on one NVIDIA H100 Tensor Core GPU to 120B-class mixture-of-experts models on one NVIDIA Blackwell Ultra B300 GPU. The implementation in vLLM used a window of L = 16 and kept the SSM state in FP32, the default precision for these models.
At batch size 256, the ReplaySSM kernel ran 1.46× to 1.80× faster than vLLM’s tuned kernels across the four models, and the full decode step ran 1.20× to 1.47× faster, according to measurements reported October 12. At batch size 1, other latency sources dominate and ReplaySSM provided no speedup in these tests. The end-to-end gain grows with batch size because the SSM share of the step grows.
Speculative Decoding Improvements
Speculative decoding has a drafter propose T tokens, verifies them in one pass through the target model, and rolls back to the last accepted token. Attention supports rollback by moving the key-value cache pointer, but once an SSM has summarized drafts into its state, the update has no easy inverse. Serving engines such as vLLM work around this by storing a state snapshot after every draft token, which is one read and T writes per step on a memory-bound path, and the snapshots cap how many requests fit in GPU memory, the report stated.
ReplaySSM appends the inputs of all T draft tokens to the buffer before verification and commits the accepted prefix by advancing the buffer pointer, with no state snapshot to save or restore. A flush folds only committed records into the checkpoint. For T drafts against h committed records, an output-only evaluation computes all T outputs using the checkpoint once, plus two small matrix products over the cached keys, cached values and draft queries, without forming intermediate state. This saves on-chip memory and floating-point operations, and Tensor Cores can accelerate these computations, according to the paper.
On GSM8K prompts with T = 4 at batch size 512 on the two 120B-class models, ReplaySSM delivered 1.85× and 2.10× the throughput of vLLM’s speculative decoding, with matched acceptance lengths, measurements from October 12 showed. The gain has two sources: faster verification—vLLM’s kernel latency grows with T and at T = 8 is almost five times that of standard decoding, while ReplaySSM verifies eight tokens in about the time of one standard step including flushes—and higher concurrency, because draft state snapshots consume GPU memory and cap concurrency. Conventional speculation stopped gaining beyond batch size 128, and on Qwen3.5-122B-A10B-NVFP4 it fell below standard decoding at batch size 512, while ReplaySSM kept scaling with batch size.
Quantization Without Error Accumulation
A lower-precision state requires fewer bytes, reducing latency when memory bandwidth limits kernel performance. On Mamba-2, the non-flush kernel ran 3.8× faster when reading NVFP4 state than FP32 state, according to the report. However, a quantized state produces the output and is the input to the next update, so its rounding error re-enters the recurrence at every step. A key-value cache has no such loop: a quantized entry is read many times but never written back into itself, which is why key-value caches tolerate lower precision than SSM states. Qwen3.5-122B-A10B-NVFP4, for example, uses NVIDIA NVFP4 4-bit floating-point format for its weights, with an FP8 key-value cache and an FP32 SSM state.
ReplaySSM’s non-flush steps read the state only to compute outputs, while flush steps update the checkpoint used in subsequent state updates. The two jobs can use different precisions. The implementation keeps a high-precision checkpoint for the recurrence and a low-precision copy for non-flush reads. A flush starts from the high-precision checkpoint, incorporates the buffered inputs, and casts the result independently into both formats. The low-precision copy is never read back into the recurrence, so its rounding error affects only the outputs of the current window and never propagates across the recurrence, the October 12 report stated.
The authors measured output error as the Kullback–Leibler divergence from FP32-state decoding under teacher forcing on PG-19 text, out to 128,000 generated tokens. A shared FP8 or NVFP4 checkpoint accumulated error with length. Keeping the recurrence in FP32 removed most of the growth: at 128,000 tokens on Qwen3.5-4B, NVFP4 reads with an FP32 recurrence cut the KL from 0.136 to 0.0038. What remained was read error, which did not grow with length. An FP16 recurrence with stochastic rounding tracked FP32 closely, the paper stated.
On GPQA-Diamond, the benchmark with the longest generations of three tested, a shared NVFP4 checkpoint lost 3 to 31 percentage points of accuracy across the four models. NVFP4 reads with an FP16 or FP8 recurrence recovered most of that and kept three out of four models within about 2 percentage points of FP32, according to results in the paper. On Nemotron-3-Super-NVFP4, the shared checkpoint tracked FP32 accuracy for shorter generations but diverged as generation length grew, while the separate-precision configuration stayed close to FP32 throughout. On MMLU-Pro and Math-L5, where generations are generally shorter, three out of four models lost at most 3 percentage points of accuracy with a shared NVFP4 checkpoint, while separate checkpoint configurations stayed within half a percentage point of FP32 on average.
Production Deployment and Open-Source Integration
GLM-5.3-Flash serves with ReplaySSM, and Z.ai’s infrastructure report describes adding it to the Kimi Delta Attention decode kernel, according to the October 12 blog post. Kimi-K3 caches the SSM inputs of draft tokens for verification in its KDA layers and cites ReplaySSM as a concurrent approach, the paper noted. The technique is integrated into TensorRT-LLM (pull requests 14203 and 16464), vLLM (pull requests 48018, 52506 and 52928), and FlashInfer (pull requests 3975 and 4081), with code available on GitHub at github.com/Johnny-Liou/ReplaySSM.
The paper, “ReplaySSM: Cache SSM Inputs, Not State,” is available on arXiv as arxiv.org/abs/2610.13778. The Dao AI Lab published an earlier walkthrough of standard and speculative decoding at tridao.me/blog/2026/replayssm/.
Source: developer.nvidia.com.
Be the first to comment