NVIDIA and Dao AI Lab speed up hybrid-model decoding
ReplaySSM, a new decoding method from NVIDIA and Princeton's Dao AI Lab, rethinks how state-space layers store state so hybrid models decode faster and handle speculative decoding and quantization better.
NVIDIA and the Dao AI Lab at Princeton have released ReplaySSM. The technique targets the state-space model (SSM) layers in hybrid architectures such as Nemotron 3, Qwen3.5 through Qwen3.8, Kimi-K3 and GLM-5.3-Flash. It is already integrated into TensorRT-LLM, vLLM and FlashInfer, and runs in production in GLM-5.3-Flash and Kimi-K3.
SSM layers are meant to make decoding cheap because their state stays the same size as context grows. NVIDIA says that at serving batch sizes, though, SSM layers can be the main decode bottleneck up to 100K tokens of context. On Nemotron-3-Super-NVFP4 at batch size 256 on a single Blackwell Ultra B300, each decode step loads and stores the full state, and one layer's FP32 state takes 4 MiB per sequence. That fused update also makes speculative-decoding rollback expensive. vLLM handles rollback by saving a state snapshot for each draft token. Quantized states add a second problem: rounding error builds up through the recurrence.
ReplaySSM stores the state as a checkpoint plus a buffer of recent inputs. On most steps it computes the output from the checkpoint and the buffer without writing the state back. Only when the buffer fills, after 16 tokens in the tested setup, does it write a new checkpoint to memory. NVIDIA says outputs match conventional decoding up to floating-point differences.
The evaluation covered four Nemotron 3 (Mamba-2) and Qwen3.5 (Gated DeltaNet) models, ranging from 4B dense models on an H100 to 120B-class MoE models on a B300. The reported results are 1.2–1.47× faster decoding, about 2× faster speculative-decoding throughput and better accuracy with quantized states. The full results are in the paper, ReplaySSM: Cache SSM Inputs, Not State, which is on arXiv, with code on GitHub.