ReplaySSM: Rethinking State Storage for Faster Hybrid-Model Decoding

extract

Hybrid models interleave state-space model (SSM) layers with a few attention layers. NVIDIA Nemotron 3, Qwen3.5 through Qwen3.8, Kimi-K3, and GLM-5.3-Flash are all built this way, with Mamba-2, Gated DeltaNet (GDN), or Kimi Delta Attention (KDA) 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, details all three benefits and presents the full evaluation. It is available on arXiv, with code on GitHub. 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 as an example. It decodes one token with the following recurrence: where \(S_t\) is the state, \(\alpha_t\) is a data-dependent decay, and \(v_t\), \(k_t\), and \(q_t\) are the token’s inputs. Each step loads \(S_{t-1}\) from high-bandwidth memory (HBM), applies the decay and a rank-1 update, computes the output with \(q_t\), and stores \(S_t\) 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 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 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 \(S_{\mathrm{ckpt}}\) is the state before the window and \(\gamma\) 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 (\(\alpha_j\), \(v_j\), \(k_j\)) 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. 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, \(o = (v k^\top)q = v(k^\top q)\). The expression on the left builds a \(d_v \times d_k\) 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, \(K^\top Q\) 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 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. 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 generally shorter than those in GPQA-Diamond. Three out of four 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 serves with ReplaySSM, and Z.ai’s infrastructure report describes adding it to the KDA decode kernel. Kimi-K3 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 - Try the code: github.com/Johnny-Liou/ReplaySSM - Serve hybrid models with ReplaySSM in TensorRT-LLM: PR 14203, PR 16464 - Serve hybrid models with ReplaySSM in vLLM: PR 48018, PR 52506, PR 52928 - Use the ReplaySSM kernels in FlashInfer: PR 3975, PR 4081 - Read the earlier post from the Dao AI Lab for a longer walkthrough of standard and speculative decoding