ReplaySSM: Rethinking State Storage for Faster Hybrid-Model Decoding

NVIDIA and Princeton's Dao AI Lab released ReplaySSM, a method that represents the SSM state as a checkpoint plus a buffer of recent inputs to cut state memory traffic, enable rollback for speculative decoding, and limit error accumulation in a quantized state. The paper, "ReplaySSM: Cache SSM Inputs, Not State" (arXiv 2610.13778), reports that SSM layers rather than attention become the latency bottleneck up to 100K tokens of context at serving batch sizes, with one layer's FP32 state at 4 MiB per sequence versus 0.5 KiB per token for an FP8 KV entry. ReplaySSM is integrated into NVIDIA TensorRT-LLM, vLLM, and FlashInfer and runs in production in GLM-5.3-Flash and Kimi-K3, with code on GitHub.

Hybrid models interleave state-space model SSM layers with a few attention layers. NVIDIA Nemotron 3 https://developer.nvidia.com/nemotron , Qwen3.5 through Qwen3.8 https://qwen.ai/blog?id=qwen3.5 , Kimi-K3 https://arxiv.org/abs/2607.24653 , and GLM-5.3-Flash https://z.ai/blog/glm-5.3-flash are all built this way, with Mamba-2 https://arxiv.org/abs/2405.21060 , Gated DeltaNet GDN https://arxiv.org/abs/2412.06464 , or Kimi Delta Attention KDA https://arxiv.org/abs/2510.26692 as the SSM layer. SSM layers are designed to make decoding fast: they use a fixed-size state in place of a key-value KV cache that grows with context length. At serving batch sizes, however, the SSM layers, not attention, can be the latency bottleneck up to 100K tokens of context. SSM layers also complicate two techniques that reliably speed up attention decoding: speculative decoding and quantization. ReplaySSM makes SSM layers the fast part again. It represents the SSM state as a checkpoint plus a buffer of recent inputs. This separates output computation from state updates and addresses three challenges in current SSM decoding: state traffic, rollback for speculative decoding, and error accumulation in a quantized state. ReplaySSM is joint work between NVIDIA and the Dao AI Lab at Princeton. The paper, ReplaySSM: Cache SSM Inputs, Not State https://arxiv.org/abs/2610.13778 , details all three benefits and presents the full evaluation. It is available on arXiv, with code on GitHub https://github.com/Johnny-Liou/ReplaySSM . It is integrated into NVIDIA TensorRT-LLM, vLLM, and FlashInfer, and also runs in production in GLM-5.3-Flash and Kimi-K3. Why SSM layers can become a decoding bottleneck Take the Mamba-2 layer used in NVIDIA Nemotron 3 Ultra https://arxiv.org/abs/2606.15007 as an example. It decodes one token with the following recurrence: where is the state, is a data-dependent decay, and , , and are the token’s inputs. Each step loads from high-bandwidth memory HBM , applies the decay and a rank-1 update, computes the output with , and stores back. GDN and KDA add a delta-rule correction, but their updates can be written in the same outer-product form, so what follows applies to them too. During decoding, the state does two jobs. It is read to produce this token’s output, and it is updated to become the next token’s state. Conventional decoding does both at once: each step updates the state, reads the output from it, and stores it for the next step. All three challenges come from this fusion. State memory traffic An SSM decoding step loads and stores the full state while performing little arithmetic on it, so the kernel is bound by HBM bandwidth, not compute. At serving batch sizes, that makes the SSM layers, not attention, the largest component of the decode step up to 100K tokens on Nemotron-3-Super-NVFP4 https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 at batch size 256 on one NVIDIA Blackwell Ultra B300 GPU Figure 2, left . Two factors drive this: - Layer count: The model has five times as many SSM layers as attention layers. - State size: One layer’s FP32 state is 4 MiB per sequence, while an FP8 KV entry is 0.5 KiB per token, so an SSM layer’s state is the size of an attention layer’s KV cache for 8K tokens. Rollback for speculative decoding Speculative decoding https://arxiv.org/abs/2211.17192 has a drafter propose T tokens, verifies them in one pass through the target model, and rolls back to the last accepted token. Attention naturally supports rollback by moving the KV cache pointer. But once an SSM has summarized the drafts into its state, the update has no easy inverse, so rollback becomes expensive. Serving engines such as vLLM work around this by storing a state snapshot after every draft token and picking the one that matches the accepted prefix. That is one read and T writes per step on a memory-bound path, and the snapshots cap how many requests fit in GPU memory. As a result, speculative decoding, a reliable speedup for Transformers, barely helps SSMs at serving batch sizes. State quantization The 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 KV cache has no such loop: a quantized entry is read many times but never written back into itself, which is why KV caches tolerate lower precision than SSM states. Qwen3.5-122B-A10B-NVFP4, for example, uses the NVIDIA NVFP4 4-bit floating-point format for its weights, with an FP8 KV cache and an FP32 SSM state. How ReplaySSM represents the state The SSM recurrence can be unrolled over a window of h tokens: where is the state before the window and is the cumulative decay. Conventional decoding follows the left side step by step. ReplaySSM uses the right side: it represents the state as a checkpoint plus a small buffer of recent input records , , for Mamba-2. On most steps, which we call non-flush steps, it appends the current inputs to the buffer and computes the output from the checkpoint and the buffer without storing a state. When the buffer of L records fills, a flush step folds the buffered inputs into the checkpoint, writes the new checkpoint to HBM once, and clears the buffer. The outputs match conventional decoding up to floating-point differences. ReplaySSM’s approach can be summarized in three questions: What does ReplaySSM do? It caches recent SSM inputs and represents the state as a checkpoint plus those inputs. What does this change? It separates output computation from state updates: outputs are computed every step, while the checkpoint is updated once per window. What does this achieve? It reduces state memory traffic, simplifies speculative rollback, and supports low-precision state reads without propagating their rounding error through recurrent updates. How we evaluated ReplaySSM We evaluate four hybrid models from two families, Nemotron 3 with Mamba-2 layers and Qwen3.5 with GDN layers, from 4B dense models on one NVIDIA H100 Tensor Core GPU to 120B-class mixture-of-experts MoE models on one NVIDIA Blackwell Ultra B300 GPU Table 1 . We implemented ReplaySSM in vLLM and compared it with vLLM’s own kernels at their best launch configuration. Performance runs keep the SSM state in FP32, the default precision for these models, and use a window of L = 16, the best-performing window size for all four models. Speculative decoding tests use each model’s multi-token prediction MTP head to generate draft tokens. | Model | SSM layer | Weights | GPU | SSM / attention layers | |---|---|---|---|---| | Nemotron-3-Nano-4B https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 | Mamba-2 | BF16 | 1×H100 | 21 / 4 | | Nemotron-3-Super-120B-A12B https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 | Mamba-2 | NVFP4 | 1×B300 | 40 / 8 | | Qwen3.5-4B https://huggingface.co/Qwen/Qwen3.5-4B | GDN | BF16 | 1×H100 | 24 / 8 | | Qwen3.5-122B-A10B https://huggingface.co/nvidia/Qwen3.5-122B-A10B-NVFP4 | GDN | NVFP4 | 1×B300 | 36 / 12 | Table 1. Evaluated models and hardware. Each model runs on one GPU without tensor parallelism Benefit 1: Reduce state memory traffic Over a window of L tokens, conventional decoding reads the state L times and writes it L times. ReplaySSM still reads the checkpoint on every step, because the output needs it, but writes it only once, during the flush. A non-flush step has one job—computing the output—and writes only the new input record, a few vectors instead of a full matrix. This roughly halves the dominant state traffic. At serving batch sizes, memory bandwidth limits the conventional kernel’s performance, so reducing state memory traffic speeds up decoding. At batch size 1, other sources of latency dominate, and ReplaySSM provides no speedup in these tests. At batch size 256, the kernel runs 1.46× to 1.80× faster, and the decode step runs 1.20× to 1.47× faster across the four models Figure 3 . The end-to-end gain grows with batch size because the SSM share of the step grows. Benefit 2: Simplify rollback and speed up verification In speculative decoding, ReplaySSM appends the inputs of all T draft tokens to the input buffer before verification and commits the accepted prefix afterward by advancing the buffer pointer. As with a KV cache, records beyond the pointer are ignored, and the next append overwrites them. There is no state snapshot to save for each draft token and no state to restore Figure 2, middle . A flush folds only committed records into the checkpoint, so it never waits for the accepted length. This separation offers another benefit. Conventional decoding must produce two things per token, the output and the updated state, because the next token consumes the state. A non-flush step in ReplaySSM produces only the output, and computing only the output opens up a different computation method. For one update from a zero state, . The expression on the left builds a state and reads it. The expression on the right forms a scalar and scales a vector without building the state. For T drafts against h committed records, the same idea gives the following expression: with D holding the cumulative decays of the draft positions and W a causally masked matrix of decays. We call this output-only evaluation. The checkpoint is used once for all T queries, and everything else is two small matrix products over the cached keys, the cached values, and the draft queries. No intermediate state is formed, which saves on-chip memory and FLOPs, and Tensor Cores can speed up these computations. Also, is shared by all value heads in a group. This is the same chunkwise computation that prefill kernels already use, with the buffer and the drafts as the chunk. GDN can also use this form. Output-only evaluation applies to standard decoding too, but it matters most in verification, where T multiplies the arithmetic per step. Figure 4 shows speculative serving throughput on GSM8K prompts with T = 4 on the two 120B-class models. At batch size 512, ReplaySSM delivers 1.85× and 2.10× the throughput of vLLM’s speculative decoding, with matched acceptance lengths. The gain has two sources. - Faster verification: vLLM’s kernel stores one snapshot per token, so its latency grows with T and at T = 8 is almost five times that of vLLM’s standard decoding step. ReplaySSM writes no state on most steps, so it verifies eight tokens in about the time of one standard step, flushes included Figure 5 . - Higher concurrency: Draft state snapshots consume GPU memory and cap concurrency. Conventional speculation stops gaining beyond batch size 128, and on Qwen3.5-122B-A10B-NVFP4, it falls below standard decoding at batch size 512, while ReplaySSM keeps scaling with batch size. Benefit 3: Limit error accumulation with separate precisions A lower-precision state requires fewer bytes, which reduces latency when memory bandwidth limits kernel performance. On Mamba-2, the non-flush kernel runs 3.8× faster when reading NVFP4 state than when reading FP32 state. ReplaySSM changes quantization in two ways. First, the checkpoint is written once per window instead of once per step, so a single low-precision checkpoint is rounded once every L steps. The Nemotron 3 Ultra report https://arxiv.org/abs/2606.15007 explores this periodic checkpointing approach. It reduces rounding frequency, but rounding error still enters later updates. Second, 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. We keep 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. Writing the additional low-precision copy during a flush reduces the amount of data read on every non-flush step. Acceptance, rollback, and buffer management are unchanged Figure 2, right . Figure 6 measures the output error as the Kullback–Leibler KL divergence from FP32-state decoding under teacher forcing on PG-19 text, out to 128K generated tokens. A shared FP8 or NVFP4 checkpoint accumulates error with length. Keeping the recurrence in FP32 removes most of the growth: at 128K tokens on Qwen3.5-4B, NVFP4 reads with an FP32 recurrence cut the KL from 0.136 to 0.0038. What remains is the read error, which does not grow with length. An FP16 recurrence with stochastic rounding SR tracks FP32 closely, so a half-precision recurrence checkpoint is enough. Better quantization of the read copy, such as separate scales along the state and channel dimensions, could cut the remaining read error further. We also checked whether the KL gap shows up in downstream accuracy. Table 2 shows GPQA-Diamond, the benchmark with the longest generations of the three we ran. A shared NVFP4 checkpoint loses 3 to 31 percentage points of accuracy across the four models. NVFP4 reads with an FP16 or FP8 recurrence recover most of that and keep three out of four models within about 2 percentage points of FP32. | GPQA-Diamond, change from FP32 points | FP32 accuracy | Shared NVFP4 checkpoint | NVFP4 read, FP16+SR recurrence | NVFP4 read, FP8+SR recurrence | |---|---|---|---|---| | Nemotron-3-Nano-4B | 56.7 | −31.3 | −6.7 | −5.6 | | Nemotron-3-Super-NVFP4 | 79.3 | −12.5 | −1.4 | −1.4 | | Qwen3.5-4B | 76.3 | −3.2 | +0.7 | +1.5 | | Qwen3.5-122B-A10B-NVFP4 | 82.7 | −5.2 | −2.0 | −1.7 | Table 2. GPQA-Diamond accuracy with an NVFP4 read checkpoint, expressed as the change from each model’s FP32 reference. The shared checkpoint uses stochastic rounding. Standard decoding, L = 16 Figure 7 shows how accuracy loss varies with generation length on GPQA-Diamond. Across the three benchmarks: - GPQA-Diamond: On Nemotron-3-Super-NVFP4, the shared checkpoint tracks FP32 accuracy for shorter generations but diverges from it as generation length grows. The separate-precision configuration stays close to FP32 throughout. - MMLU-Pro and Math-L5: Generations are several times shorter than those in GPQA-Diamond. The three larger models lose at most 3 percentage points of accuracy with a shared NVFP4 checkpoint, while the separate checkpoint configurations stay within half a percentage point of FP32 on average. The paper includes results for all three benchmarks and a long-input RULER test. Get started with ReplaySSM ReplaySSM is in production. GLM-5.3-Flash https://z.ai/blog/glm-5.3-flash serves with ReplaySSM, and Z.ai’s infrastructure report https://z.ai/blog/glm-built-its-inference-infrastructure describes adding it to the KDA decode kernel. Kimi-K3 https://arxiv.org/abs/2607.24653 caches the SSM inputs of draft tokens for verification in its KDA layers and cites ReplaySSM as a concurrent approach. ReplaySSM is also integrated into open-source inference frameworks such as TensorRT-LLM, vLLM, SGLang, and FlashInfer. To learn more, check out these resources: - Read the paper: ReplaySSM: Cache SSM Inputs, Not State TODO: Link - Try the code: github.com/Johnny-Liou/ReplaySSM https://github.com/Johnny-Liou/ReplaySSM - Serve hybrid models with ReplaySSM in TensorRT-LLM: PR 14203 https://github.com/NVIDIA/TensorRT-LLM/pull/14203 , PR 16464 https://github.com/NVIDIA/TensorRT-LLM/pull/16464 - Serve hybrid models with ReplaySSM in vLLM: PR 48018 https://github.com/vllm-project/vllm/pull/48018 , PR 52506 https://github.com/vllm-project/vllm/pull/52506 , PR 52928 https://github.com/vllm-project/vllm/pull/52928 - Use the ReplaySSM kernels in FlashInfer: PR 3975 https://github.com/flashinfer-ai/flashinfer/pull/3975 , PR 4081 https://github.com/flashinfer-ai/flashinfer/pull/4081 - Read the earlier post from the Dao AI Lab https://tridao.me/blog/2026/replayssm/ for a longer walkthrough of standard and speculative decoding