ReplaySSM：缓存 SSM 输入加速混合模型解码
来自 NVIDIA Developer · 期号 2026-10-13
摘要
混合模型把 SSM 层和注意力层拼在一起，SSM 层用固定大小的状态代替随上下文增长的 KV 缓存，本应让解码更快。但在服务端批量推理时，SSM 层反而成了瓶颈——它每步都要读写完整状态，而计算量很小，被 HBM 带宽卡住。 NVIDIA 和普林斯顿 Dao AI Lab 提出的 ReplaySSM 换了一种思路：不再每步更新状态，而是把状态拆成一个检查点加一个近期输入缓存。大部分步骤只往缓存里追加输入、从检查点算输出，不写状态；缓存满了才折叠进检查点，写一次。这样状态内存流量大约减半，在 Blackwell Ultra B300 上批量 256 时，SSM 解码内核快了 1.46 到 1.80 倍。 推测解码的另一个麻烦是回退：SSM 把草稿 token 收进状态后没法撤销，传统做法每步存快照，既慢又占显存。ReplaySSM 把草稿输入追加到缓存里，验证后只推进指针，不存状态，回退就是忽略指针后的记录。120B 级模型的推测吞吐量因此提高了 1.85 到 2.10 倍。 量化 SSM 状态时，误差每步都会循环进入下一次更新，不像 KV 缓存读一次就完。ReplaySSM 用高精度检查点做递推、低精度副本做读取，把误差隔离开。四个模型在 GPQA-Diamond 上的精度损失控制在 FP32 的 2 个百分点以内。 ReplaySSM 已集成到 TensorRT-LLM、vLLM 和 FlashInfer，GLM-5.3-Flash 和 Kimi-K3 已在生产中使用。论文和代码都已开源。
深读
导读 混合模型将状态空间模型(SSM)层与注意力层结合，SSM 层用固定大小的状态替代随上下文增长的 KV 缓存，理论上应加速解码。但在服务端批量推理时，SSM 层反而因每步读写完整状态、计算量小而被 HBM 带宽卡住，成为瓶颈。NVIDIA 与普林斯顿 Dao AI Lab 提出的 ReplaySSM 通过将状态拆分为检查点加近期输入缓存，将状态内存流量减半，并简化推测解码的回退与状态量化，在 Blackwell Ultra B300 上批量 256 时 SSM 解码内核加速 1.46 到 1.80 倍。 关键信息 - ReplaySSM 将 SSM 状态表示为一个检查点(checkpoint)加一个近期输入缓存(buffer)，大部分步骤只往缓存追加输入、从检查点算输出，不写状态；缓存满时才折叠进检查点，写一次。这使状态内存流量大约减半。 - 在 Blackwell Ultra B300 GPU 上、批量 256 时，四个模型的 SSM 解码内核加速 1.46× 到 1.80×，解码步骤整体加速 1.20× 到 1.47×。 - 在推测解码中，ReplaySSM 将草稿 token 的输入追加到缓存，验证后只推进指针，不存状态快照。在 120B 类 MoE 模型上、批量 512 时，推测吞吐量是 vLLM 基线方案的 1.85× 到 2.10×。 - ReplaySSM 支持对状态使用不同精度：用高精度检查点做递推更新，用低精度副本做读取。在 GPQA-Diamond 上，三个模型在 NVFP4 读取加 FP16 递推下，精度与 FP32 基线相差约 2 个百分点以内。 逐节讲解 SSM 层为何成为解码瓶颈 以 Nemotron 3 Ultra 使用的 Mamba-2 层为例，每步解码执行递推：从 HBM 加载上一状态，做衰减和秩一更新，计算输出，再写回状态。状态同时承担两个角色，被读取以产生当前输出，被更新以成为下一步的状态。传统方法将两者融合在一起。 两个因素使 SSM 层成为瓶颈：一是层数，模型中的 SSM 层数通常是注意力层的 5 倍；二是状态大小，单层 FP32 状态每序列 4 MiB，而 FP8 KV 条目每 token 仅 0.5 KiB，一个 SSM 层的状态大小相当于注意力层 8K token 的 KV 缓存。由于每步读写完整状态而计算量很小，内核被 HBM 带宽卡住。 推测解码的回退在 SSM 上也很麻烦：注意力层只需移动 KV 缓存指针即可回退，但 SSM 将草稿 token 收进状态后无法撤销。vLLM 等引擎的做法是每步存状态快照，这既慢又占显存，导致推测解码在 SSM 上收益甚微。 状态量化同样困难：量化后的状态既产生输出又作为下一步的输入，其舍入误差每步都会重新进入递推循环，而 KV 缓存没有这种循环。 ReplaySSM…