ReplaySSM cuts memory traffic to speed up hybrid-model decoding

NVIDIA and Princeton’s Dao AI Lab have developed ReplaySSM, a method for speeding up inference in hybrid AI models by changing how their state is stored. NVIDIA reports up to 1.80× faster SSM decoding kernels at batch size 256, with larger gains for speculative decoding on tested 120B-class models.

NVIDIA Watch analysis

What happened

Hybrid models combine state-space model (SSM) layers with attention layers. ReplaySSM represents an SSM’s state as a checkpoint plus a buffer of recent inputs, rather than repeatedly writing the full state to memory. In most decoding steps, it adds inputs to the buffer; it updates the checkpoint when that buffer fills.

NVIDIA says this roughly halves the dominant state-memory traffic. In its tests across four models, the method made SSM decoding kernels 1.46× to 1.80× faster at batch size 256, and the full decode step 1.20× to 1.47× faster. At batch size one, the tests showed no speed-up.

For speculative decoding, the buffer also avoids saving and restoring a full state snapshot for every draft token. NVIDIA reports 1.85× and 2.10× the throughput of vLLM’s speculative decoding on two 120B-class models at batch size 512, with matched acceptance lengths. The NVIDIA technical blog links to the paper and implementation.

Why it matters

Inference speed is not only about how quickly a model can calculate. For these SSM layers, moving large states to and from GPU memory can be the bottleneck. ReplaySSM targets that traffic, and its reported gains suggest that a change in data handling can improve both ordinary decoding and the more complex business of checking draft tokens.

The results are workload-dependent: the reported speed-ups come at larger batch sizes, not batch size one. That makes this especially relevant to teams serving many simultaneous requests, rather than a blanket promise of faster responses for every user.

Our read

This is a useful systems idea with numbers attached, not just a new acronym released into the wild. The reported gains are substantial for the tested serving workloads, and the batch-size caveat matters. Operators should compare it on their own models and traffic before expecting the same improvement; the useful next step is to see how the paper’s evaluation holds up beyond the configurations NVIDIA reports.

What to watch

• Whether independent evaluations reproduce the throughput gains on other hybrid models and serving setups.
• How performance changes at smaller batch sizes and across different sequence lengths.
• Whether the integration in vLLM, TensorRT-LLM and FlashInfer makes the technique practical for more operators.

Discussion spark: For AI services handling heavy traffic, should teams prioritise techniques that boost throughput at large batch sizes, even if they offer no speed-up for single requests?

Independent WittyWires tracker for public updates about NVIDIA. Not affiliated with or endorsed by NVIDIA; this is not an official account.