ReplaySSM Accelerates Hybrid-Model Decoding by Caching Inputs Instead of Full State

ReplaySSM Accelerates Hybrid-Model Decoding by Caching Inputs Instead of Full State Developed jointly by NVIDIA and Princeton's Dao AI Lab, ReplaySSM restructures state storage in hybrid SSM-attention models to eliminate decoding latency bottlenecks. By expressing state as an intermittent checkpoint paired with a short input buffer, the method curtails high-bandwidth memory traffic, streamlines speculative decoding rollbacks, and mitigates quantization error accumulation. At serving batch sizes, conventional state-space model (SSM) layers can become memory-bandwidth bottlenecks up to 100K context tokens due to reading and writing the entire state at every step. ReplaySSM represents the state as a checkpoint plus a buffer of recent inputs, computing outputs on non-flush steps and only updating the checkpoint to memory once the buffer fills. Standard decoding with ReplaySSM achieves 1.46× to 1.80× kernel speedups and 1.20× to 1.47× step speedups at batch size 256 across Nemotron 3 and Qwen3.5 models. In speculative decoding, ReplaySSM commits drafts by advancing a pointer and uses output-only evaluation, delivering up to 2.10× higher serving throughput at batch size 512 compared to vLLM. Decoupling precision allows low-precision (such as NVFP4) state reads for output calculations alongside higher-precision recurrence updates, halting error compounding across extended generation contexts. SSM Decoding Bottlenecks in Hybrid Architectures Hybrid architectures such as Nemotron 3, Qwen3.5, Kimi-K3, and GLM-5.3-Flash interleave SSM layers like Mamba-2, Gated DeltaNet, or Kimi Delta Attention with sparse attention layers. Although SSMs replace growing KV caches with fixed-size states, conventional decoding fuses output calculation with state updates at every step. This makes SSMs memory-bound: a single layer's FP32 state consumes 4 MiB per sequence, roughly matching an attention layer's KV cache for 8K tokens. With high SSM-to-attention layer ratios, loading and storing the state creates heavy high-bandwidth memory (HBM) traffic that bottlenecks latency at serving batch sizes. State Checkpointing and Buffered Input Formulation ReplaySSM unrolls the recurrent SSM formulation across a window of tokens, defining the state as an initial checkpoint plus a buffer of input records. During intermediate non-flush steps, the model appends incoming inputs to the buffer and evaluates outputs directly from the checkpoint and buffer without storing a new state. When the buffer reaches capacity, a flush step folds the buffered records into the checkpoint and writes it to HBM once. This design maintains numerical consistency with conventional decoding while cutting write traffic during non-flush steps down to a few input vectors. Standard Decoding Gains and Bandwidth Relief Evaluations across 4B dense models on NVIDIA H100 GPUs and 120B-class mixture-of-experts models on Blackwell Ultra B300 GPUs demonstrated substantial improvements when evaluated with a window size of 16. While batch size 1 shows no speedup due to other latency factors, batch size 256 yields 1.46× to 1.80× faster kernel execution and 1.20× to 1.47× faster decode steps across Nemotron 3 and Qwen3.5. Because the share of step latency occupied by SSM layers expands alongside batch size, end-to-end performance scaling improves concurrently. Streamlined Speculative Verification and Output-Only Evaluation Under conventional speculative decoding, serving engines save state snapshots per draft token to support rollback, taxing memory bandwidth and capping concurrency. ReplaySSM appends draft token inputs to the buffer and handles rollbacks by adjusting the buffer pointer, eliminating snapshot storage entirely. It also applies output-only evaluation to compute outputs via small matrix products over cached keys, values, and draft queries using Tensor Cores without building intermediate states. On GSM8K benchmarks with draft length 4 at batch size 512, ReplaySSM achieved 1.85× to 2.10× the throughput of vLLM while maintaining verifying performance for 8 tokens in approximately the time of a single standard step. Decoupled Precision to Prevent Quantization Drift In recurrent updates, state quantization errors repeatedly feed back into subsequent steps. ReplaySSM decouples precision roles: a high-precision checkpoint (such as FP16 with stochastic rounding or FP32) is preserved for recurrent updates, while a low-precision copy (such as NVFP4) is generated during flushes specifically for non-flush output reads. Because the low-precision copy never recirculates into the recurrence, error does not accumulate over generation length. On 128K generated tokens of PG-19 with Qwen3.5-4B, this separation reduced Kullback–Leibler divergence from 0.136 to 0.0038 against FP32 baselines, protecting downstream benchmark performance across GPQA-Diamond, MMLU-Pro, and Math-L5. Production Deployment and Open-Source Integration ReplaySSM is currently deployed in production engines for GLM-5.3-Flash and Kimi-K3's KDA layers. The framework is also merged and available across inference ecosystems, including TensorRT-LLM, vLLM, SGLang, and FlashInfer, alongside public research code.

