ReplaySSM: Cache SSM Inputs, Not State

📄 arXiv: 2610.13778v1 📥 PDF

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

分类: cs.LG, cs.AI

发布日期: 2026-10-09

备注: 11 pages (35 pages counting references and appendix), 6 figures, 2 tables. Code: https://github.com/Johnny-Liou/ReplaySSM

💡 一句话要点

提出ReplaySSM以解决状态空间模型的瓶颈问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 状态空间模型 解码优化 内存管理 推测回滚 混合模型 高效计算 深度学习

📋 核心要点

  1. 现有的状态空间模型在每个解码步骤中频繁读写状态，导致性能瓶颈和高内存流量。
  2. ReplaySSM通过将状态表示为检查点和输入缓冲区，分离输出计算与状态更新，减少了状态写入频率。
  3. 在实验中，ReplaySSM在多个模型上实现了1.20到1.47倍的加速，并在120B级模型上提高了推测解码的吞吐量约1.8倍。

📝 摘要（中文）

状态空间模型（SSMs）在每个解码步骤中都需要读取和写入固定大小的递归状态，这会导致混合模型的瓶颈。为了缓解这一问题，本文提出ReplaySSM，通过将递归展开为检查点和最近SSM输入的缓冲区，分离输出计算与检查点更新。该方法在大多数步骤中避免了状态写入，减少了状态内存流量，并使得推测回滚变得更为高效。ReplaySSM在vLLM中实现，显著提高了标准解码速度，并在多个混合模型中表现出色。

🔬 方法详解

问题定义：本文旨在解决状态空间模型在解码过程中频繁读写状态所带来的性能瓶颈和高内存流量问题。现有方法在状态管理上过于急切，导致低算术强度和推测解码的回滚成本高。

核心思路：ReplaySSM的核心思想是将递归展开，使用检查点和最近的SSM输入缓冲区来分离输出计算与状态更新。这种设计使得在大多数步骤中避免状态写入，从而降低内存流量。

技术框架：ReplaySSM的整体架构包括两个主要模块：检查点和输入缓冲区。检查点定期更新，而输入缓冲区则存储最近的SSM输入，支持快速的输出计算和推测回滚。

关键创新：ReplaySSM的主要创新在于将状态管理与输出计算分离，使得低精度检查点可以在大多数步骤中用于输出计算，而高精度检查点则用于递归更新，从而避免低精度带来的误差影响后续更新。

关键设计：在实现中，ReplaySSM采用了不同精度的检查点设计，优化了内存使用和计算效率。此外，缓冲区的设计使得推测回滚操作变得简单高效，进一步提升了模型的性能。

🖼️ 关键图片

img_0
img_1
img_2

📊 实验亮点

在实验中，ReplaySSM在批量大小为256的情况下，相较于调优后的vLLM基线，在四个混合模型上实现了1.20到1.47倍的加速。此外，在120B级模型上，其推测解码的吞吐量约为vLLM的1.8倍，显著提升了性能。

🎯 应用场景

ReplaySSM的设计可以广泛应用于需要高效状态管理的混合模型，尤其是在自然语言处理和计算机视觉等领域。其减少内存流量和提高解码速度的能力，使其在实时应用中具有重要价值，未来可能推动更复杂模型的开发和应用。

📄 摘要（原文）

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.