Reproduce “ReplaySSM: Cache SSM Inputs, Not State” on OpenResearch

About this paper

Ze-Wei Liou, Harris Nover, Amey Naik, Andrii Skliar, Igor Shovkun, Ido Shahaf, Tri Dao

ReplaySSM reduces recurrent-state writes during state-space model decoding by keeping a checkpoint and buffering recent inputs, then periodically folding them into the checkpoint. It also supports pointer-based speculative rollback and separate low- and high-precision checkpoints. At batch size 256, whole-model decode latency improved 1.20–1.47× across four hybrid models; gains depend on setup and batch size, and checkpoint reads remain.

Key insights

• The approach trades fewer full-state writes for additional replay-buffer processing; the best balance depends on window length, batch size, and hardware behavior.
• Benefits are workload-dependent: gains are strongest when state traffic is a bottleneck at large batches, while batch-one kernel performance is at or below parity in the reported tests.
• Separating output precision from recurrence precision can reduce long-run quantization error without requiring all output computations to use full-precision state.

Summary by alphaXiv.