论文

ReplaySSM:缓存SSM输入而非状态

ReplaySSM: Cache SSM Inputs, Not State

精选理由

ReplaySSM通过缓存输入而非状态,大幅提升SSM模型性能,特别适合大模型推理场景。

ReplaySSM是一种新型状态空间模型优化方法,通过将状态表示为检查点和最近输入的缓冲区,分离输出计算与检查点更新。该技术在vLLM中实现,针对Mamba-2和Gated DeltaNet模型,在批量大小256下,比优化后的vLLM基准加速1.20-1.47倍,在两个120B级模型上提供约1.8倍的推测解码吞吐量。

原文 · arXiv cs.AI

ReplaySSM: Cache SSM Inputs, Not State

State-space models (SSMs) maintain a fixed-size recurrent state, but reading and writing it at every decoding step can bottleneck hybrid models. This eager state management causes large state I/O with low arithmetic intensity, leaves speculative decoding no cheap rollback, and feeds state quantization error into subsequent recurrent updates. We introduce ReplaySSM, which unrolls the recurrence to represent the state as a checkpoint and a buffer of recent SSM inputs. This separates output computation from checkpoint updates. By updating the checkpoint periodically, ReplaySSM avoids state writes on most steps and roughly halves the dominant state memory traffic. Keeping recent inputs also makes speculative rollback a buffer-pointer operation and enables output-only verification without materializing intermediate states. The same separation allows a low-precision checkpoint to serve output computation on most steps while a higher-precision checkpoint serves recurrent updates, so rounding error in the low-precision copy never enters later updates. We implement ReplaySSM in vLLM for Mamba-2 and Gated DeltaNet. At serving batch size 256, it accelerates standard decoding by $1.20$--$1.47\times$ over a tuned vLLM baseline across four hybrid models, and delivers about $1.8\times$ the throughput of vLLM's speculative decoding on the two 120B-class models. Separate state precisions also reduce output KL divergence and downstream accuracy losses compared with a shared low-precision checkpoint.