ReplaySSM caches SSM inputs to cut state traffic and speed hybrid decoding

NVIDIA and the Dao AI Lab at Princeton University propose ReplaySSM, a way to represent state-space model (SSM) state as a checkpoint plus a buffer of recent inputs. The technique separates output computation from state updates, addressing three bottlenecks in hybrid-model decoding: state memory traffic, rollback for speculative decoding, and error accumulation in quantized state.

Hybrid models interleave SSM layers with a few attention layers. NVIDIA Nemotron 3, Qwen3.5 through Qwen3.8, Kimi-K3, and GLM-5.3-Flash all use this design, with Mamba-2, Gated DeltaNet (GDN), or Kimi Delta Attention (KDA) as the SSM layer. SSM layers use a fixed-size state instead of a growing key-value (KV) cache, but at serving batch sizes they can become the latency bottleneck up to 100K tokens of context.

Why SSM layers slow down decoding

An SSM decoding step loads and stores the full state while performing little arithmetic, making the kernel bound by high-bandwidth memory (HBM) traffic rather than compute. Two factors drive this. First, models like Nemotron-3-Super-NVFP4 have five times as many SSM layers as attention layers. Second, 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 equals an attention layer's KV cache for 8K tokens.

Speculative decoding also suffers. Attention supports rollback by moving a KV cache pointer, but once an SSM summarizes draft tokens into its state, the update has no easy inverse. Serving engines like vLLM work around this by storing a state snapshot after every draft token, which adds reads and writes on a memory-bound path and caps how many requests fit in GPU memory.

Quantization compounds the problem. A quantized SSM state produces the output and feeds the next update, so rounding error re-enters the recurrence at every step. A KV cache has no such loop, which is why it tolerates lower precision.

How ReplaySSM represents state

ReplaySSM unrolls the SSM recurrence over a window of tokens. On most steps, called non-flush steps, it appends current inputs to a buffer and computes the output from the checkpoint plus buffer without storing a state. When the buffer fills, a flush step folds the buffered inputs into the checkpoint, writes it to HBM once, and clears the buffer. Outputs match conventional decoding up to floating-point differences.

This roughly halves dominant state traffic. At batch size 256 on one NVIDIA Blackwell Ultra B300 GPU, the kernel runs 1.46x to 1.80x faster, and the decode step runs 1.20x to 1.47x faster across four models. At batch size 1, other latency sources dominate and ReplaySSM provides no speedup in these tests.

Simpler rollback and low-precision reads

In speculative decoding, ReplaySSM appends all draft token inputs to the buffer before verification and commits the accepted prefix afterward by advancing the buffer pointer. Records beyond the pointer are ignored, and the next append overwrites them. No state snapshot is needed per draft token. A flush folds only committed records into the checkpoint.

Output-only evaluation computes verification results without forming an intermediate state, using two small matrix products over cached keys, cached values, and draft queries. This saves on-chip memory and FLOPs, and Tensor Cores can speed up the computation.

ReplaySSM also supports a high-precision checkpoint for recurrence and a low-precision copy for reads, limiting quantization error accumulation. Three of four evaluated models stayed within about 2 percentage points of FP32 accuracy on GPQA-Diamond.

The paper, "ReplaySSM: Cache SSM Inputs, Not State," is available on arXiv with code on GitHub. It is integrated into NVIDIA TensorRT-LLM, vLLM, and FlashInfer, and runs in production in GLM-5.3-Flash and Kimi-K3.