NVIDIA and Dao AI Lab Present ReplaySSM for Faster Hybrid Model Decoding

How ReplaySSM Works

Hybrid models interleave state-space model (SSM) layers with a few attention layers, and NVIDIA Nemotron 3, Qwen3.5, Kimi-K3, and GLM-5.3-Flash follow this design. SSM layers use a fixed-size state instead of a KV cache, but at serving batch sizes the SSM layers, not attention, can become the latency bottleneck up to 100K tokens of context. Conventional decoding loads and stores the full state at every step.

ReplaySSM represents the state as a checkpoint plus a buffer of recent input records. On most steps, called non-flush steps, it appends the current inputs to the buffer and computes the output from the checkpoint and buffer without storing a state. When the buffer of 16 records fills, a flush step folds the buffered inputs into the checkpoint, writes the new checkpoint once, and clears the buffer. The outputs match conventional decoding up to floating-point differences.

Cutting State Traffic and Simplifying Rollback

Over a window of 16 tokens, conventional decoding reads and writes the state 16 times each. ReplaySSM writes the checkpoint only once per window, which roughly halves the dominant state traffic. At batch size 256 on NVIDIA Blackwell Ultra B300 GPUs, the decoding kernel runs 1.46x to 1.80x faster and the decode step runs 1.20x to 1.47x faster across four evaluated models. At batch size 1, other latency sources dominate and ReplaySSM provides no speedup in these tests.

For speculative decoding, ReplaySSM appends draft token inputs to the buffer before verification and commits the accepted prefix by advancing the buffer pointer, with no state snapshots to save or restore. This enables an output-only evaluation that forms two small matrix products without building an intermediate state. On GSM8K prompts with 120B-class models, ReplaySSM delivers 1.85x and 2.10x the speculative throughput of vLLM at batch size 512, with matched acceptance lengths.

Limiting Quantization Error

A quantized SSM state feeds back into the recurrence, so rounding error accumulates with length, unlike a KV cache entry. ReplaySSM keeps a high-precision checkpoint for the recurrence and a low-precision copy for reads, so read error never propagates across the recurrence. On Qwen3.5-4B at 128K tokens, NVFP4 reads with an FP32 recurrence cut the output KL divergence from 0.136 to 0.0038. On GPQA-Diamond, three of four models stay within about 2 percentage points of FP32 accuracy with NVFP4 reads and an FP16 or FP8 recurrence, while a shared NVFP4 checkpoint loses 3 to 31 points.

Availability

ReplaySSM is joint work between NVIDIA and the Dao AI Lab at Princeton University, detailed in the paper ReplaySSM: Cache SSM Inputs, Not State, with code on GitHub. It is integrated into TensorRT-LLM, vLLM, SGLang, and FlashInfer, and runs in production in GLM-5.3-Flash and Kimi-K3, which caches the SSM inputs of draft tokens for verification in its KDA layers and cites ReplaySSM as a concurrent approach.

Sources: Nvidia Dev