ReplaySSM Caches Inputs Instead of State to Speed Up Hybrid-Model Decoding

Why SSM layers slow down hybrid models

Hybrid models interleave SSM layers with a few attention layers. NVIDIA names Nemotron 3, Qwen3.5 through Qwen3.8, Kimi-K3 and GLM-5.3-Flash as models built this way, using Mamba-2, Gated DeltaNet (GDN) or Kimi Delta Attention (KDA) as the SSM layer. SSM layers use a fixed-size state instead of a key-value (KV) cache that grows with context length, which is meant to make decoding fast.

At serving batch sizes, however, NVIDIA says the SSM layers rather than attention can be the latency bottleneck up to 100K tokens of context. The post describes three causes, all tied to the fact that conventional decoding updates the state, reads the output from it and stores it again at every step.

• State memory traffic: each step loads and stores the full state while doing little arithmetic, so the kernel is bound by high-bandwidth memory (HBM) bandwidth. In Nemotron-3-Super-NVFP4 the model has five times as many SSM layers as attention layers, and one layer's FP32 state is 4 MiB per sequence.
• Speculative decoding rollback: once an SSM has folded draft tokens into its state, the update has no easy inverse. Engines such as vLLM store a state snapshot after every draft token, which means one read and T writes per step and limits how many requests fit in GPU memory.
• Quantization error: a quantized state feeds the next update, so its rounding error re-enters the recurrence at every step. A KV cache has no such loop.

How ReplaySSM works

ReplaySSM unrolls the SSM recurrence over a window of tokens. Instead of updating the state each step, it keeps a checkpoint and a small buffer of recent input records. On most steps, which NVIDIA calls non-flush steps, it appends the current inputs to the buffer and computes the output from the checkpoint and the buffer without storing a state. When the buffer of L records fills, a flush step folds the buffered inputs into the checkpoint, writes the new checkpoint to HBM once and clears the buffer.

According to the post, outputs match conventional decoding up to floating-point differences. In short, output computation is separated from state updates: outputs are produced every step, while the checkpoint is updated once per window.

For speculative decoding, this means draft inputs can be appended to the buffer rather than snapshotting state. For quantization, a high-precision checkpoint handles the recurrence while a low-precision copy is used for reads, which limits error accumulation.

Evaluation and reported results

The team evaluated four hybrid models from two families: Nemotron 3 with Mamba-2 layers and Qwen3.5 with GDN layers. They range from 4B dense models on one NVIDIA H100 GPU to 120B-class mixture-of-experts models on one NVIDIA Blackwell Ultra B300 GPU. ReplaySSM was implemented in vLLM and compared with vLLM's own kernels at their best launch configuration. Performance runs kept the SSM state in FP32 and used a window of L = 16, which NVIDIA says was the best-performing size for all four models. Speculative decoding tests used each model's multi-token prediction head to generate drafts.

NVIDIA reports the following results:

• SSM decoding kernels run 1.46x to 1.80x faster at batch size 256 on B300 GPUs, with state memory traffic reduced by roughly half.
• Speculative throughput on 120B-class models is 1.85x to 2.10x higher.
• Three of four evaluated models stay within about 2 percentage points of FP32 accuracy on GPQA-Diamond when using the low-precision state read.

Availability and what it means

The paper, titled ReplaySSM: Cache SSM Inputs, Not State, is available on arXiv, with code on GitHub. NVIDIA says the technique is integrated into TensorRT-LLM, vLLM and FlashInfer, and that it also runs in production in GLM-5.3-Flash and Kimi-K3.

For teams serving hybrid models, the work targets the layers that the post identifies as a bottleneck at realistic batch sizes. If the reported gains hold in other deployments, speculative decoding and lower-precision state could become practical for SSM-based models, which NVIDIA says have so far benefited little from them.