ReplaySSM

Code for the paper ReplaySSM: Cache SSM Inputs, Not State by Ze-Wei Liou, Harris Nover, Amey Naik, Andrii Skliar, Igor Shovkun, Ido Shahaf and Tri Dao (Princeton University and NVIDIA).

arXiv · NVIDIA Technical Blog · Dao AI Lab blog

The NVIDIA post summarizes the paper and its deployments, and the earlier Dao AI Lab post gives a longer walkthrough of standard and speculative decoding.

ReplaySSM runs in production in GLM-5.3-Flash and Kimi-K3 and is integrated in TensorRT-LLM, vLLM, SGLang and FlashInfer. See Adoption for details.

ReplaySSM represents an SSM state as a checkpoint and a bounded buffer of recent inputs. This separates output computation from checkpoint updates, with three benefits:

• Reduced memory traffic: most decoding steps compute outputs without writing the state. Periodic flushes update the checkpoint, roughly halving the dominant state memory traffic.
• Cheap rollback: speculative decoding discards rejected inputs by moving a buffer pointer. Output-only verification avoids materializing intermediate recurrent states.
• Separate state precisions: non-flush steps use a low-precision checkpoint for output computation, while flushes use a higher-precision checkpoint for recurrent updates. Rounding error in the low-precision copy never enters later recurrent updates.

This repository is a vLLM fork with ReplaySSM implementations for Mamba2 and Gated DeltaNet (GDN), covering standard and speculative decoding, along with standalone precision kernels. Upstream vLLM has merged the Mamba2 standard decoding kernel, and its FlashInfer backend covers Mamba2 standard and speculative decoding. The remaining Triton kernels, the emulated quantization and the low-precision CuTe kernels are only available here.

The vLLM integration uses --use-replayssm for standard decoding and --use-replayssm-spec alongside a speculative decoding configuration. The flags are mutually exclusive. --replayssm-buffer-len 16 selects the paper's window length. Relevant metadata and dispatch live in mamba_attn.py, gdn_attn.py, and the Mamba layers.

The vLLM integration is based on upstream commit ee1d99636702b8ed9ad88c2c2b833d331dce01c1, with ReplaySSM extensions.

Build the serving engine from this checkout using the upstream source-build instructions. Set VLLM_VERSION_OVERRIDE=0.23.1+replayssm when installing this snapshot without upstream Git tags. The original project overview is retained in README.vllm.md. The two precision folders have separate setup instructions and do not require a vLLM build.

The emulation kernels implement separate read and recurrence precision and expose their quantization behavior. The CuTe kernels support the latency benefit of reading the checkpoint in low precision on B300. Both folders include focused correctness tests and usage instructions.

e2e_decode_speedup.py times standard decoding with the conventional kernel and with ReplaySSM. e2e_spec_decode_throughput.py compares speculative decoding throughput against standard decoding and vLLM's native speculative decoding.

python benchmarks/replayssm/e2e_decode_speedup.py --model-id nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16
python benchmarks/replayssm/e2e_decode_speedup.py --model-id Qwen/Qwen3.5-4B
python benchmarks/replayssm/e2e_spec_decode_throughput.py --batch-size 512 \
    --model-id nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4

ReplaySSM is used in production models. GLM-5.3-Flash runs it in its serving stack, and the Z.ai infrastructure report describes adding it to the KDA decode kernel. Kimi-K3 (Section 5.4.2) caches the projected draft inputs for MTP verification in its KDA layers and cites ReplaySSM as a concurrent approach.

ReplaySSM is also integrated in TensorRT-LLM, vLLM, SGLang and FlashInfer, as summarized in the NVIDIA Technical Blog post. The paper authors contributed the merged pull requests below, except the row marked "NVIDIA", which was contributed independently.

Pull requests in progress: FlashInfer #4140 and #6274; vLLM #48792, #49847 and #49887

If you use ReplaySSM in your work, please cite:

@misc{liou2026replayssmcachessminputs,
      title={ReplaySSM: Cache SSM Inputs, Not State},
      author={Ze-Wei Liou and Harris Nover and Amey Naik and Andrii Skliar and Igor Shovkun and Ido Shahaf and Tri Dao},
      year={2026},
      eprint={2610.13778},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2610.13778},
}