ReplaySSM：缓存 SSM 输入，而非状态

arXiv:2610.13778v1 公告类型：新论文 摘要：状态空间模型（SSM）维护一个固定大小的循环状态，但在每个解码步骤中读写该状态可能成为混合模型的性能瓶颈。这种即时状态管理会产生算术强度较低的大量状态 I/O，使推测解码缺乏低成本的回滚方式，并将状态量化误差传递到后续循环更新中。我们提出 ReplaySSM，通过展开递推过程，将状态表示为一个检查点和一个存储近期 SSM 输入的缓冲区，从而将输出计算与检查点更新分离。通过周期性更新检查点，ReplaySSM 避免了大多数步骤中的状态写入，并将占主导地位的状态内存流量减少约一半。保留近期输入还使推测回滚简化为缓冲区指针操作，并支持仅验证输出，无需显式生成中间状态。同样，这种分离也允许在大多数步骤中使用低精度检查点进行输出计算，同时使用更高精度的检查点进行循环更新，因此低精度副本中的舍入误差永远不会进入后续更新。我们在 vLLM 中为 Mamba-2 和 Gated DeltaNet 实现了 ReplaySSM。在推理服务批大小为 256 时，相较于经过调优的 vLLM 基线，它在四个混合模型上将标准解码速度提升至 $1.20$--$1.47$ 倍，并在两个 120B 级别模型上实现了约为 vLLM 推测解码 $1.8$ 倍的吞吐量。与共享低精度检查点相比，为状态分别采用不同精度还能降低输出 KL 散度，并减少下游任务的准确率损失。