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

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