ReplaySSM: Cache SSM Inputs, Not State
cs.LG, cs.AI
Submitted: 2026-10-09
Updated: 2026-10-09
Comments: 11 pages (35 pages counting references and appendix), 6 figures, 2 tables. Code: https://github.com/Johnny-Liou/ReplaySSM
Code: https://github.com/Johnny-Liou/ReplaySSM
License: http://creativecommons.org/licenses/by/4.0/
The gist: State-space models (SSMs) maintain a fixed-size recurrent state, but reading and writing it at every decoding step can bottleneck hybrid models.
Abstract
State-space models (SSMs) maintain a fixed-size recurrent state, but reading and writing it at every decoding step can bottleneck hybrid models. This eager state management causes large state I/O with low arithmetic intensity, leaves speculative decoding no cheap rollback, and feeds state quantization error into subsequent recurrent updates. We introduce ReplaySSM, which unrolls the recurrence to represent the state as a checkpoint and a buffer of recent SSM inputs. This separates output computation from checkpoint updates. By updating the checkpoint periodically, ReplaySSM avoids state writes on most steps and roughly halves the dominant state memory traffic. Keeping recent inputs also makes speculative rollback a buffer-pointer operation and enables output-only verification without materializing intermediate states. The same separation allows a low-precision checkpoint to serve output computation on most steps while a higher-precision checkpoint serves recurrent updates, so rounding error in the low-precision copy never enters later updates. We implement ReplaySSM in vLLM for Mamba-2 and Gated DeltaNet. At serving batch size 256, it accelerates standard decoding by 1.20 -- 1.47 times over a tuned vLLM baseline across four hybrid models, and delivers about 1.8 times the throughput of vLLM's speculative decoding on the two 120B-class models. Separate state precisions also reduce output KL divergence and downstream accuracy losses compared with a shared low-precision checkpoint.
Sources
- Training Verifiers to Solve Math Word Problems
- Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality
- Mamba: Linear-Time Sequence Modeling with Selective State Spaces
- Measuring Mathematical Problem Solving With the MATH Dataset
- RULER: What's the Real Context Size of Your Long-Context Language Models?
- Repeat After Me: Transformers are Better than State Space Models at Copying
- Kimi Linear: An Expressive, Efficient Attention Architecture
- Kimi K3: Open Frontier Intelligence
- NVIDIA Nemotron 3: Efficient and Open Intelligence
- Nemotron 3 Ultra: Open, Efficient Mixture-of-Experts Hybrid Mamba-Transformer Model for Agentic Reasoning
- Compressive Transformers for Long-Range Sequence Modelling
- GPQA: A Graduate-Level Google-Proof Q&A Benchmark
- An Empirical Study of Mamba-based Language Models
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks