Loading post…
NVIDIAMichelle HortonProblem & solution
ReplaySSM: Rethinking State Storage for Faster Hybrid-Model Decoding
Summary
Hybrid models interleave state-space model (SSM) layers with attention layers, but SSM layers can become a decoding bottleneck at serving batch sizes due to memory-bound state updates. ReplaySSM represents the SSM state as a checkpoint plus a buffer of recent inputs, separating output computation from state updates. This reduces state memory traffic by roughly half, simplifies rollback for speculative decoding by avoiding state snapshots, and supports low-precision state reads without propagating rounding error through the recurrence, enabling separate precision for output computation and state updates. Evaluated on Nemotron 3 and Qwen3.5 models from 4B to 120B parameters on H100 and B300 GPUs, ReplaySSM achieves 1.20× to 1.47× faster decode steps at batch size 256, with end-to-end gains growing with batch size. The technique is integrated into TensorRT-LLM, vLLM, SGLang, and FlashInfer, and runs in production in GLM-5.3-Flash and Kimi-K3.
Context
Hybrid models that interleave state-space model (SSM) layers with attention layers are used in models like NVIDIA Nemotron 3, Qwen3.5, Kimi-K3, and GLM-5.3-Flash. SSM layers use a fixed-size state instead of a growing KV cache, but at serving batch sizes, the SSM layers can become the latency bottleneck due to memory-bound state updates. Additionally, SSM layers complicate speculative decoding and state quantization. ReplaySSM addresses these challenges by representing the state as a checkpoint plus a buffer of recent inputs.
Approach / What changed
ReplaySSM represents the SSM state as a checkpoint plus a small buffer of recent input records (α_j, v_j, k_j for Mamba-2). On non-flush steps, it appends inputs to the buffer and computes the output from the checkpoint and buffer without storing a state. When the buffer fills (L=16), a flush step folds the buffered inputs into the checkpoint, writes the new checkpoint once, and clears the buffer. This separates output computation from state updates, reducing memory traffic, simplifying speculative rollback, and enabling separate precision for output reads and recurrent updates.
Takeaways
SSM layers in hybrid models can become the decoding bottleneck at serving batch sizes because each step loads and stores the full state from HBM, making the kernel memory-bandwidth-bound rather than compute-bound.
ReplaySSM reduces state memory traffic by roughly half by writing the checkpoint only once per window of L tokens instead of every step, while non-flush steps only write the new input record.
ReplaySSM enables low-precision state reads without propagating rounding error through the recurrence by using separate precision for the read copy and the recurrent update, keeping the recurrence in FP32 while reading a quantized copy.
FAQ
How does ReplaySSM reduce state memory traffic?
ReplaySSM represents the SSM state as a checkpoint plus a buffer of recent inputs. On non-flush steps, it only writes the new input record instead of the full state. The checkpoint is written only once per window of L tokens, roughly halving the dominant state memory traffic.
Why does speculative decoding barely help SSMs at serving batch sizes?
Once an SSM has summarized draft tokens into its state, the update has no easy inverse, so rollback becomes expensive. Serving engines work around this by storing a state snapshot after every draft token, which adds memory traffic and caps how many requests fit in GPU memory.
How does ReplaySSM handle state quantization without error accumulation?
ReplaySSM separates the two jobs of the state: output computation and recurrent update. It writes a low-precision copy of the checkpoint for output reads, but keeps the recurrence in FP32. The low-precision copy is never read back into the recurrence, so its rounding error affects only the current window's outputs and does not propagate.
Source and freshness
- Source
- developer.nvidia.com
- Published
- Added to Yomu
Summaries are generated from the original by a language model. Check the original before relying on it.
Keep reading
Related reading
The Good Stuff
Paid picks we use ourselves. Links carry a referral code — a subscription through one earns us a commission at no extra cost to you.
The best engineering reads, once a week
One useful email. Unsubscribe anytime.