ReplaySSM is a decoding method for hybrid models that interleave state-space model (SSM) layers with attention layers. Models such as NVIDIA Nemotron 3, Qwen3.5–Qwen3.8, Kimi-K3 and GLM-5.3-Flash use this hybrid design, with SSM layers implemented as Mamba-2, Gated DeltaNet (GDN), or Kimi Delta Attention (KDA). ReplaySSM was developed jointly by NVIDIA and the Dao AI Lab at Princeton.
In conventional SSM decoding each token step loads the full SSM state from high-bandwidth memory (HBM), applies a decay and a rank-1 update, computes the token output, and writes the full state back. Because the kernel performs little arithmetic relative to the amount of state transferred, it becomes memory-bandwidth bound; at serving batch sizes this can make SSM layers the latency bottleneck up to 100K tokens of context. Conventional SSM updates also complicate two common acceleration techniques: speculative decoding (which needs cheap rollback) and low-precision state reads (because quantization error feeds back into the recurrence).
ReplaySSM represents the SSM state as a persistent checkpoint plus a short buffer of recent input records. On most steps (non-flush steps) it appends the current inputs to the buffer and computes outputs from checkpoint+buffer without writing a new full state. When the buffer reaches length L, a flush folds buffered inputs into the checkpoint, writes the updated checkpoint once to HBM, and clears the buffer. In effect output computation is separated from state updates.
Why this matters
ReplaySSM addresses three practical decoding problems:
-
State memory traffic: Over an L-token window conventional decoding reads and writes the full state L times. ReplaySSM still reads the checkpoint each step for output computation but writes the checkpoint only once per window; non-flush steps write only the small input record. This substantially reduces HBM traffic and thus speeds up kernels that are memory-bound.
-
Rollback for speculative decoding: In speculative decoding a drafter proposes T tokens and the verifier checks them; rollback must return to the last accepted token. Attention supports rollback cheaply by moving the KV cache pointer. Conventional SSMs require snapshots per draft token, imposing an extra read and T writes on a memory-bound path and limiting concurrency. ReplaySSM instead appends draft inputs to the buffer and advances the buffer pointer for commits; no per-draft state snapshot is needed, and flushes only incorporate committed records.
-
Quantization error accumulation: A quantized SSM state is written back into the recurrence, so rounding errors re-enter the state at every step and can accumulate. ReplaySSM separates the recurrence checkpoint (kept in high precision) from a low-precision copy used only for non-flush reads. The low-precision copy is never read back into the recurrence, so its rounding error affects only outputs within the current window and does not accumulate across windows.
Technical summary
The SSM recurrence (for Mamba-2) is S_t = α_t S_{t-1} + v_t k_t^⊤, o_t = S_t q_t. Unrolling over an h-token window yields S_t = γ_h S_ckpt + Σ_{j=1}^{h} γ_{j+1:h} v_j k_j^⊤, where S_ckpt is the checkpoint before the window and γ are cumulative decays. ReplaySSM uses this unrolled form: keep S_ckpt and a small buffer of recent (α_j, v_j, k_j) records, compute outputs from S_ckpt plus the buffer, and only occasionally fold the buffer into the checkpoint.
In verification phases ReplaySSM can use an output-only evaluation that avoids forming intermediate states: for a single update o = (v k^⊤) q = v (k^⊤ q). For T drafts against h committed records a compact form O = S_ckpt Q D + V (W ⊙ (K^⊤ Q)) computes all draft outputs efficiently; the checkpoint is used once for all T queries and the rest are small matrix products on the cached keys, values and draft queries. This saves FLOPs and on-chip memory and benefits from Tensor Core acceleration.
Evaluated models and hardware
The paper evaluates four hybrid models across two families, running each model on a single GPU without tensor parallelism (window L = 16 was used throughout and FP32 checkpoint recurrence was the default for performance tests):
- Nemotron-3-Nano-4B — Mamba-2, BF16 weights, 1×H100, 21 SSM / 4 attention layers
- Nemotron-3-Super-120B-A12B — Mamba-2, NVFP4 weights, 1×B300, 40 SSM / 8 attention layers
- Qwen3.5-4B — GDN, BF16 weights, 1×H100, 24 SSM / 8 attention layers
- Qwen3.5-122B-A10B — GDN, NVFP4 weights, 1×B300, 36 SSM / 12 attention layers
ReplaySSM was implemented in vLLM for comparison against vLLM’s tuned kernels. Speculative decoding tests used each model’s multi-token prediction head to generate drafts.
Measured benefits
-
Reduced state memory traffic and speedups: At serving batch size 256 the ReplaySSM kernel runs 1.46×–1.80× faster and the full decode step runs 1.20×–1.47× faster across the four tested models. Gains increase with batch size because the SSM portion of the step grows.
-
Faster verification and higher concurrency in speculative decoding: On GSM8K prompts with T = 4 and L = 16, for the two 120B-class models at batch size 512 ReplaySSM achieved 1.85× and 2.10× the throughput of vLLM’s speculative decoding while matching acceptance lengths. The speedup comes from both faster verification (no per-draft snapshot writes) and higher concurrency because snapshots no longer consume GPU memory.
-
Quantization behavior and accuracy: ReplaySSM’s separate-precision approach keeps the recurrence in high precision and uses a lower-precision read copy for non-flush reads. This eliminates most accumulation of rounding error across long generations. For example, on Qwen3.5-4B at 128K generated tokens NVFP4 reads with an FP32 recurrence reduced KL divergence from 0.136 to 0.0038 compared with FP32-state decoding. Downstream accuracy (GPQA-Diamond) showed that a shared NVFP4 checkpoint could lose 3–31 percentage points versus FP32, while NVFP4 reads combined with an FP16 or FP8 recurrence (with stochastic rounding) recovered most loss and kept three out of four models within ~2 percentage points of FP32.
Availability and adoption
The paper (ReplaySSM: Cache SSM Inputs, Not State) is on arXiv and the implementation is on GitHub (github.com/Johnny-Liou/ReplaySSM). ReplaySSM kernels are integrated into NVIDIA TensorRT-LLM, vLLM and FlashInfer and the method is deployed in production for GLM-5.3-Flash and used in Kimi-K3. Z.ai’s infrastructure report documents adding the approach to a KDA decode kernel.
Conclusion
By caching recent SSM inputs and treating state as a checkpoint plus a buffer, ReplaySSM reduces memory traffic, makes speculative rollback cheap, and confines quantization error to a window. The approach yields measurable decode-step speedups at large batch sizes, substantial speculative-decoding throughput gains for large models, and practical quantization strategies, and it is already integrated into several inference stacks.
