混合模型将状态空间模型（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）缓存。然而在服务批大小下，当上下文达到 100K token 时，SSM 层（而非注意力）可能成为延迟瓶颈。SSM 层还使两种能可靠加速注意力解码的技术变得复杂：推测解码和量化。
ReplaySSM 让 SSM 层重新成为快速的部分。它将 SSM 状态表示为一个检查点加上近期输入的缓冲区。这将输出计算与状态更新分离，并解决了当前 SSM 解码中的三个挑战：状态流量、推测解码的回滚，以及量化状态中的误差累积。
ReplaySSM 是 NVIDIA 与普林斯顿大学 Dao AI Lab 的联合工作。论文 ReplaySSM: Cache SSM Inputs, Not State 详细介绍了全部三项优势并给出了完整评估。该论文可在 arXiv 上获取，代码在 GitHub 上。它已集成到 NVIDIA TensorRT-LLM、vLLM 和 FlashInfer 中，并在 GLM-5.3-Flash 和 Kimi-K3 中投入生产运行。
图 1。传统解码在每一步都加载并存储完整状态。ReplaySSM 缓存近期输入，从检查点加缓冲区计算输出，并仅在缓冲区填满时写入检查点
为什么 SSM 层可能成为解码瓶颈
以 NVIDIA Nemotron 3 Ultra 中使用的 Mamba-2 层为例。它通过以下递推关系解码一个 token：
其中 是状态， 是数据相关的衰减，而 、 和 是该 token 的输入。每一步从高带宽内存（HBM）加载 ，应用衰减和秩-1 更新，用 计算输出，并将 存回。GDN 和 KDA 增加了 delta 规则修正，但它们的更新可以写成相同的外积形式，因此以下内容同样适用于它们。
在解码过程中，状态承担两项工作。它被读取以产生当前 token 的输出，并被更新以成为下一个 token 的状态。传统解码同时完成这两项工作：每一步更新状态，从中读取输出，并为下一步存储它。所有三个挑战都源于这种融合。
图 2. 传统 SSM 解码的三个挑战以及 ReplaySSM 对每个挑战的应对方案。左图：在单块 NVIDIA Blackwell Ultra B300 GPU 上，批大小为 256 时 Nemotron-3-Super-NVFP4 的解码步延迟分解
状态内存流量
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 的窗口上展开：
其中是窗口之前的状态，是累积衰减。传统解码逐步遵循左侧。ReplaySSM 使用右侧：它将状态表示为一个检查点加上一个由近期输入记录（对于 Mamba-2 为 、 、 ）组成的小缓冲区。
在大多数步骤中——我们称之为非刷新步骤——它将当前输入追加到缓冲区，并根据检查点和缓冲区计算输出，而不存储状态。当 L 条记录的缓冲区填满时，一个刷新步骤会将缓冲的输入折叠进检查点，将新检查点一次性写入 HBM，并清空缓冲区。其输出与传统解码一致，仅存在浮点差异。ReplaySSM 的方法可以概括为三个问题：
ReplaySSM 做什么？
它缓存近期的 SSM 输入，并将状态表示为一个检查点加上这些输入。
这改变了什么？
它将输出计算与状态更新分离：输出每一步都计算，而检查点每个窗口更新一次。
这实现了什么？
它减少了状态内存流量，简化了推测性回滚，并支持低精度状态读取，而不会通过循环更新传播其舍入误差。
我们如何评估 ReplaySSM
我们评估了来自两个系列的四个混合模型：具有 Mamba-2 层的 Nemotron 3 和具有 GDN 层的 Qwen3.5，从单个 NVIDIA H100 Tensor Core GPU 上的 4B 密集模型到单个 NVIDIA Blackwell Ultra B300 GPU 上的 120B 级混合专家（MoE）模型（表 1）。我们在 vLLM 中实现了 ReplaySSM，并将其与 vLLM 自有的内核在其最佳启动配置下进行了比较。
性能运行将 SSM 状态保持在 FP32（这些模型的默认精度），并使用 L = 16 的窗口，这是所有四个模型的最佳性能窗口大小。推测解码测试使用每个模型的多令牌预测（MTP）头来生成草稿令牌。
| 模型 | SSM 层 | 权重 | GPU | SSM / 注意力层 |
|---|
Nemotron-3-Super-120B-A12BQwen3.5-4BQwen3.5-122B-A10B表 1. 评估的模型和硬件。每个模型在一个 GPU 上运行，无张量并行
优势 1：减少状态内存流量
在 L 个令牌的窗口内，传统解码读取状态 L 次并写入 L 次。ReplaySSM 仍然在每一步读取检查点，因为输出需要它，但只在刷新期间写入一次。非刷新步骤只有一项工作——计算输出——并且只写入新的输入记录，即几个向量而不是完整矩阵。这大约将主要状态流量减半。
在服务批处理大小下，内存带宽限制了传统内核的性能，因此减少状态内存流量可以加速解码。在批处理大小为 1 时，其他延迟来源占主导，ReplaySSM 在这些测试中没有提供加速。在批处理大小为 256 时，内核运行速度提高 1.46 倍至 1.80 倍，解码步骤在四个模型上运行速度提高 1.20 倍至 1.47 倍（图 3）。端到端增益随着批处理大小的增加而增长，因为 SSM 在步骤中的份额增加。
图 3. 在批处理大小为 256 时相对于调优后的 vLLM 的标准解码加速。内核数字包括刷新成本
优势 2：简化回滚并加速验证
在推测解码中，ReplaySSM 在验证之前将所有 T 个草稿令牌的输入追加到输入缓冲区，然后通过推进缓冲区指针来提交接受的令牌前缀。与 KV 缓存一样，指针之后的记录被忽略，下一次追加会覆盖它们。不需要为每个草稿令牌保存状态快照，也不需要恢复状态（图 2，中间）。刷新仅将已提交的记录合并到检查点中，因此它从不等待接受的长度。
这种分离带来了另一个好处。传统解码每个 token 必须产生两样东西：输出和更新后的状态，因为下一个 token 要消费该状态。ReplaySSM 中的非刷新步骤只产生输出，而只计算输出开辟了一种不同的计算方法。对于从零状态出发的一次更新，。左侧的表达式构建一个状态并读取它。右侧的表达式形成一个标量并缩放一个向量，而不构建状态。对于针对 h 条已提交记录的 T 个草稿，同样的思路给出以下表达式：
其中 D 保存草稿位置的累积衰减，W 是一个因果掩码的衰减矩阵。我们称之为仅输出求值。检查点对所有 T 个查询只使用一次，其余一切都只是对缓存键、缓存值和草稿查询的两个小矩阵乘积。不形成任何中间状态，这节省了片上内存和 FLOPs，而且 Tensor Cores 可以加速这些计算。此外，被一个组中的所有值头共享。这与预填充内核已经使用的分块计算相同，只是把缓冲区和草稿作为块。GDN 也可以使用这种形式。仅输出求值也适用于标准解码，但它在验证中最重要，因为 T 会使每步的算术量成倍增加。
图 4 展示了在两个 120B 级模型上、T = 4 时 GSM8K 提示上的推测服务吞吐量。在批大小为 512 时，ReplaySSM 的吞吐量分别是 vLLM 推测解码的 1.85 倍和 2.10 倍，且接受长度相当。
图 4. GSM8K 提示上的推测服务吞吐量，T = 4，L = 16，FP32 状态，每个模型一块 B300
这一收益有两个来源。
- 更快的验证：vLLM 的内核为每个 token 存储一个快照，因此其延迟随 T 增长，在 T = 8 时几乎是 vLLM 标准解码步骤的五倍。ReplaySSM 在大多数步骤上不写入状态，因此它验证八个 token 的时间大约相当于一个标准步骤（包括刷新）（图 5）。- 更高的并发度：草稿状态快照消耗 GPU 内存并限制并发度。传统推测在批大小超过 128 后不再有收益，而在 Qwen3.5-122B-A10B-NVFP4 上，它在批大小为 512 时低于标准解码，而 ReplaySSM 则随批大小持续扩展。
图 5. Nemotron-3-Super-NVFP4 上批大小为 256 时的验证内核延迟，相对于 vLLM 的标准解码步骤，包括刷新
好处 3：用分离的精度限制误差累积
较低精度的状态需要更少的字节，这在内存带宽限制内核性能时降低延迟。在 Mamba-2 上，读取 NVFP4 状态时非刷新内核的运行速度是读取 FP32 状态时的 3.8 倍。
ReplaySSM 从两个方面改变了量化方式。第一，检查点按窗口写入一次，而不是每步写入一次，因此单个低精度检查点每 L 步才舍入一次。Nemotron 3 Ultra 报告 探讨了这种周期性检查点方法。它降低了舍入频率，但舍入误差仍会进入后续更新。
第二，非刷新步骤只读取状态以计算输出，而刷新步骤会更新用于后续状态更新的检查点。这两项工作可以使用不同的精度。我们为递归保留一个高精度检查点，并为非刷新读取保留一个低精度副本。刷新从高精度检查点开始，纳入缓冲的输入，并将结果独立地转换为两种格式。低精度副本永远不会被读回递归中，因此其舍入误差只影响当前窗口的输出，绝不会在递归中传播。在刷新期间写入额外的低精度副本，减少了每个非刷新步骤读取的数据量。接受、回滚和缓冲区管理保持不变（图 2，右）。
图 6 测量了在 PG-19 文本上、教师强制下、相对于 FP32 状态解码的输出误差，以 Kullback–Leibler（KL）散度表示，生成 token 数最多到 128K。共享的 FP8 或 NVFP4 检查点会随长度累积误差。将递归保持在 FP32 可消除大部分增长：在 Qwen3.5-4B 上、128K token 时，使用 FP32 递归的 NVFP4 读取将 KL 从 0.136 降至 0.0038。剩下的是读取误差，它不随长度增长。
使用随机舍入（SR）的 FP16 递归与 FP32 非常接近，因此半精度递归检查点就足够了。对读取副本进行更好的量化，例如沿状态和通道维度使用单独的缩放因子，可以进一步降低剩余的读取误差。
图 6. Qwen3.5-4B 上相对于传统 FP32 状态解码的输出 KL，L = 16。两项工作（输出计算和递归更新）使用分离的精度，消除了大部分累积误差
我们还检查了 KL 差距是否体现在下游准确率上。表 2 显示了 GPQA-Diamond，这是我们运行的三个基准中生成最长的。共享的 NVFP4 检查点在四个模型上损失了 3 到 31 个百分点的准确率。使用 FP16 或 FP8 递归的 NVFP4 读取恢复了大部分准确率，并使四个模型中的三个与 FP32 的差距保持在约 2 个百分点以内。
| GPQA-Diamond，相对于 FP32 的变化（百分点） | FP32 准确率 | 共享 NVFP4 检查点 | 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 |
表 2. 使用 NVFP4 读取检查点的 GPQA-Diamond 准确率，表示为相对于每个模型 FP32 参考值的变化。共享检查点使用随机舍入。标准解码，L = 16
图 7 显示了在 GPQA-Diamond 上准确率损失如何随生成长度变化。在三个基准上：
GPQA-Diamond：在 Nemotron-3-Super-NVFP4 上，共享检查点在较短的生成中与 FP32 准确率保持一致，但随着生成长度增加而偏离。分离精度配置则始终接近 FP32。MMLU-Pro 和 Math-L5：生成内容通常比 GPQA-Diamond 中的更短。四个模型中有三个在使用共享 NVFP4 检查点时准确率损失最多 3 个百分点，而分离检查点配置平均与 FP32 的差距保持在半个百分点以内。
论文包含了全部三个基准测试以及长输入 RULER 测试的结果。
图 7. 在 Nemotron-3-Super-NVFP4 上按 FP32 参考生成长度四分位数划分的 GPQA-Diamond 准确率
开始使用 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 内核：
Dao AI Lab 的早期文章 以获取关于标准和推测解码的更详细讲解