NVIDIA, Princeton Lab Introduce ReplaySSM for Faster AI Decoding

NVIDIA and Princeton University’s Dao AI Lab introduced ReplaySSM, a method that speeds inference in hybrid language models by reducing repeated memory writes during token decoding.

At a serving batch size of 256, ReplaySSM improved the state-space-model kernel by 1.46× to 1.80× and accelerated full decode steps by 1.20× to 1.47× versus tuned conventional vLLM kernels. The results were measured across Nemotron-3-Nano-4B, Nemotron-3-Super-120B-A12B, Qwen3.5-4B and Qwen3.5-122B-A10B.

ReplaySSM represents the recurrent state as a checkpoint plus a buffer of recent state-space-model inputs. Instead of rewriting the full state at every decoding step, the system updates the checkpoint when the buffer fills. The tests used a replay window of 16 inputs.

The 4B models ran on one NVIDIA H100 GPU, while the 120B-class models ran on one NVIDIA Blackwell Ultra B300 GPU. At batch size one, end-to-end performance stayed within a few percent of the conventional baseline because other sources of latency dominated.

The method also targets speculative serving, where a model produces draft tokens before a larger model verifies them. In tests using four draft tokens and a batch size of 512, ReplaySSM delivered 1.85× and 2.10× the throughput of conventional vLLM speculative decoding on the two 120B-class models. The tests used GSM8K prompts and matched acceptance lengths.

Those same tests produced throughput of 2.16× and 1.90× conventional standard decoding on the two large models. The figures reflect the specified models, hardware, batch sizes and benchmark configuration rather than a general speedup across all deployments.

ReplaySSM has been integrated into TensorRT-LLM, vLLM, FlashInfer and SGLang, and is used in production serving for GLM-5.3-Flash and Kimi-K3.

The method separates the precision used for recurrent updates from the precision used to read outputs. In a Qwen3.5-4B test generating 128,000 tokens, NVFP4 reads with an FP32 recurrence reduced output KL divergence from 0.136 to 0.0038 compared with FP32-state decoding.

The work was submitted to arXiv on Oct. 9, and NVIDIA announced the collaboration and released links to the paper and code on Oct. 12.