arXiv cs.AI· Ze-Wei Liou, Harris Nover, Amey Naik, Andrii Skliar, Igor Shovkun, Ido Shahaf, Tri Dao·· 6 天前AI 评分32
ReplaySSM：缓存 SSM 输入而非状态
ReplaySSM: Cache SSM Inputs, Not State
AI 导读
ReplaySSM 将 SSM 的循环展开为「检查点 + 近期输入缓冲区」，把输出计算与检查点更新分离，从而在多数解码步上避免状态写入，将主要状态内存流量减少约一半。
正文 · AI 翻译
摘要：状态空间模型（SSM）维护固定大小的循环状态，但在每个解码步骤读写该状态会成为混合模型的瓶颈。这种急切的状态管理导致大量状态 I/O 且算术强度低，使推测解码没有廉价的回滚，并将状态量化误差传入后续的循环更新。我们提出 ReplaySSM，它展开循环，将状态表示为检查点和近期 SSM 输入的缓冲区。这将输出计算与检查点更新分离。通过定期更新检查点，ReplaySSM 在大多数步骤上避免了状态写入，并将主要的状态内存流量大致减半。保留近期输入还使推测回滚成为缓冲区指针操作，并支持仅输出的验证而无需物化中间状态。同样的分离允许低精度检查点在大多数步骤上服务于输出计算，而更高精度的检查点服务于循环更新，因此低精度副本中的舍入误差永远不会进入后续更新。我们在 vLLM 中为 Mamba-2 和 Gated DeltaNet 实现了 ReplaySSM。在服务批大小为 256 时，它在四个混合模型上比调优后的 vLLM 基线加速标准解码 $1.20$--$1.47\times$，并在两个 120B 级模型上提供约 $1.8\times$ 于 vLLM 推测解码的吞吐量。与共享低精度检查点相比，分离的状态精度还降低了输出 KL 散度和下游准确率损失。
| 评论： | 11 页（含参考文献和附录共 35 页），6 幅图，2 张表。代码：此 https URL |
| 主题： | 机器学习（cs.LG）；人工智能（cs.AI） |
| 引用为： | arXiv:2610.13778 [cs.LG] |
| （或 arXiv:2610.13778v1 [cs.LG] 用于此版本） | |
| https://doi.org/10.48550/arXiv.2610.13778 arXiv 颁发的 DOI 通过 DataCite（待注册） |
提交历史
来自：Ze-Wei Liou [查看电子邮件]
[v1]
2026 年 10 月 9 日星期五 15:59:40 UTC（691 KB）
来源：arXiv cs.AI · arxiv.org