混合模型将状态空间模型（SSM）层与少量注意力层交替排列。NVIDIA Nemotron 3、Qwen3.5 到 Qwen3.8、Kimi-K3 以及 GLM-5.3-Flash 均采用这种架构，其 SSM 层分别使用 Mamba-2、门控 DeltaNet（GDN）或 Kimi Delta Attention（KDA）实现。
SSM 层的设计初衷是让解码更快：它们使用固定大小的状态，取代随上下文长度增长的键值（KV）缓存。但在实际服务批处理规模下，SSM 层而非注意力层，往往才是上下文长度达到 10 万 Token 以内时的延迟瓶颈所在。此外，SSM 层还给两项能够稳定提升注意力解码速度的技术——推测解码和量化——带来了复杂性问题。
ReplaySSM 让 SSM 层重新成为解码中的"快速部分"。它将 SSM 状态表示为一个检查点加上一个最近输入的缓冲区，这种做法将输出计算与状态更新分离开，从而解决了当前 SSM 解码中的三大难题：状态数据传输量、推测解码的回滚机制，以及量化状态中的误差累积。
ReplaySSM 是 NVIDIA 与普林斯顿大学 Dao AI 实验室的联合研究成果。论文《ReplaySSM：缓存 SSM 输入，而非状态》详细阐述了这三方面优势并呈现了完整的评估结果，论文已发布在 arXiv，代码托管于 GitHub。该技术已集成至 NVIDIA TensorRT-LLM、vLLM 和 FlashInfer 中，同时也在 GLM-5.3-Flash 和 Kimi-K3 的生产环境中实际运行。
为什么 SSM 层会成为解码瓶颈
以 NVIDIA Nemotron 3 Ultra 中使用的 Mamba-2 层为例，它通过以下递推公式解码一个 Token：
St = αt·St-1 + vt·kt?， ot = St·qt
其中 St 是状态，αt 是数据相关的衰减系数，vt、kt 和 qt 是该 Token 的输入。每一步都需要从高带宽内存（HBM）中加载 St-1，应用衰减并进行秩-1 更新，用 qt 计算输出，再将 St 存回内存。GDN 和 KDA 增加了 delta 规则修正项，但它们的更新同样可以写成相同的外积形式，因此以下讨论同样适用于它们。
在解码过程中，状态要承担两项任务：被读取以生成当前 Token 的输出，以及被更新成为下一个 Token 的状态。传统解码方式将这两项任务一起完成：每一步都更新状态、从中读出输出，再将其存储供下一步使用。三大难题正是源于这种任务融合。
状态内存流量
一个 SSM 解码步骤要加载并存储整个状态，但对其执行的运算量却很少，因此该计算核心受限于 HBM 带宽，而非计算能力。在服务批处理规模下，这使得 SSM 层而非注意力层，在上下文长度达到 10 万 Token 时成为解码步骤中占比最大的部分（在一块 NVIDIA Blackwell Ultra B300 GPU 上、批大小为 256 时运行 Nemotron-3-Super-NVFP4 的测试结果显示，见图 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 位浮点格式，KV 缓存为 FP8，而 SSM 状态为 FP32。
ReplaySSM 如何表示状态
SSM 的递推过程可以在一个长度为 h 个 Token 的窗口内展开：
St = γh·Sckpt + Σj=1h γj+1:h·vj·kj?
其中 Sckpt 是窗口之前的状态，γ 是累积衰减系数。传统解码方式按照左侧的公式逐步执行，而 ReplaySSM 使用右侧的表达方式：将状态表示为一个检查点，加上一个记录最近输入（针对 Mamba-2 的 αj、vj、kj）的小型缓冲区。
在大多数步骤（我们称之为"非刷新步骤"）中，系统将当前输入追加到缓冲区中，并直接根据检查点和缓冲区计算输出，而不存储状态。当缓冲区中累积了 L 条记录后填满时，会触发一次"刷新步骤"，将缓冲的输入折叠进检查点，将新检查点一次性写入 HBM，再清空缓冲区。这样得到的输出结果（在浮点误差范围内）与传统解码方式一致。ReplaySSM 的方法可以归纳为三个问题：
1. ReplaySSM 做了什么？它缓存最近的 SSM 输入，并将状态表示为检查点加这些输入的组合。
2. 这改变了什么？它将输出计算与状态更新分离开来：每一步都计算输出，而检查点只在每个窗口内更新一次。
3. 这实现了什么？它减少了状态内存流量，简化了推测解码的回滚机制，并支持在不将舍入误差传播到递推更新中的前提下进行低精度状态读取。
评估方法
我们在两个模型系列的四个混合模型上进行了评估：带有 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，这是对全部四个模型表现最佳的窗口大小。推测解码测试使用每个模型的多 Token 预测（MTP）头来生成草稿 Token。
优势一：减少状态内存流量
在一个长度为 L 个 Token 的窗口内，传统解码方式要对状态读取 L 次、写入 L 次。而 ReplaySSM 每一步仍需读取检查点（因为计算输出需要用到它），但只在刷新时写入一次。一个非刷新步骤只需完成一项任务——计算输出，且只需写入新的输入记录（几个向量），而非完整的矩阵。这大致将占主导的状态流量减半。
在服务批处理规模下，内存带宽限制了传统计算核心的性能，因此减少状态内存流量能加速解码。在批大小为 1 时，其他延迟来源占主导，ReplaySSM 在这些测试中没有提速效果。而在批大小为 256 时，计算核心运行速度提升 1.46 倍至 1.80 倍，解码步骤整体速度提升 1.20 倍至 1.47 倍，覆盖全部四个模型（见图 3）。随着批大小增加，端到端收益也随之增长，因为 SSM 在整个解码步骤中的占比也随之增大。
优势二：简化回滚并加速验证
在推测解码中，ReplaySSM 在验证之前，将全部 T 个草稿 Token 的输入追加到输入缓冲区中，验证之后通过移动缓冲区指针来提交被接受的前缀。与 KV 缓存类似，指针之后的记录会被忽略，下一次追加操作会将其覆盖。无需为每个草稿 Token 保存状态快照，也无需恢复状态（见图 2 中）。刷新操作只会将已提交的记录折叠进检查点，因此永远不需要等待被接受的长度确定。
这种任务分离还带来了另一项好处。传统解码方式每个 Token 都必须产生两项结果——输出和更新后的状态，因为下一个 Token 要消耗该状态。而 ReplaySSM 中的非刷新步骤只产生输出，而只计算输出这一点开启了另一种计算方式。以从零状态开始的一次更新为例：o=(vk?)q=v(k?q)。左侧表达式需要构建一个 dv×dk 大小的状态并读取它，而右侧表达式只需形成一个标量并对一个向量进行缩放，无需构建状态。对于针对 h 条已提交记录验证 T 个草稿的情形，同样的思路可以得到如下表达式：
O = Sckpt·Q·D + V(W⊙(K?Q))
其中 D 保存草稿位置的累积衰减系数，W 是经过因果掩码处理的衰减矩阵。我们将这种方式称为"仅输出评估"。检查点对所有 T 个查询只使用一次，剩余的计算只是针对已缓存的键、已缓存的值以及草稿查询进行的两次小规模矩阵乘法。整个过程不会形成任何中间状态，从而节省了芯片内存和浮点运算量，Tensor Core 还能进一步加速这些计算。此外，K?Q 可以被同一组内的所有值头共享使用。这与预填充（prefill）计算核心已经采用的分块计算方式相同，只是将缓冲区和草稿作为分块单位。GDN 同样可以采用这种形式。"仅输出评估"同样适用于标准解码，但在验证环节中意义最大，因为 T 会使每一步的运算量成倍增加。
图 4 展示了在 GSM8K 提示集上、T=4 条件下两个 120B 级别模型的推测解码服务吞吐量。在批大小为 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 则能随批大小持续扩展性能。
优势三：通过分离精度限制误差累积
更低精度的状态所需字节数更少，在内存带宽限制计算核心性能的情况下，这能够降低延迟。在 Mamba-2 上，当读取 NVFP4 状态而非 FP32 状态时，非刷新计算核心的运行速度提升了 3.8 倍。
ReplaySSM 通过两种方式改进了量化。首先，检查点改为每个窗口写入一次，而非每步写入一次，这意味着单个低精度检查点每隔 L 步才进行一次舍入。Nemotron 3 Ultra 的报告中也探讨了这种周期性检查点方案，它降低了舍入频率，但舍入误差仍会进入后续的更新中。
其次，非刷新步骤只在计算输出时读取状态，而刷新步骤则更新用于后续状态更新的检查点。这两项任务可以使用不同的精度。我们为递推过程保留一份高精度检查点，同时为非刷新读取保留一份低精度副本。每次刷新都从高精度检查点开始，纳入缓冲的输入，并将结果独立地分别转换为两种格式。低精度副本永远不会被读回递推过程中，因此其舍入误差只影响当前窗口内的输出，永远不会在递推过程中传播累积。在刷新时额外写入低精度副本，减少了每个非刷新步骤需要读取的数据量。接受、回滚和缓冲区管理机制均保持不变（见图 2 右）。
图 6 测量了在 PG-19 文本上使用教师强制（teacher forcing）方式、生成多达 128K 个 Token 时，输出误差相对 FP32 状态解码的 KL 散度（Kullback–Leibler divergence）。共享的 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 个百分点以内。
图 7 展示了在 GPQA-Diamond 上，准确率损失如何随生成长度变化。综合三个基准来看：
GPQA-Diamond：在 Nemotron-3-Super-NVFP4 上，共享检查点在较短生成长度下能跟上 FP32 的准确率水平，但随生成长度增长会逐渐偏离。而分离精度配置则在整个过程中始终贴近 FP32 表现。
MMLU-Pro 和 Math-L5：生成长度通常比 GPQA-Diamond 更短。四个模型中有三个在采用共享 NVFP4 检查点时，准确率损失最多为 3 个百分点，而分离检查点配置的平均准确率差距保持在 FP32 基准的 0.5 个百分点以内。
论文中包含了全部三个基准测试的完整结果，以及一项长输入 RULER 测试结果。
开始使用 ReplaySSM
ReplaySSM 已投入生产使用。GLM-5.3-Flash 已采用 ReplaySSM 进行服务，Z.ai 的基础设施报告中描述了将其加入 KDA 解码计算核心的过程。Kimi-K3 在其 KDA 层中缓存草稿 Token 的 SSM 输入以供验证使用，并将 ReplaySSM 列为同期提出的类似方案。ReplaySSM 还已集成至 TensorRT-LLM、vLLM、SGLang 和 FlashInfer 等开源推理框架中。
Q&A
Q1：ReplaySSM 是什么？它解决了什么问题？
A：ReplaySSM 是由 NVIDIA 和普林斯顿大学 Dao AI 实验室联合提出的技术，它将 SSM 状态表示为检查点加最近输入缓冲区的形式，解决了混合模型中 SSM 层解码时的状态内存流量过大、推测解码回滚代价高、以及状态量化误差累积这三大难题。
Q2：ReplaySSM 能带来多大的性能提升？
A：在批大小为 256 时，ReplaySSM 使计算核心运行速度提升 1.46 倍至 1.80 倍，解码步骤整体提速 1.20 倍至 1.47 倍。在推测解码场景中，批大小为 512 时吞吐量可达传统方案的 1.85 倍至 2.10 倍。
Q3：ReplaySSM 目前已经在哪些产品中使用？
A：ReplaySSM 已在 GLM-5.3-Flash 和 Kimi-K3 的生产环境中实际运行，并已集成至 NVIDIA TensorRT-LLM、vLLM、SGLang 和 FlashInfer 等主流开源推理框架中。
热门跟贴