ReplaySSM：把 SSM 状态换成检查点加输入缓冲，vLLM 上解码提速…
论文AI 评分 78/100arXiv cs.AI
AI 聚合
本条为 AI 依据下方公开信源自动整理生成的摘要，不构成转载，可能存在偏差，请以原文为准。整理者：贝果科技 AI 资讯助手。
一篇 arXiv 论文提出 ReplaySSM，针对状态空间模型每步解码都要读写固定大小循环状态、导致状态 I/O 大而算术强度低的问题，改为把递推展开，用「检查点 + 近期 SSM 输入缓冲」表示状态，从而把输出计算与检查点更新分开。检查点定期更新，多数步骤不再写状态，主要状态内存流量约减半；保留近期输入还让投机解码的回滚变成缓冲区指针操作，并支持只验证输出而不物化中间状态。该分离也让低精度检查点负责多数步骤的输出计算、高精度检查点负责递推更新，低精度舍入误差不会进入后续更新。作者在 vLLM 中为 Mamba-2 和 Gated DeltaNet 实现该方案：服务批大小 256 时，四个混合模型上标准解码比调优后的 vLLM 基线快 1.20–1.47 倍；两个 120B 级模型上吞吐约为 vLLM 投机解码的 1.8 倍。分离状态精度还降低了输出 KL 散度与下游准确率损失。
全文梳理
AI 摘要依据下方信源原文自动整理，非原文转载状态读写为何成为瓶颈
状态空间模型在解码时维持一个固定大小的循环状态，但每一步都要读取并写回这个状态。论文指出，这种即时状态管理带来两个后果：一是状态 I/O 量大而算术强度低，成为混合模型的瓶颈；二是投机解码缺少廉价回滚手段。此外，状态量化误差会进入后续的递推更新，逐步累积。这三类问题都源于把状态当作每步必须落地的实体来维护。
用检查点加输入缓冲替代状态
ReplaySSM 的做法是把递推展开，不再直接保存状态，而是用一份检查点加上一段近期 SSM 输入的缓冲来表示它。这样输出计算与检查点更新被拆成两件事：检查点按周期更新，大部分解码步骤因此不需要写状态，论文称主要的状态内存流量约减少一半。保留近期输入还带来两个附带效果——投机解码的回滚只需移动缓冲区指针，输出验证也不必物化中间状态。
精度分离与实测结果
同一套分离结构还允许两种精度并存：低精度检查点承担多数步骤的输出计算，高精度检查点负责递推更新，于是低精度副本的舍入误差不会进入后续更新。作者在 vLLM 中为 Mamba-2 和 Gated DeltaNet 实现了该方案。在服务批大小 256 下，四个混合模型的标准解码相比调优后的 vLLM 基线快 1.20 至 1.47 倍；在两个 120B 级模型上，吞吐约为 vLLM 投机解码的 1.8 倍。相比共享低精度检查点，分离状态精度还降低了输出 KL 散度与下游准确率损失。
为什么值得看
混合 SSM 模型的推理瓶颈在状态读写，这里给出了可落地的缓存改法，并附 vLLM 实现与实测加速比。
论文
信源1 家
关键事实
相关 · 论文
本页内容由 AI 自动聚合公开信源生成，仅供了解行业动态参考，不构成任何投资或决策建议。如需引用请以原文出处为准。