NVIDIA与普林斯顿联手推出ReplaySSM：重构状态存储以加速混合模型解码 针对混合模型（如Nemotron 3、Qwen3.5、GLM-5.3-Flash等）在服务批处理下SSM层引发的内存瓶颈、投机解码回滚困难及量化误差累积问题，NVIDIA与普林斯顿大学Dao AI Lab联合开发了ReplaySSM技术。该方案通过缓存最近输入并将状态拆分为检查点与输入缓冲区，解耦了输出计算与状态更新。 在服务批处理场景下，由于SSM层数远多于注意力层且状态内存流量受限于HBM带宽，SSM层在长达10万Token的上下文中反而成为解码延迟的主要瓶颈。 ReplaySSM采用“检查点+最近输入缓冲区”的形式表示SSM状态，大多数非刷新步骤仅计算输出并写入少量输入向量，仅在缓冲区填满时触发刷新回写，将状态内存流量大致减半。 在投机解码中，ReplaySSM通过仅移动缓冲区指针实现无需快照的快速回滚，并引入免建中间状态的“纯输出计算”形式，在批处理512时使120B级模型服务吞吐量提升至vLLM的1.85至2.10倍。 量化状态更新与读取解耦为高低双精度，利用高精度检查点进行递归更新、低精度副本执行非刷新读取，阻断了舍入误差在递归中的连续累积。 该技术已集成至TensorRT-LLM、vLLM、SGLang与FlashInfer，并在GLM-5.3-Flash与Kimi-K3等模型中投入生产部署。 传统SSM解码面临的三大瓶颈 在常规解码中，SSM步骤必须同时读取状态以输出当前Token，并执行外积更新后存回HBM。这导致了三个核心问题：首先是内存带宽受限，例如单个FP32状态达到4 MiB，导致其在批处理256时的延迟占据整个解码步的大头；其次是投机解码回滚困难，传统引擎必须在每个草稿Token后保存完整状态快照，消耗巨大内存且难以扩展；最后是量化误差累积，量化状态的舍入误差会循环进入后续迭代，使得低精度量化难以适用。 ReplaySSM的状态拆分机制与流量削减 ReplaySSM将长度为L的滑动窗口展开计算。在普通的非刷新步骤中，它直接基于检查点与缓存输入生成输出，仅向缓冲区追加输入向量而不写回完整状态；只有在缓冲区满后的刷新步骤中，才将累积输入折叠并一次性写回新检查点。在批处理256的测试中，该设计令核心运行速度提高1.46至1.80倍，整个解码步骤提速1.20至1.47倍。 投机解码验证加速与纯输出计算 在投机解码验证时，ReplaySSM直接将T个草稿Token输入追加至缓冲区，并根据接受的前缀移动指针完成回滚，彻底摒弃了快照保存操作。同时，非刷新计算利用结合律采用“纯输出计算”公式，直接基于缓存的键、值与查询进行矩阵乘法，无需构建中间状态并能调用Tensor Core加速。这使得8个草稿Token的验证耗时接近单步标准解码，大幅提升了系统的并发处理能力。 分离精度消除量化误差扩散 为了降低状态读取开销同时避免精度崩塌，ReplaySSM将递归更新与输出读取采用不同精度管理：刷新步骤生成高精度检查点（如FP16加随机舍入或FP8）维持状态递推，并独立铸造NVFP4低精度副本专供非刷新步骤读取。测试表明，该架构将12.8万Token下的KL散度从0.136大幅降低至0.0038，并在GPQA-Diamond等长序列基准评测中使模型精度保持在接近FP32的2个百分点以内。

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. Figure 1. Conventional decoding loads and stores the full state at every step. ReplaySSM caches recent inputs, computes outputs from the checkpoint plus buffer, and writes the checkpoint only when the buffer fills 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: S t = α t S t − 1 + v t k t ⊤ , o t = S t q t , S_t = \alpha_t S_{t-1} + v_t k_t^\top, \qquad o_t = S_t q_t, where S t S_t is the state, α t \alpha_t is a data-dependent decay, and v t v_t , k t k_t , and q t q_t are the token’s inputs. Each step loads S t − 1 S_{t-1} from high-bandwidth memory (HBM), applies the decay and a rank-1 update, computes the output with q t q_t , and stores S t 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. Figure 2. Three challenges of conventional SSM decoding and the ReplaySSM answer to each. Left: decode-step latency breakdown for Nemotron-3-Super-NVFP4 at batch size 256 on one NVIDIA Blackwell Ultra B300 GPU 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: S t = α t S t − 1 + v t k t ⊤ = γ h S c k p t + ∑ j = 1 h γ j + 1 : h v j k j ⊤ , S_t = \alpha_t S_{t-1} + v_t k_t^\top = \gamma_h\, S_{\mathrm{ckpt}} + \sum_{j=1}^{h} \gamma_{j+1:h}\, v_j k_j^\top, where S c k p t 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 ( α j \alpha_j , v j v_j , k 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. Model SSM layer Weights GPU SSM / attention layers Nemotron-3-Nano-4B Mamba-2 BF16 1×H100 21 / 4 Nemotron-3-Super-120B-A12B Mamba-2 NVFP4 1×B300 40 / 8 Qwen3.5-4B GDN BF16 1×H100 24 / 8 Qwen3.5-122B-A10B 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. Figure 3. Standard decoding speedup over tuned vLLM at batch size 256. Kernel numbers include flush costs 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 ⊤ ) q = v ( k ⊤ q ) o = (v k^\top)q = v(k^\top q) . The expression on the left builds a d v × d k 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: O = S c k p t Q D + V ( W ⊙ ( K ⊤ Q ) ) , O = S_{\mathrm{ckpt}}\, Q D + V \bigl(W \odot (K^\top Q)\bigr), 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 ⊤ Q 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. Figure 4. Speculative serving throughput on GSM8K prompts, T = 4, L = 16, FP32 state, one B300 per model 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. Figure 5. Verification kernel latency on Nemotron-3-Super-NVFP4 at batch size 256, relative to vLLM’s standard decoding step, flushes included 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. Figure 6. Output KL from conventional FP32-state decoding on Qwen3.5-4B, L = 16. Separate precisions for the two jobs, output computation and recurrent update, remove most of the accumulated error 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 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. Figure 7. GPQA-Diamond accuracy by quartile of FP32 reference generation length on Nemotron-3-Super-NVFP4 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