状态空间模型（SSM）在每个解码步都读取和写入固定大小的循环状态，这会成为混合模型的瓶颈。频繁的状态 I/O 计算强度低，导致投机解码缺乏廉价回滚，并且状态量化误差会传递到后续循环更新。ReplaySSM 将递归展开为检查点加最近 SSM 输入的缓冲区，把状态表示为 checkpoint 与输入缓冲。输出计算与检查点更新分离，周期性更新检查点，使大多数步骤无需写状态，约减半主要的状态内存流量。保留最近输入使得投机回滚仅需移动缓冲指针，并可在不生成中间状态的情况下进行仅输出验证。同一分离还允许低精度检查点用于大多数步骤的输出计算，而高精度检查点用于循环更新，低精度的舍入误差不会进入后续更新。我们在 vLLM 中实现了 ReplaySSM，支持 Mamba-2 与 Gated DeltaNet。批量大小 256 时，相比调优的 vLLM 基线，标准解码加速 $1.20$–$1.47\times$，在两个 120B 级模型上实现约 $1.8\times$ 的投机解码吞吐。不同精度的检查点还能降低输出 KL 散度和下游精度损失。
点评