在早期 ReplaySSM 工作的基础上，NVIDIA 与普林斯顿大学 Dao AI Lab 探索了缓存 SSM 输入如何加速解码，并以更少的误差累积支持低精度状态。
混合模型将状态空间模型（SSM）层与少量注意力层交错排列。NVIDIA Nemotron 3、Qwen3.5 至 Qwen3.8、Kimi-K3 和 GLM-5.3-Flash 均采用这种方式构建，其中 SSM 层使用 Mamba-2、Gated DeltaNet（GDN）或 Kimi Delta Attention（KDA）。
SSM 层旨在使解码快速：它们使用固定大小的状态，替代随上下文长度增长的键值（KV）缓存。然而，在服务批大小下，SSM 层（而非注意力）可能成为上下文长度达到 100K token 时的延迟瓶颈。SSM 层还使两种可靠加速注意力解码的技术变得复杂：推测解码和量化。
ReplaySSM 让 SSM 层重新成为快速的部分。它将 SSM 状态表示为检查点加上近期输入的缓冲区。这将输出计算与状态更新分离，并解决了当前 SSM 解码中的三个挑战：状态流量、推测解码的回滚，以及量化状态中的误差累积。
ReplaySSM 是 NVIDIA 与普林斯顿大学 Dao AI Lab 的联合工作。论文《ReplaySSM: Cache SSM Inputs, Not State》详细介绍了全部三项优势并给出了完整评估。该论文可在 arXiv 上获取，代码在 GitHub 上。它已集成到 NVIDIA TensorRT（NVIDIA 深度学习推理优化库）-LLM（大语言模型）、vLLM 和 FlashInfer 中，并在 GLM-5.3-Flash 和 Kimi-K3 中投入生产运行。
为什么 SSM 层可能成为解码瓶颈
以 NVIDIA Nemotron 3 Ultra 中使用的 Mamba-2 层为例。它通过以下递推关系解码一个 token：
其中 S t S_t 是状态，α t \alpha_t 是数据相关的衰减，v t v_t、k t k_t 和 q t q_t 是该 token 的输入。每一步从高带宽内存（HBM）加载 S t − 1 S_{t-1}，应用衰减和秩-1 更新，用 q t q_t 计算输出，并将 S t S_t 存回。GDN 和 KDA 增加了 delta 规则校正，但它们的更新可以写成相同的外积形式，因此以下内容同样适用于它们。
在解码过程中，状态承担两项任务。它被读取以产生当前 token 的输出，并被更新以成为下一个 token 的状态。传统解码同时完成这两项：每一步更新状态，从中读取输出，并将其存储以供下一步使用。所有三个挑战都源于这种融合。
状态内存流量
SSM 解码步骤在加载和存储完整状态时，对其执行的算术运算很少，因此该内核受 HBM 带宽而非计算能力限制。在服务批大小下，这使得 SSM 层（而非注意力）成为解码步骤中最大的组成部分，在单块 NVIDIA Blackwell Ultra B300 GPU（图形处理器） 上以批大小 256 运行 Nemotron-3-Super-NVFP4 时，上下文长度可达 100K token（图 2，左）。
有两个因素驱动了这一点：
层数：该模型的 SSM 层数量是注意力层的五倍。
状态大小：一层的 FP32 状态为每个序列 4 MiB，而一个 FP8 KV 条目为每个 token 0.5 KiB，因此一个 SSM 层的状态大小相当于一个注意力层在 8K token 下的 KV 缓存大小。
投机解码的回滚
投机解码由一个草稿模型提出 T 个 token，通过目标模型一次前向传播进行验证，并回滚到最后一个被接受的 token。注意力机制通过移动 KV 缓存指针天然支持回滚。但一旦 SSM 将草稿内容汇总到其状态中，该更新便没有简单的逆操作，因此回滚变得代价高昂。诸如 vLLM 之类的服务引擎通过在每个草稿 token 之后存储状态快照，并选择与已接受前缀匹配的快照来绕过这一问题。在内存受限路径上，这意味着每步一次读取和 T 次写入，而且这些快照限制了 GPU 内存中可容纳的请求数量。因此，投机解码——对 Transformer（深度学习中的注意力模型架构） 而言是一种可靠的加速手段——在服务批大小下对 SSM 几乎毫无帮助。
状态量化
量化后的状态产生输出，并作为下一次更新的输入，因此其舍入误差在每一步都会重新进入递归。KV 缓存没有这样的循环：量化条目会被多次读取，但永远不会被写回自身，这就是为什么 KV 缓存比 SSM 状态更能容忍较低精度。例如，Qwen3.5-122B-A10B-NVFP4 对其权重使用 NVIDIA NVFP4 4 位浮点格式，并搭配 FP8 KV 缓存和 FP32 SSM 状态。
ReplaySSM 如何表示状态
SSM 递归可以在 h 个 token 的窗口上展开：
其中 S c k p t S_{\mathrm{ckpt}} 是窗口之前的状态，γ \gamma 是累积衰减。传统解码逐步遵循左侧。ReplaySSM 使用右侧：它将状态表示为一个检查点加上最近输入记录的小缓冲区（对于 Mamba-2 为 α j \alpha_j 、 v j v_j 、 k j k_j ）。
在大多数步骤中（我们称之为非刷新步骤），它将当前输入追加到缓冲区，并根据检查点和缓冲区计算输出，而不存储状态。当 L 条记录的缓冲区填满时，一个刷新步骤会将缓冲的输入折叠到检查点中，将新检查点一次性写入 HBM，并清空缓冲区。其输出与传统解码一致，仅存在浮点差异。ReplaySSM 的方法可以概括为三个问题：
ReplaySSM 做了什么？
它缓存最近的 SSM 输入，并将状态表示为一个检查点加上这些输入。
这改变了什么？
它将输出计算与状态更新分离：每一步都计算输出，而检查点每个窗口更新一次。
这实现了什么？
它减少了状态内存流量，简化了推测性回滚，并支持低精度状态读取，而不会通过循环更新传播其舍入误差。
我们如何评估 ReplaySSM
我们评估了来自两个系列的四个混合模型，Nemotron 3 使用 Mamba-2 层，Qwen3.5 使用 GDN 层，从单个 NVIDIA H100 Tensor Core GPU 上的 4B 稠密模型到单个 NVIDIA Blackwell Ultra B300 GPU 上的 120B 级混合专家（MoE）模型（表 1）。我们在 vLLM 中实现了 ReplaySSM，并将其与 vLLM 自有的内核在其最佳启动配置下进行了比较。
性能测试将 SSM 状态保持为 FP32，这是这些模型的默认精度，并使用 L = 16 的窗口，这是所有四个模型的最佳性能窗口大小。推测解码测试使用每个模型的多 token 预测（MTP）头来生成草稿 token。
| 模型 | SSM 层 | 权重 | GPU | SSM / 注意力层 | | --- | --- | --- | --- | --- | | Nemotron-3-Nano-4B | Mamba-2 | BF16 | 1×H100 | 21 / 4 | | Nemotron-3-Super-120B-A12B | Mamba-2 | NVFP4 | 1×B300 | 40 / 8 | | Qwen3.5-4B | GDN | BF16 | 1×H100 | 24 / 8 | | Qwen3.5-122B-A10B | GDN | NVFP4 | 1×B300 | 36 / 12 |
优势 1：减少状态内存流量
在 L 个 token 的窗口内，传统解码读取状态 L 次并写入 L 次。ReplaySSM 仍然在每一步读取检查点，因为输出需要它，但只写入一次，即在 flush 期间。非 flush 步骤只有一项工作——计算输出——并且只写入新的输入记录，即几个向量而不是完整矩阵。这大约将主要状态流量减半。
在服务批处理大小下，内存带宽限制了传统内核的性能，因此减少状态内存流量可以加速解码。在批处理大小为 1 时，其他延迟来源占主导，ReplaySSM 在这些测试中没有提供加速。在批处理大小为 256 时，内核运行速度提高 1.46× 到 1.80×，解码步骤在四个模型上运行速度提高 1.20× 到 1.47×（图 3）。端到端增益随批处理大小增长，因为 SSM 在步骤中的占比增加。
优势 2：简化回滚并加速验证
在推测解码中，ReplaySSM 在验证之前将所有 T 个草稿 token 的输入追加到输入缓冲区，然后通过推进缓冲区指针来提交接受的 prefix。与 KV 缓存一样，指针之外的记录被忽略，下一次追加会覆盖它们。不需要为每个草稿 token 保存状态快照，也不需要恢复状态（图 2，中间）。flush 仅将已提交的记录折叠到检查点中，因此它从不等待接受的长度。
这种分离带来了另一个好处。传统解码必须为每个 token 产生两样东西：输出和更新后的状态，因为下一个 token 会消费该状态。ReplaySSM 中的非刷新步骤只产生输出，而仅计算输出开启了一种不同的计算方法。对于从零状态开始的一次更新，o = ( v k ⊤ ) q = v ( k ⊤ q ) o = (v k^\top)q = v(k^\top q) 。左侧表达式构建一个 d v × d k d_v \times d_k 状态并读取它。右侧表达式形成一个标量并缩放一个向量，而不构建状态。对于针对 h 条已提交记录的 T 个草稿，同样的思路给出以下表达式：
其中 D 保存草稿位置的累积衰减，W 是因果掩码的衰减矩阵。我们称之为仅输出评估。检查点对所有 T 个查询只使用一次，其余部分都是对缓存键、缓存值和草稿查询的两个小矩阵乘积。不形成中间状态，这节省了片上内存和 FLOPs，并且 Tensor Cores 可以加速这些计算。此外，K ⊤ Q K^\top Q 由组内所有值头共享。这与预填充内核已经使用的分块计算相同，以缓冲区和草稿作为块。GDN 也可以使用这种形式。仅输出评估也适用于标准解码，但它在验证中最重要，因为 T 会倍增每步的算术量。
图 4 展示了在两个 120B 级模型上，使用 GSM8K 提示、T = 4 时的推测服务吞吐量。在批量大小为 512 时，ReplaySSM 的吞吐量分别是 vLLM 推测解码的 1.85 倍和 2.10 倍，且接受长度匹配。
这一增益有两个来源。
更快的验证：vLLM 的内核为每个 token 存储一个快照，因此其延迟随 T 增长，在 T = 8 时几乎是 vLLM 标准解码步骤的五倍。ReplaySSM 在大多数步骤中不写入状态，因此它在大约一个标准步骤的时间内验证八个 token，包括刷新在内（图 5）。
更高的并发性：草稿状态快照消耗 GPU 内存并限制并发性。传统推测在批量大小超过 128 后不再获得增益，并且在 Qwen3.5-122B-A10B-NVFP4 上，在批量大小为 512 时低于标准解码，而 ReplaySSM 则随批量大小持续扩展。
优势 3：通过分离精度限制误差累积
较低精度的状态需要更少的字节，这在内存带宽限制内核性能时减少延迟。在 Mamba-2 上，读取 NVFP4 状态时，非刷新内核的运行速度是读取 FP32 状态时的 3.8 倍。
ReplaySSM 以两种方式改变量化。首先，检查点每个窗口写入一次，而不是每步写入一次，因此单个低精度检查点每 L 步舍入一次。Nemotron 3 Ultra 报告探讨了这种周期性检查点方法。它降低了舍入频率，但舍入误差仍会进入后续更新。
其次，非刷新步骤读取状态仅用于计算输出，而刷新步骤更新用于后续状态更新的检查点。这两项工作可以使用不同的精度。我们为递归保留一个高精度检查点，并为非刷新读取保留一个低精度副本。刷新从高精度检查点开始，纳入缓冲输入，并将结果独立地转换为两种格式。低精度副本永远不会被读回递归中，因此其舍入误差仅影响当前窗口的输出，而绝不会在递归中传播。在刷新期间写入额外的低精度副本减少了每个非刷新步骤读取的数据量。接受、回滚和缓冲区管理保持不变（图 2，右）。
图 6 测量了在 PG-19 文本上、教师强制下，相对于 FP32 状态解码的 Kullback–Leibler (KL) 散度所表示的输出误差，直到生成 128K 个 token。共享的 FP8 或 NVFP4 检查点会随长度累积误差。将递归保持在 FP32 可消除大部分增长：在 Qwen3.5-4B 上，128K token 时，使用 FP32 递归的 NVFP4 读取将 KL 从 0.136 降至 0.0038。剩下的就是读取误差，它不随长度增长。
使用随机舍入 (SR) 的 FP16 递归与 FP32 紧密吻合，因此半精度递归检查点就足够了。对读取副本进行更好的量化，例如沿状态和通道维度使用单独的缩放因子，可以进一步减少剩余的读取误差。
我们还检查了 KL 差距是否会在下游准确率中显现。表 2 显示了 GPQA-Diamond，这是我们运行的三个基准中生成最长的基准。共享的 NVFP4 检查点在四个模型上损失了 3 到 31 个百分点的准确率。使用 FP16 或 FP8 递归的 NVFP4 读取恢复了大部分准确率，并使四个模型中的三个保持在 FP32 约 2 个百分点以内。
| GPQA-Diamond，相对 FP32 的变化（百分点） | FP32 准确率 | 共享 NVFP4 checkpoint | NVFP4 读取，FP16+SR 递推 | NVFP4 读取，FP8+SR 递推 | | --- | --- | --- | --- | --- | | Nemotron-3-Nano-4B | 56.7 | −31.3 | −6.7 | −5.6 | | Nemotron-3-Super-NVFP4 | 79.3 | −12.5 | −1.4 | −1.4 | | Qwen3.5-4B | 76.3 | −3.2 | +0.7 | +1.5 | | Qwen3.5-122B-A10B-NVFP4 | 82.7 | −5.2 | −2.0 | −1.7 |
图 7 展示了在 GPQA-Diamond 上准确率损失随生成长度的变化情况。在三个基准测试中：
GPQA-Diamond：在 Nemotron-3-Super-NVFP4 上，共享 checkpoint 在较短生成中与 FP32 准确率保持一致，但随着生成长度增加而出现偏离。分离精度配置则始终接近 FP32。
MMLU-Pro 和 Math-L5：生成长度通常短于 GPQA-Diamond。四个模型中有三个在使用共享 NVFP4 checkpoint 时准确率损失最多为 3 个百分点，而分离 checkpoint 配置平均与 FP32 的差距保持在半个百分点以内。
论文包含了全部三个基准测试的结果以及一项长输入 RULER 测试。
开始使用 ReplaySSM
ReplaySSM 已投入生产使用。GLM-5.3-Flash 使用 ReplaySSM 提供服务，Z.ai 的基础设施报告描述了将其添加到 KDA 解码内核中。Kimi-K3 在其 KDA 层中缓存草稿 token 的 SSM 输入以进行验证，并将 ReplaySSM 列为一种并行方法。ReplaySSM 也已集成到开源推理框架中，如 TensorRT-LLM、vLLM、SGLang 和 FlashInfer。
如需了解更多信息，请查看以下资源：
阅读论文：ReplaySSM: Cache SSM Inputs, Not State
试用代码：github.com/Johnny-Liou/ReplaySSM
在 TensorRT-LLM 中使用 ReplaySSM 服务混合模型：PR 14203 , PR 16464
在 vLLM 中使用 ReplaySSM 服务混合模型：PR 48018 , PR 52506 , PR 52928
在 FlashInfer 中使用 ReplaySSM 内核：PR 3975 , PR 4081
阅读 Dao AI Lab 的早期文章，了解标准解码和推测解码的更详细讲解
转载声明：本文为转载整理，原文版权归原作者及原站点所有。
原文作者：Michelle Horton；如有侵权，请联系管理员删除。
