ReplaySSM: Cache SSM Inputs, Not State

发表时间: 2026-10 · arXiv:2610.13778 (Princeton University & NVIDIA)

原文: https://arxiv.org/abs/2610.13778

Ze-Wei Liou (Princeton University), Harris Nover (NVIDIA), Amey Naik (NVIDIA), Andrii Skliar (NVIDIA), Igor Shovkun (NVIDIA), Ido Shahaf (NVIDIA), Tri Dao (Princeton University)

速读

一句话结论 ReplaySSM 通过将状态空间模型（SSM）的隐藏状态展开为“历史检查点加近期输入缓存”的表示，解耦了输出计算与状态更新，在几乎不改变数值结果的前提下将常规解码加速了 1.20–1.47 倍，并将投机解码吞吐提升约 1.8 倍。

要解决什么问题 在常规的混合架构（SSM 混合注意力）解码中，系统每一步都必须完整加载、更新并写回定长的递归状态，这种激进的状态管理引入了三个严重卡点。第一，计算强度极低：每个解码步为了极少量的运算搬运几兆字节的完整状态，算力利用率不足 1 FLOP/字节，导致服务批大小下即便上下文达 100K tokens，SSM 层的显存 I/O 依然是整体时延的最大瓶颈。第二，投机解码无法廉价回滚：当目标模型拒绝多余的草稿 token 时，不可逆的递归更新迫使服务系统为每个草稿保留状态快照，引入庞大的状态写流量并限制并发，使投机解码在服务端难以产生收益。第三，状态量化误差随时间步累积：若采用低精度状态存储，单步的舍入误差会重新喂入后续的循环更新中，导致长文本生成时输出分布严重漂移，下游任务准确率大幅崩塌。

怎么做的 核心思路是展开递归式：当前状态完全等价于一个历史检查点衰减后的结果加上近期输入的加权和，因此不必在每个时间步都物化并写回完整状态。ReplaySSM 将状态拆分为常驻显存的检查点 $S_{\mathrm{ckpt}}$ 与一个容量为 $L$ 的循环输入缓冲区 $B$。通过这一解耦，系统将绝大部分时间步变成非刷写步（non-flush step），仅将极小的输入向量与衰减系数写入缓冲区，直接结合检查点计算输出；只有当缓冲区填满时才触发刷写步（flush step），将累积的输入一次性折叠进检查点并写回显存，使主导性的状态写操作由每步一次摊销为每 $L$ 步一次。具体设计包含三个关键部件：

  1. 状态重放与标准解码：在非刷写步，检查点与输入记录联合重建状态以产生输出，定义性公式为：

    $$S_t = \gamma_h S_{\mathrm{ckpt}} + \sum_{j=1}^h \gamma_{j+1:h} v_j k_j^\top$$


    其中 $\gamma$ 为累积衰减系数。常规解码每 $L$ 步需要 $L$ 次全状态读和 $L$ 次写，而 ReplaySSM 仅需 $L$ 次读和 1 次写，使主要的状态显存带宽开销几乎减半。

  2. 纯输出投机验证与指针回滚：验证 $T$ 个草稿 token 时，利用结合律将外积改写为内积形式，计算仅需单次检查点投影而无需物化任何中间状态：

    $$\operatorname{OUTPUTONLY}(S_{\mathrm{ckpt}}, B_{1:h+T}, Q) = S_{\mathrm{ckpt}} Q D + V \big( W \odot (K^\top Q) \big)$$


    草稿被部分拒绝时，系统直接移动缓冲区指针覆盖废弃记录，省去了状态回滚的显存拷贝。

  3. 双精度检查点解耦：系统在显存中分别保存高精度检查点 $S_{\mathrm{ckpt}}^{\mathrm{HP}}$ 与低精度检查点 $S_{\mathrm{ckpt}}^{\mathrm{LP}}$。非刷写步只读取 $S_{\mathrm{ckpt}}^{\mathrm{LP}}$ 计算当前步输出；刷写步则读取 $S_{\mathrm{ckpt}}^{\mathrm{HP}}$ 承接真正的循环累积，并在更新后独立向下投射生成新的低精度与高精度副本。这切断了量化舍入误差重返递归主干的反馈回路。

效果如何 实验涵盖 Nemotron-3-Nano-4B（Mamba-2，BF16 权重）、Nemotron-3-Super-120B-A12B（Mamba-2，NVFP4 权重）、Qwen3.5-4B（GDN，BF16 权重）和 Qwen3.5-122B-A10B（GDN，NVFP4 权重）四款混合模型，分别在单卡 H100 与单卡 B300 上运行，并在 vLLM 框架内与开启 CUDA Graph 的高度调优版常规 vLLM 基线对比；投机解码测试中还对比了基于每 draft 保留快照的 vLLM、融合恢复核的 The Mamba in the Llama，以及依赖重放重跑递归的 Snakes and Ladders 与 STree 等路线。在批大小 256 下，标准解码实现了 1.20–1.47 倍的整模型端到端加速，孤立 SSM 核加速比达 1.46–1.80 倍；在两款 120B 级模型上，投机解码吞吐达到常规 vLLM 投机解码的 1.85 倍与 2.10 倍。在低精度状态评估中，对于 128K 超长文本，采用 NVFP4 读、FP32 循环更新的分离策略使输出 KL 散度较共享 NVFP4 检查点降低了 36 倍；在 GPQA-Diamond 推理基准的长文本分块上，共享 NVFP4 会导致高达 28.2 个百分点的准确率崩塌，而双精度解耦策略成功将差距抹平在 3.4 个百分点以内。其局限在于：在批大小为 1 的小并发场景下，解码耗时主要被模型权重搬运统治而非状态 I/O，此时引入输入缓存不仅没有加速，甚至会因微小开销导致吞吐略低于基线（0.96–0.99 倍）；此外，非刷写步仍需按步读取历史检查点，尚无法完全消除状态读取流量。

主要贡献

状态空间模型（State-Space Models, SSM）如 Mamba-2、Gated DeltaNet (GDN) 以及 Kimi Delta Attention (KDA) 通过将历史 token 压缩进一个固定大小的循环状态（recurrent state）中，避免了传统 Transformer 随着序列长度增长而线性膨胀的键值缓存（KV 缓存）。然而，现有的 SSM 解码系统在每一个生成步中都急切地（eagerly）加载状态、执行更新、计算输出并存回显存（HBM）。这种状态管理策略导致了三大核心瓶颈：

  1. 算术强度极低且显存 I/O 成为系统瓶颈：在批处理服务场景下，每步仅仅执行逐元素衰减和秩-1 更新，却需要完整读写巨大的状态矩阵，使得在混合架构（Hybrid Models）中即使上下文达到 100K tokens，SSM 层的更新依然占据主要的生成延迟。
  2. 循环状态无法低成本回滚（Rollback）：在投机解码（Speculative Decoding）中，当验证失败需要舍弃部分草稿 token 时，传统实现必须为每个草稿 token 存储独立的状态快照，或通过重运行循环更新来恢复状态，这严重消耗显存带宽并限制了并发度。
  3. 量化舍入误差在循环中累积扩散：低精度状态存储下，量化误差在每一步都会重新反馈进下一步的循环状态更新中，随着生成长度增加，导致显著的分布偏移和模型精度下降。

针对上述瓶颈，论文提出了 ReplaySSM，其核心理念为“缓存 SSM 输入，而非状态（Cache SSM inputs, not state）”。论文的核心贡献如下：

  • 解耦输出计算与检查点更新（ReplaySSM 标准解码）：通过展开状态递推关系，将当前状态表示为“历史检查点（Checkpoint）+ 近期输入缓冲（Replay Buffer）”。在大部分非刷新步（non-flush steps）仅追加微小的输入向量并直接计算输出，彻底消除了绝大多数步骤的状态写回操作，将 SSM 主导的状态显存流量减半。
  • 基于指针的零成本投机回滚与仅输出验证（Output-Only Verification）：将草稿 token 的输入直接存入输入缓冲，验证时的回滚操作简化为移动缓冲区指针；同时利用代数结合律推导出的仅输出验证形式，在不物化（materialize）任何中间状态的情况下并行计算所有草稿 token 的输出。
  • 解耦精度的状态管理机制（Separate State Precisions）：在低精度状态服务中，采用高精度检查点维持跨刷新窗口的循环状态递推，采用低精度副本仅用于非刷新步的输出读取，彻底切断了量化舍入误差重新进入循环递推的反馈回路。
  • 系统实现与全面评测：在 vLLM 中为 Mamba-2 和 Gated DeltaNet 实现了 ReplaySSM，并支持 CUDA Graphs。在 Serving 批大小为 256 时，相比高度优化的 vLLM 基线，标准解码实现了 $1.20\times - 1.47\times$ 的端到端加速；在两个 120B 级 MoE 混合模型上，投机解码吞吐量达到 vLLM 投机基线的约 $1.8\times$；同时在高精度/低精度解耦配置下，输出 KL 散度显著降低（降低 $13\times - 36\times$），有效消除了长序列推理中的准确率溃败。

ReplaySSM 架构与核心思想概览
图 1：ReplaySSM 概览。上方：传统解码每步读写完整状态。ReplaySSM 展开循环，将输出计算与检查点更新解耦：缓存近期 SSM 输入，仅周期性将其折叠进检查点，消除了绝大部分步骤的完整状态写入。下方：传统解码面临的三大挑战及 ReplaySSM 的解决方案。(1) 即使在注意力成本随长度增长的混合模型中，SSM 层也会成为长达 100K 上下文时的系统瓶颈；ReplaySSM 将其状态流量减半。(2) 传统投机解码中昂贵的状态回滚变成了轻量级的缓冲区指针移动。(3) 低精度检查点仅负责输出计算，高精度检查点承载循环递推，阻断舍入误差重复进入循环。

背景知识、关键 Observation 与设计原则

1. 状态空间模型（SSMs）的基础递推结构

SSM 通过固定大小的循环状态对历史上下文进行压缩。

  • Mamba-2 递推：其状态更新与输出映射为：

    $$S_t = \alpha_t S_{t-1} + v_t k_t^\top, \qquad o_t = S_t q_t$$


    其中状态 $S_t \in \mathbb{R}^{d_v \times d_k}$，输入向量为 $v_t, o_t \in \mathbb{R}^{d_v}$，$k_t, q_t \in \mathbb{R}^{d_k}$，$\alpha_t \in (0, 1)$ 是由数据决定的标量衰减因子（scalar decay）。累积衰减定义为 $\gamma_{i:j} = \prod_{\ell=i}^j \alpha_\ell$（当 $i > j$ 时 $\gamma_{i:j} = 1$），并将 $\gamma_{1:h}$ 简记为 $\gamma_h$。传统解码在每个步均需从 HBM 加载 $S_{t-1}$，计算得到 $S_t$ 和 $o_t$，然后将 $S_t$ 存回 HBM。

  • Delta-rule 模型（Gated DeltaNet, GDN）：在写入新信息前更新已有记忆，其递推式为：

    $$S_t = \alpha_t S_{t-1} \left( I - \beta_t k_t k_t^\top \right) + \beta_t v_t k_t^\top$$


    其中 $\beta_t \in (0, 1)$ 控制写入强度。引入中间校正项：

    $$u_t = \beta_t \left( v_t - \alpha_t S_{t-1} k_t \right)$$
    递推式可重写为标准外积形式：
    $$S_t = \alpha_t S_{t-1} + u_t k_t^\top$$
    该形式在数学结构上与 Mamba-2 等价，仅需先利用上一时刻状态计算出 $u_t$。

2. 传统解码策略的三大核心挑战（Observation）

  • 观察 1：极低的算术强度与显存带宽受限（Low Arithmetic Intensity）
    在传统 Mamba-2 解码中，单头每步加载和存储状态需传输 $2 b d_v d_k$ 字节（$b$ 为每元素字节数）。而计算仅包含逐元素缩放、外积求和与矩阵-向量乘法，共执行约 $5 d_v d_k - d_v$ FLOPs。其算术强度仅为：

    $$I_{\mathrm{conv}} \approx \frac{5}{2b} \text{ FLOP/byte}$$


    对于 FP32 状态（$b=4$），算术强度仅为 0.63 FLOP/byte；对于 FP16（$b=2$），仅为 1.25 FLOP/byte。由于每个序列都维护独立的状态，增加 batch size 无法提升算术强度。现代加速器（如 NVIDIA H100）保持算力核心饱和需要约 20（FP32 CUDA 核心）至 300（BF16 Tensor 核心）FLOP/byte。因此传统 SSM 解码绝对受限于 HBM 带宽。
    在生产级混合架构（如 Nemotron-3-Super-120B-A12B，包含 40 个 Mamba-2 层和 8 个 Attention 层）中，每个 SSM 层的 FP32 状态为 4 MiB/seq，等价于一个注意力层 8K 长度的 FP8 KV 缓存。由于 SSM 每步进行读写（8 MiB），而注意力仅读取已有缓存，在上下文达到约 80K~100K 之前，SSM 层的显存传输量始终压倒注意力层。

  • 观察 2：循环状态难以逆转，制约投机解码（No Cheap Rollback）
    投机解码需要一次前向传播验证 $T$ 个草稿 token，并根据验证结果接受前 $a$ 个 token（丢弃末尾 $T-a$ 个）。Attention 仅需回退 KV 指针，而 SSM 由于每步将输入破坏性地折叠进固定大小状态中，无天然指针可用。为支持回滚，vLLM 等系统对每个草稿 token 强制写出一份状态快照（每步需 $1\text{R} + T\text{W}$ 的全状态 I/O），极大消耗带宽并成倍削减并发请求量；先前的恢复方案（如重跑递推）计算与显存开销巨大，在 serving 批大小下甚至慢于快照方案。

  • 观察 3：低精度量化误差沿循环链式传递（Quantization Error Propagation）
    若对存储状态进行低精度量化，每步量化后的近似状态 $\hat{S}_t$ 既用于输出，又作为下一步的输入。舍入误差在每一步被重新注入循环系统。若采用最近舍入（RTN），固定符号的偏置误差按 $\frac{b}{1-\alpha}$ 累加；若采用随机舍入（SR），方差按 $\frac{\sigma^2}{1-\alpha^2}$ 累积。对于慢衰减头（$\alpha \to 1$），误差在长上下文下发散。

3. 核心设计原则：展开递推与输入回放

循环公式自身证明了“每步写回状态”是非必要的。设 $S_{\mathrm{ckpt}}$ 为包含 $h$ 个输入的窗口之前的状态，则经过 $h$ 步后，展开公式为：

$$S_t = \gamma_h S_{\mathrm{ckpt}} + \sum_{j=1}^h \gamma_{j+1:h} v_j k_j^\top$$


当前状态可严格由历史检查点 $S_{\mathrm{ckpt}}$ 与近期输入集合 $\{(\alpha_j, v_j, k_j)\}_{j=1}^h$ 线性重构。因此，系统无须每步将状态持久化，只需将小尺寸输入写入环形缓冲，把全状态写入平摊至窗口刷新步。

方法细节

1. 标准解码中的 ReplaySSM

状态重构算子与算法流程

在标准单步自回归解码中，核心操作定义为利用检查点与输入缓冲重构当前状态的 $\mathrm{REPLAY}$ 算子：

$$\mathrm{REPLAY}\left(S_{\mathrm{ckpt}}, \mathcal{B}_{1:h}\right) = \gamma_h S_{\mathrm{ckpt}} + \sum_{j=1}^h \gamma_{j+1:h} v_j k_j^\top$$


对于长度为 $L$ 的回放窗口，前 $L-1$ 步为非刷新步（non-flush steps），最后一步为刷新步（flush step）。其标准解码逻辑严格遵循 Algorithm 1。

Algorithm 1: ReplaySSM 标准解码 (Mamba-2)
输入: 当前 token 的 x = (α, v, k) 与查询向量 q
1: B[h + 1] ← x;  h ← h + 1
2: S ← REPLAY(S_ckpt, B[1:h])
3: o ← S * q
4: if h == L then             ▷ 触发检查点刷新
5:     S_ckpt ← S; 将 S_ckpt 写入 HBM
6:     B ← ∅;  h ← 0
7: end if
8: return o

显存流量削减机理

在大小为 $L$ 的窗口中：

  • 传统解码每步进行一次状态读取和一次状态写入，总计执行 $L$ 次读和 $L$ 次写，移动显存流量为 $8 L d_v d_k$ 字节（FP32 下每元素 4 字节）。
  • ReplaySSM 每步仍需读取检查点 $S_{\mathrm{ckpt}}$ 用于输出计算，总计 $L$ 次状态读取；但在整个窗口内仅在第 $L$ 步写回一次更新后的检查点（1 次写）。
  • 非刷新步仅向环形缓冲追加输入记录 $x = (\alpha, v, k)$。由于单个记录仅包含标量衰减和向量，其体积 $r = 2 d_v + 2 d_k + 4$ 字节远小于二维状态矩阵大小 $4 d_v d_k$。
  • 在 $L=16$ 时，检查点读写项的比例为 $\frac{L+1}{2L} \approx \frac{17}{32} \approx 0.53$。加入缓冲区读写后，在 Nemotron-3-Super 维度下，总显存传输量降低至传统方法的 $0.58\times$（减少约 $1.7\times$）。

标准解码加速比对比
图 2：在 Serving 批大小 $B=256$、窗口 $L=16$、FP32 状态下，ReplaySSM 相比调优后的 vLLM 基线在四种混合模型上的标准解码加速比。深色为纯 SSM 内核级加速（包含刷新开销，达 $1.46\times - 1.80\times$），浅色为整网端到端解码步加速（达 $1.20\times - 1.47\times$）。

2. 投机解码中的 ReplaySSM

指针级回滚（Cheap Rollback）

投机解码每步由草稿模型生成 $T$ 个候选 tokens，目标模型需在一趟前向传播中验证它们。由于被接受的前缀长度 $a \in [1, T]$ 必须在验证完成后才能获知，传统架构无法确定应保留哪个中间状态。ReplaySSM 将回滚机制简化为环形缓冲指针的操作：

  1. 验证前：直接将这 $T$ 个草稿 token 的输入记录批量追加在已提交的 $h$ 条记录之后（位置 $h+1$ 至 $h+T$）。
  2. 状态验证：利用全部 $h+T$ 条记录计算 $T$ 个输出 token。
  3. 验证后：根据采样返回的接受长度 $a$，只需执行标量赋值 $h \leftarrow h + a$。
  4. 被拒绝的 $T-a$ 个记录自然停留在指针 $h$ 之后，在下一次生成步追加输入时被直接就地覆盖。

仅输出验证（Output-Only Verification）

传统投机验证需计算并保留每个草稿位置的状态矩阵 $S_t$。ReplaySSM 发现非刷新步不需要中间状态，只需输出 $o_t$。对于单个无初态更新：

$$o = \left(v k^\top\right) q = v \left(k^\top q\right)$$


通过交换运算结合律，先进行向量内积 $k^\top q$ 得到标量，再缩放向量 $v$，直接避免了生成 $d_v \times d_k$ 大小的中间外积矩阵。
推广至包含 $h$ 个历史提交记录和 $T$ 个草稿 token 的验证序列，令矩阵 $V, K$ 的列向量分别包含这 $h+T$ 个缓冲记录的键和值，$Q = [q_1, \dots, q_T]$，则所有草稿位置的输出向量矩阵 $O \in \mathbb{R}^{d_v \times T}$ 可由下式直接并行求得：

$$\mathrm{OUTPUTONLY}\left(S_{\mathrm{ckpt}}, \mathcal{B}_{1:h+T}, Q\right) = S_{\mathrm{ckpt}} Q D + V \Big( W \odot \left(K^\top Q\right) \Big)$$

  • 其中 $D \in \mathbb{R}^{T \times T}$ 是对角矩阵，对角元为 $D_{ss} = \gamma_{h+s}$；
  • $\odot$ 表示 Hadamard 逐元素积；
  • $W \in \mathbb{R}^{(h+T) \times T}$ 为包含因果掩码的累积衰减矩阵：当 $j \le h + s$ 时 $W_{js} = \gamma_{j+1:h+s}$，否则为 0。
  • 该形式具有显著计算优势：彻底消除了状态的物化；对于共享键和查询头（GQA/MQA 特性）的一组值头，小 GEMM 操作 $K^\top Q$ 只需计算一次并在头之间共享。
Algorithm 2: ReplaySSM 投机解码 (Mamba-2)
输入: T 个草稿 token 的输入 X = {x_s}_{s=1}^T 与查询矩阵 Q
1: B[h + 1 : h + T] ← X                  ▷ 写入草稿记录，h 暂不改变
2: 从 HBM 加载 S_ckpt 到片上 S
3: O ← OUTPUTONLY(S, B[1:h+T], Q)         ▷ 并行计算 T 个草稿的输出
4: if h + T > L then                     ▷ 刷新已提交的记录
5:     S ← REPLAY(S, B[1:h])
6:     S_ckpt ← S; 将 S_ckpt 存回 HBM
7:     丢弃 B[1:h];  h ← 0                ▷ 仅保留 B[h+1:h+T]
8: end if
9: return O
-- 验证判定产生接受长度 a 之后 --
10: h ← h + a                            ▷ 指针前进，正式提交被接受的草稿

投机解码服务吞吐量与内核延迟
图 3：在单张 B300 上，$L=16, T=4$, FP32 状态下的投机服务吞吐量对比。针对 Nemotron-3-Super-120B (上方) 和 Qwen3.5-122B (下方)，传统 vLLM 投机在并发较大时性能停滞甚至下降，而 ReplaySSM 展现出优异扩展性，在 $B=512$ 时达到基线投机吞吐量的 $1.85\times - 2.10\times$。

投机验证内核单步延迟
图 4：在 Nemotron-3-Super-120B 上不同草稿长度 $T$ 的投机内核耗时对比（$B=256, L=16$）。传统 vLLM 随 $T$ 线性写回快照，延迟急剧上升（$T=8$ 时达标准步的 $4.66\times$）；ReplaySSM 利用 Output-Only 形式与稀疏写回，验证 8 个 token 的开销仅为标准解码步的 $0.89\times - 1.21\times$。

3. 解耦精度的低精度状态管理（Separate State Precisions）

误差隔离原理

为了打破量化舍入误差每步反馈回递推链的恶性循环，ReplaySSM 利用输出计算与状态更新的物理分离，在 HBM 中维护两套不同精度的检查点：

  • 高精度检查点（High-Precision Checkpoint, $S_{\mathrm{ckpt}}^{\mathrm{HP}}$）：如 FP32 或 FP16+随机舍入（SR），专门用于在刷新步承载跨窗口的长程循环递推。
  • 低精度检查点（Low-Precision Checkpoint, $S_{\mathrm{ckpt}}^{\mathrm{LP}}$）：如 FP8 或 NVFP4，专门用于在非刷新步为 $\mathrm{OUTPUTONLY}$ 提供低带宽消耗的矩阵乘法读取。
Algorithm 3: ReplaySSM 双精度投机解码 (Mamba-2)
输入: T 个 token 的输入 X 与查询矩阵 Q
1: B[h + 1 : h + T] ← X
2: π ← (若 h + T > L 则为 HP，否则为 LP)   ▷ 动态选择读取的检查点精度
3: 从 HBM 加载 S_ckpt^π 到片上 S
4: O ← OUTPUTONLY(S, B[1:h+T], Q)
5: if h + T > L then                     ▷ 刷新分支使用高精度路径
6:     S ← REPLAY(S, B[1:h])
7:     S_ckpt^HP ← Q_HP(S)
8:     S_ckpt^LP ← Q_LP(S)
9:     将 S_ckpt^HP 与 S_ckpt^LP 分别存回 HBM
10:    丢弃 B[1:h];  h ← 0
11: end if
12: return O
-- 获得接受长度 a 之后 --
13: h ← h + a

在刷新时，算法从 $S_{\mathrm{ckpt}}^{\mathrm{HP}}$ 出发，用高精度累加重构状态，再独立执行两次类型转换：$\mathcal{Q}_{\mathrm{HP}}$（采用随机舍入保持无偏）和 $\mathcal{Q}_{\mathrm{LP}}$（采用最近舍入降低当前误差）。低精度副本 $S_{\mathrm{ckpt}}^{\mathrm{LP}}$ 从不作为递推初态输入回放，其截断误差完全被局限在当前的非刷新步输出中，彻底避免了随步数发散。

Gated DeltaNet (GDN) 的两精度自洽修正

GDN 的非刷新步输入缓冲记录包含校正项 $u_t$。若直接存储 $u_t$，由于非刷新步从 $S_{\mathrm{ckpt}}^{\mathrm{LP}}$ 计算得到 $u_t$，该校正项已被低精度污染。若刷新时将带有低精度误差的 $u_t$ 折叠进高精度检查点，误差仍会渗入递推链。
解决方案：双精度 GDN 记录扩展为包含原始输入 $\{(\alpha_j, \beta_j, u_j, v_j, k_j)\}$。非刷新步直接使用 $u_j$；在刷新时刻，算法抛弃缓冲中缓存的 $u_j$，而是基于高精度初态 $S_{\mathrm{ckpt}}^{\mathrm{HP}}$ 和原始输入 $(v_j, k_j)$，沿纯净的高精度路径重新推导精准校正项：

$$u_j^{\mathrm{HP}} = \beta_j \left(v_j - \alpha_j S_{j-1}^{\mathrm{HP}} k_j\right)$$


从而确保高精度状态的数学自洽性与误差纯洁度。

实验环境

1. 评测模型与参数规格

实验选取两大家族共 4 个主流生产级/开源混合架构大模型，全部采用单卡无张量并行（Tensor Parallelism）模式运行：

  • Nemotron-3-Nano-4B：BF16 权重；SSM 架构为 Mamba-2；21 个 SSM 层 / 4 个 Attention 层；SSM 头配置为 8 个 Key 头，96 个 Value 头，$d_k=128, d_v=80$。
  • Nemotron-3-Super-120B-A12B：NVFP4 权重；SSM 架构为 Mamba-2；40 个 SSM 层 / 8 个 Attention 层；8 个 Key 头，128 个 Value 头，$d_k=128, d_v=64$。
  • Qwen3.5-4B：BF16 权重；SSM 架构为 Gated DeltaNet (GDN)；24 个 SSM 层 / 8 个 Attention 层；16 个 Key 头，32 个 Value 头，$d_k=128, d_v=128$。
  • Qwen3.5-122B-A10B：NVFP4 权重；SSM 架构为 Gated DeltaNet (GDN)；36 个 SSM 层 / 12 个 Attention 层；16 个 Key 头，64 个 Value 头，$d_k=128, d_v=128$。
  • 投机解码机制：两个 120B 级模型均自带 Multi-Token-Prediction (MTP) 投机头作为 Drafter。

2. 硬件与软件配置

  • 硬件平台：4B 参数量模型部署于单张 NVIDIA H100 (80GB HBM3) GPU；120B 级 MoE 模型部署于单张 NVIDIA Blackwell B300 (1100W 功耗受限) GPU。
  • 软件栈与系统实现：基于 vLLM 统一构建；使用 CUDA Graphs 捕获解码流；算子核心基于 Triton 和 CuTe DSL 编写；底层显存分析借助 NVIDIA Nsight Compute。

3. 基线与量化设置

  • 基线对比：对比调优后的 vLLM 官方标准与投机解码算子（参数均经过全网格 Launch 搜索，超越官方默认配置）；投机解码对比了 The Mamba in the Llama、Snakes and Ladders、STree 等最前沿方案。
  • 低精度量化格式：FP8 采用 E4M3（每 128 元素块分配一个 FP32 缩放因子）；NVFP4 采用 E2M1（每 16 元素块分配一个 E4M3 缩放因子，整个状态矩阵分配一个 FP32 缩放因子）。

4. 评估数据集与协议

  • 长文本分布偏离度（KL 散度）：采用 PG-19 测试集（长篇书籍文本），在 Teacher Forcing 条件下评估多步解码输出分布相对 FP32 基线的 $D_{\mathrm{KL}}(p_{\mathrm{FP32}} \parallel p_{\mathrm{config}})$，序列长度延伸至 128K tokens。
  • 下游推理与逻辑基准：MMLU-Pro（多任务理解与推理，12,032 题）、GPQA-Diamond（极难研究生级常识推理，198 题重复 3 次评测）、Math-L5（MATH 数据集 Level 5 最难题集，1,324 题）、RULER-128K（长上下文检索与聚合任务）。

实验结果

1. 标准解码吞吐与内核加速

  • 内核级与端到端性能：在常规服务并发 $B=256$、窗口 $L=16$、FP32 状态下，SSM 内核计算（含刷新成本）实现 $1.46\times - 1.80\times$ 的加速；全模型单步端到端推理实现 $1.20\times - 1.47\times$ 加速（见图 2 与表 6、表 7）。
  • 批大小扩展规律（Batch-Size Sweep）：在 $B=1$ 时内核受限于延迟而非显存带宽，基线带宽利用率仅为 $9.3\% - 21.6\%$，此时 ReplaySSM 与基线持平或稍逊；当 $B \ge 128$ 时基线迅速打满 HBM 带宽（达到峰值的 $84\% - 86\%$），此时 ReplaySSM 显存流量减半的优势完全兑现，加速比达到平台期（见图 8）。

2. 投机解码服务吞吐与内核延迟

  • 服务吞吐量表现：在 120B 级 NVFP4 模型（$L=16, T=4$）的实际 GSM8K 请求服务中，传统 vLLM 投机解码由于 $O(T)$ 的状态写回开销及快照导致的显存膨胀，在 $B \ge 128$ 时吞吐量饱和，在 $B=512$ 时其性能甚至跌落至普通标准解码以下。而 ReplaySSM 吞吐量持续线性爬升，在 $B=512$ 时达到 vLLM 投机基线的 $1.85\times$（Nemotron）与 $2.10\times$（Qwen）（见图 3）。
  • 单步验证内核开销：传统验证随 $T$ 呈急剧上升趋势，$T=8$ 时内核延迟达到常规单步解码的 $4.66\times$（Mamba-2）与 $6.40\times$（GDN）；ReplaySSM 依靠仅输出推导和缓冲区回放，在验证 8 个 token 时的整体内核延迟仅为普通解码单步的 $0.89\times - 1.21\times$（图 4、图 10）。

3. 低精度状态的 KL 散度与长序列稳定性

  • 累积误差分析：在 PG-19 数据集上随着生成长度 $N$ 增加至 128K，采用共享单一检查点的 FP8+SR 和 NVFP4+SR 误差呈数个数量级发散；而采用“NVFP4 读取 / FP32 递推”的双精度方案，KL 散度在 128K 时较共享 NVFP4 降低了 $36\times$（从 0.136 降至 0.0038），且全程保持水平平稳（图 5）。
  • FP16+SR 替代可行性：“NVFP4 读取 / FP16+SR 递推”与 FP32 递推曲线重合度极高，表明半精度随机舍入足以满足长序列循环递推的数值精度要求。

PG-19 长文本输出 KL 散度随生成长度变化
图 5：在 Qwen3.5-4B 上以 FP32 传统解码为基准，Teacher Forcing 下不同量化配置的输出分布 KL 散度。共享检查点随上下文爆炸式发散；解耦精度方案（蓝线）成功切断误差传播链，在 128K 长度时实现 36 倍的 KL 散度压制。

4. 下游评测基准准确率（Accuracy）

  • 长思维链任务崩溃现象：如表 2 及图 6 所示，共享 NVFP4+SR 检查点在 GPQA-Diamond 上平均崩塌 13.05 个百分点（Nemotron-4B 上大跌 31.3 点，Nemotron-120B 上跌 12.5 点），在 MMLU-Pro 和 Math-L5 上分别损失 5.88 和 10.92 个百分点。错误的主要诱因是长生成过程中的累积误差导致模型无法正常收敛输出而耗尽 Token Budget（图 14）。
  • 双精度机制的精度挽救：采用“NVFP4 读取 / FP16+SR 递推”后，GPQA、MMLU-Pro 和 Math-L5 上的准确率损失被分别大幅收敛至 2.36、0.15 和 0.10 个百分点；在更大模型上几乎完全抹平了与 FP32 的性能代差。
  • 超长输入评测（RULER-128K）：在输入长达 130K、输出仅数十 token 的 RULER 任务中，所有低精度配置（包括共享 NVFP4）与 FP32 的分差均在 0.3 个百分点以内（表 13）。这实证证明了：低精度状态的精度衰减是由长生成累积造成的，而非源自长输入前缀。

GPQA-Diamond 准确率按输出长度四分位数切分
图 6：在 Nemotron-3-Super-120B 上按参考生成长度分桶的 GPQA-Diamond 准确率。短输出桶中共享 NVFP4 与基线差异微弱；而在超长输出桶（中位数 61K tokens）中，共享 NVFP4 暴跌 28.2 个百分点，而双精度解耦方案（NVFP4 读取 / FP8+SR 递推，蓝线）与 FP32 紧密贴合。

结论

传统的状态空间模型自回归解码每步强制执行状态的读取与写回，将输出计算与循环状态的更新紧紧绑定，从而引发了显存带宽受限、投机回滚代价高昂以及量化误差无限累积三大死结。ReplaySSM 通过代数展开，将状态表示解构为“历史检查点 + 输入回放缓冲”，实现了两者的彻底解耦：

  1. 显存优化：大部分生成步不写回状态，核心状态显存传输流量近似减半，在主流批处理规模下获得 $1.20\times - 1.47\times$ 的端到端加速。
  2. 投机优化：将昂贵的状态快照与递推恢复转变为轻量级的缓冲区指针移动，结合仅输出验证消除中间状态物化，投机吞吐量提升至 vLLM 的 $1.8\times - 2.1\times$。
  3. 量化优化：高精度检查点维系跨窗口无偏递推，低精度检查点支撑高吞吐非刷新步读取，消除了量化误差在时序上的滚雪球效应，使极低精度状态在长序列推理中的应用成为可能。

未来工作展望

在 ReplaySSM 中，每步依然必须读取一次检查点。由于检查点是通过累积衰减因子 $\gamma_h$ 乘入输出的，对于具有极快时间衰减特性的注意力/SSM 头，$\gamma_h$ 在进入窗口数步后即可衰减至数值精度不可分辨的微小范围。未来工作将探索基于衰减强度的检查点跳读机制（Skipping Checkpoint Reads by Decay）：快衰减头可直接略过非刷新步的检查点加载，并在刷新步变为纯写入操作；或仅读取极粗糙的超低精度检查点副本，进一步释放显存带宽潜能（参见附录 F）。

附录内容详细解析

附录 A：投机解码对比前序方案的深入分析

在 Nemotron-3-Super-120B 形状下（$B=256$, FP32 状态，单状态 $4.2\text{ MB}$，全 Batch 单层状态达 $1.07\text{ GB}$），附录深入评测了 4 种现有投机状态处理机制（见图 7 与表 3）：

投机解码单步内核延迟与已有方案对比
图 7：Nemotron-3-Super-120B-A12B (Mamba-2) 上单步投机内核耗时对比（归一化到标准解码 $T=1$）。展示了全拒绝（$a=1$）到全接受（$a=T$）的延迟带。STree 与 The Mamba in the Llama 受限于串行递推与寄存器常驻限制，在大 Batch 下均无法击败 vLLM Snapshot 基线，而 ReplaySSM 实现了数倍的延迟压缩。

  • vLLM Snapshots：内核单次启动加载 1 次状态，在计算每个草稿 token 后强制写出该位置的完整状态快照（$1\text{R} + T\text{W}$）。当 $T=8$ 时单层单步移动 $9.7\text{ GB}$ 显存，导致严重的显存带宽挤占。
  • The Mamba in the Llama (Wang et al., 2024a)：单次融合内核启动，重新从缓存输入把前一步已接受的 token 回放到检查点上，再运行当前 $T$ 个 token（$1\text{R} + 2\text{W}$）。虽然减少了写回次数，但在 GPU 上形成了长达 $2T-1$ 步的串行递归链，指令级并行度极低。

  • Snakes and Ladders (Wu et al., 2024)：

    • Activation Replay：启动 3 次内核，执行重算、拷贝与验证，单步产生高达 $3\text{R} + 3\text{W}$（$6.4\text{ GB}$）的状态数据流动。
    • Joint Attainment：不依赖激活缓存，单请求维护两行执行填充回溯与推导，需要 $3\text{R} + 4\text{W}$（$7.5\text{ GB}$）流量，并受限于捕获的 CUDA 图尺寸。

  • STree (Wu et al., 2026)：执行掩码重放后进行 Chunkscan 验证，算法本身需 $2\text{R} + 1\text{W}$（$3.2\text{ GB}$），其实际开源路径包含冗余拷贝，达到 $5\text{R} + 4\text{W}$（$9.7\text{ GB}$）。

  • ReplaySSM 的断代优势：非刷新步仅为 $1\text{R}$（$1.1\text{ GB}$），仅刷新步发生 $1\text{R} + 1\text{W}$（$2.1\text{ GB}$）。在 $a=T$ 时，ReplaySSM 比 STree 快 $3.3\times - 3.9\times$，比 The Mamba in the Llama 快 $2.5\times - 3.6\times$；在 $a=1$ 时优势进一步扩大至 $3.5\times - 5.3\times$。实验揭示了一个重要底层洞察：更少的状态字节并不直接带来更低的延迟，先前基于串行递归展开的重算核因寄存器常驻和占用率受限，大部分 DRAM 带宽处于空转饥饿状态，而 ReplaySSM 的仅输出形式以矩阵乘法组织，能高效填满硬件流水线。

附录 B：GDN 完整数学推导与专用算法

1. 外积与展开形式推导

从 GDN 原始递推出发：

$$S_t = \alpha_t S_{t-1} \left( I - \beta_t k_t k_t^\top \right) + \beta_t v_t k_t^\top = \alpha_t S_{t-1} - \alpha_t \beta_t \left( S_{t-1} k_t \right) k_t^\top + \beta_t v_t k_t^\top$$


利用结合律将中间项重整为 $\beta_t \left( v_t - \alpha_t S_{t-1} k_t \right) k_t^\top$。令 $u_t = \beta_t \left( v_t - \alpha_t S_{t-1} k_t \right) \in \mathbb{R}^{d_v}$，成功转化为 Mamba-2 相同的标准外积更新结构：

$$S_t = \alpha_t S_{t-1} + u_t k_t^\top$$
在此结构下，整个回放窗口 $t=1, \dots, h$ 展开后与检查点 $S_0 = S_{\mathrm{ckpt}}$ 的关系为：
$$S_t = \gamma_t S_{\mathrm{ckpt}} + \sum_{j=1}^t \gamma_{j+1:t} u_j k_j^\top$$

2. GDN 标准解码流程

Algorithm 4: ReplaySSM 标准解码 (Gated DeltaNet)
输入: 当前 token 的 (α, β, v, k) 与查询向量 q
1: S ← REPLAY(S_ckpt, B[1:h])              ▷ 先重构当前步之前的状态 S_{t-1}
2: u ← β * (v - α * S * k)                 ▷ 在片上高精度计算校正项
3: B[h + 1] ← (α, u, k);  h ← h + 1        ▷ 将校正项代替 v 存入缓冲区
4: S ← α * S + u * k^T                     ▷ 更新得到当前步状态 S_t
5: o ← S * q                               ▷ 产生输出
6: if h == L then
7:     S_ckpt ← S; 存回 HBM
8:     B ← ∅;  h ← 0
9: end if
10: return o

3. GDN 投机验证中的上三角求解（UT Transform）

在投机验证中，草稿位置的校正项 $u_t$ 依赖于包含前面草稿影响的状态。设草稿处于窗口位置 $t = h + s$：

$$\alpha_t S_{t-1} k_t = \gamma_t S_{\mathrm{ckpt}} k_t + \sum_{j=1}^{t-1} \gamma_{j+1:t} \left( k_j^\top k_t \right) u_j$$


其中 $j \le h$ 的历史项使用已缓存的校正项 $U_c$；而 $h < j < t$ 的草稿项未知。把全部 $T$ 个未知草稿校正项组合为矩阵 $U_d \in \mathbb{R}^{d_v \times T}$，可建立紧致线性方程组：

$$U_d \left( I + A B_d \right) = R B_d$$

  • 其中 $B_d = \mathrm{diag}(\beta_{h+1}, \dots, \beta_{h+T})$；
  • $A \in \mathbb{R}^{T \times T}$ 为严格上三角矩阵，当 $j < s$ 时 $A_{js} = \gamma_{h+j+1:h+s} k_{h+j}^\top k_{h+s}$，其余位置为 0；
  • 右端项 $R = V_d - S_{\mathrm{ckpt}} K_d D - U_c \Big( W_c \odot \left(K_c^\top K_d\right) \Big)$。
    由于 $I + A B_d$ 是单位上三角矩阵（Unit Upper Triangular），通过单次轻量级的 $T \times T$ 矩阵三角求解（Triangular Solve）即可完全并行解析出所有草稿校正项 $U_d$。求解后，输出计算等价转为：
    $$O = S_{\mathrm{ckpt}} Q D + U \Big( W \odot \left(K^\top Q\right) \Big), \quad U = [U_c, U_d]$$
    $S_{\mathrm{ckpt}} K_d$ 与 $S_{\mathrm{ckpt}} Q$ 在对检查点的单次流式加载过程中合并完成。

附录 C：计算开销与数值误差详尽推导

1. 算术强度精确建模（C.1）

传统 Mamba-2 单步单头计算：

  • 显存流量：加载 $S_{t-1}$（$b d_v d_k$）+ 写出 $S_t$（$b d_v d_k$）= $2 b d_v d_k$ 字节。
  • 浮点操作：衰减 $\alpha_t S_{t-1}$（$d_v d_k$）+ 外积 $v_t k_t^\top$（$d_v d_k$）+ 状态求和（$d_v d_k$）+ 输出投影 $S_t q_t$（$2 d_v d_k - d_v$）= $5 d_v d_k - d_v$ FLOPs。
  • 算术强度极限：$I_{\mathrm{conv}} \approx \frac{5}{2b}$ FLOP/byte。

2. 混合架构中 SSM 与 Attention 的流量临界点（C.2）

以 Nemotron-3-Super-120B 为例：

  • SSM 层（40 层）：每层 128 个值头（$d_v=64, d_k=128$），FP32 状态。单序列单层状态大小 $128 \times 64 \times 128 \times 4\text{ B} = 4\text{ MiB}$。每步读+写消耗 $8\text{ MiB}$。40 层总计消耗 $320\text{ MiB/step}$。
  • 注意力层（8 层）：每层 2 个 KV 头（维度 128），FP8 缓存。单 token 的 KV 占用为 $2 \times 2 \times 128 \times 1\text{ B} = 0.5\text{ KiB}$。在上下文长度 $N$ 时，8 个注意力层单步读取显存为 $8 \times 0.5 N\text{ KiB} = 4 N\text{ KiB/step}$。
  • 流量临界点方程：$4 N = 320 \times 1024 \implies N = 81,920$ tokens。
    在长达 8.2 万 tokens 之前，SSM 状态搬运的数据量始终压倒全部注意力 KV 缓存的总读取量。

3. 显存流量封闭解（C.3）

在长度为 $L$ 的回放窗口中，传统解码与 ReplaySSM 总流量方程对比：

$$M_{\mathrm{conv}} = 8 L d_v d_k$$

$$M_{\mathrm{replay}} = 4 (L + 1) d_v d_k + \frac{(L - 1)(L + 2)}{2} r$$
其中 $r = 2 d_v + 2 d_k + 4$ 为单条缓冲记录字节数。

4. 仅输出计算的 FLOPs 对比（C.4）

对于包含 $h$ 条缓冲记录与当前输入的单步计算：

  • 物化状态路线：重构状态消耗 $2(h+2)d_v d_k$，计算输出消耗 $2d_v d_k$，合计约 $2(h+3)d_v d_k$ FLOPs。

  • 仅输出路线：

    • 检查点投影 $S_{\mathrm{ckpt}} q_t$：$2 d_v d_k$
    • 内积求和 $k_j^\top q_t$（每个 Key 组计算 1 次）：$2 (h+1) d_k$
    • 加权求和：$2 (h+1) d_v$
    • 合计仅需 $2 d_v d_k + 2 (h+1)(d_v + d_k)$ FLOPs。

  • 节省比例：计算量从 $O(h \cdot d_v d_k)$ 级骤降，在常见维度下物化项被压缩了 $\frac{d_v d_k}{d_v + d_k} \approx 43\times$。

5. 最佳回放窗口的理论求解（C.5）

设平均单步接受草稿数为 $\bar{a}$，单周期总步数为 $L/\bar{a}$。单 token 均摊显存传输开销函数定义为：

$$f(L) = \frac{M_S}{\bar{a}} + \frac{M_S}{L} + \frac{r L}{2\bar{a}} - \frac{r}{2} + \frac{r T}{\bar{a}} - \frac{r T}{L}$$


求导并令 $f'(L) = -\frac{M_S - rT}{L^2} + \frac{r}{2\bar{a}} = 0$，解得最佳理论窗口：

$$L^* = \sqrt{2\bar{a}\left(\frac{M_S}{r} - T\right)}$$
在 Nemotron-3-Super 维度下（$M_S / r \approx 84$），标准解码（$\bar{a}=T=1$）解得 $L^* \approx 13$；实测最优值为 16（最贴近的 2 的幂次）。投机解码（$T=4, \bar{a}=3.5$）解得 $L^* \approx 24$。

6. 量化舍入误差时序传递理论（C.6）

设第 $t$ 步的量化舍入误差为 $e_t$。累积偏差状态方程为：

$$D_t = \alpha_t D_{t-1} + e_t = \sum_{j \le t} \gamma_{j+1:t} e_j$$


若 $e_t$ 具有偏置 $b$ 与方差 $\sigma^2$：

  • 均值极限定理：$\mathbb{E}[D_t] \to \frac{b}{1-\alpha}$。
  • 方差极限定理：$\mathrm{Var}(D_t) \to \frac{\sigma^2}{1-\alpha^2}$。
    在传统逐步量化中，若 $\alpha \to 1$，偏置与方差无界放大；在共享低精度检查点 ReplaySSM 中，由于刷新间隔为 $L$，衰减步长变为 $\alpha^L$，极限收缩为 $\frac{b}{1-\alpha^L}$ 与 $\frac{\sigma^2}{1-\alpha^{2L}}$；而在解耦双精度下，输出误差退化为单步无累加的单次舍入，误差彻底与序列长度 $N$ 解绑。

附录 D：Serving 与推理系统工程实现

1. 投机解码刷新准则（D.1）

维护容量为 $L+T$ 的环形缓冲（存储 $L$ 个已提交记录与 $T$ 个临时草稿）。为了防止极端情况下（当前草稿全接受且下一步再次进入 $T$ 个新草稿）发生缓冲区指针回绕覆盖未提交数据，刷新判定严格设置为 $h + T > L$。一旦满足该条件，在草稿验证返回前立即将现存 $h$ 个已提交记录无条件折叠进检查点。由于已提交记录在语义上具备永久确定性，该前置刷新策略完全解除了对采样结果的回调依赖。

2. CUDA Graphs 动态控制流设计（D.2）

在连续批处理（Continuous Batching）中，各序列进入时间与接受长度 $a$ 各异，导致 Batch 内不同行在同一解码时刻的刷新状态产生分歧。为了适配固化执行路径的 CUDA Graphs：

  • 行级动态分支：将刷新决定作为张量标志在 GPU 显存内动态维护。单步 Kernel 内部根据标志动态跳转，非刷新行在遇到刷新逻辑时快速早退，避免重捕获 Graph。
  • 设备端状态提交 Kernel：采样产生每个序列的接受长度 $a$ 后，不将 $a$ 拷贝回 CPU Host 端，而是直接由常驻在 GPU 上的轻量 Commit 内核推进环形指针 $h \leftarrow h + a$ 并计算下一步的刷新掩码，彻底消除了 Host-Device 同步屏障。

附录 E：补充评测数据与参数扫描

1. 批大小全面扫描细化数据（E.2.1）

表 6 和表 7 展示了从 $B=1$ 至 $B=512$ 的详尽加速对比。在 B300 硬件功耗上限为 1100W 时，大 Batch 下传统基线在整网中发热严重发生降频节流，而 ReplaySSM 显著降低了 DRAM 搬运能耗，端到端加速进一步飙升至 $1.72\times$。

Nemotron-3-Nano-4B 上的标准解码 Batch-Size 扫描
图 8：Nemotron-3-Nano-4B 在单张 H100 上的批大小扫描。左：传统内核的 DRAM 峰值带宽占用率（小 Batch 时低，在 $B \ge 128$ 逼近饱和）。中：独立 SSM 内核加速比。右：整网单步端到端解码加速比。

2. 回放窗口网格搜索（E.2.2）

图 9 在所有 4 个模型上验证了 $L \in \{8, 16, 32\}$ 的表现。较小窗口（$L=8$）由于刷新频繁，检查点写出占比过高；较大窗口（$L=32$）单步加载的输入记录过多，耗尽片上共享内存与寄存器。四款模型均在 $L=16$ 处取得最优均摊延迟。

回放窗口尺寸扫描
图 9：在 $B=256$, FP32 状态下，不同模型均摊内核延迟随回放窗口大小 $L$ 的变化对比。所有模型均在 $L=16$ 处取得最低均摊延迟。

3. 投机解码全网格扫描（E.3.1）

表 8、9、10 详细记录了不同 $T$（2, 4, 6, 8）与 $B$ 组合下的吞吐数据。在 Nemotron-3-Super 上，$T=8, B=512$ 时 ReplaySSM 较基线投机获得高达 $2.34\times$ 的吞吐飞跃。两者的草稿接受长度误差小于 0.05 个 token，证实加速完全来自于计算与访存体系优化，而非采样分布漂移。图 10 展现了 GDN 模型同样显著的投机验证耗时抑制。

GDN 架构在 Qwen3.5-122B 上的投机内核延迟对比
图 10：Qwen3.5-122B-A10B (GDN) 上的投机验证内核延迟对比（$B=256, L=16$）。ReplaySSM 将验证 8 个草稿的开销压缩至单步基线解码的 $1.11\times - 1.33\times$。

4. 低精度读取内核纯延迟（E.4.2）

图 11 采用 CuTe DSL 内核测算了在非刷新步纯粹降低读取检查点精度时的内核加速比。在 Mamba-2 架构上，将读取精度从 FP32 降至 FP16、FP8、NVFP4 时，非刷新内核分别提速 $1.61\times$、$2.82\times$ 与 $3.83\times$。

低精度读取检查点内核加速比
图 11：在 $B=256, L=16$, 单张 B300 上，改变非刷新步读取检查点精度时获得的纯内核级算子加速比。NVFP4 读取在 Mamba-2 上达到 $3.83\times$ 飞跃。

5. 全模型输出 KL 散度与窗口敏感度分析（E.4.3 & E.4.4）

图 12 与图 13 表明，四款模型在 KL 散度上均表现出高度一致的物理规律。无论采用何种窗口尺寸（$L=4$ 到 64），解耦精度的方案其 KL 散度均纹丝不动贴近 FP32 底噪，而共享检查点方案随上下文拉长迅速劣变。

另外三款模型在 PG-19 上的输出 KL 散度曲线
图 12：PG-19 测试集上另外三款模型的 KL 散度。右侧箭头展示了在最大生成步长处，解耦精度方案对共享 NVFP4 误差的巨大压制效果（降低 $13\times - 29\times$）。

KL 散度随回放窗口尺寸的变化
图 13：在 Qwen3.5-4B 上扫描窗口大小 $L$ 对 KL 散度的影响。解耦双精度配置在所有 $L$ 下均高度平稳，对窗口调参表现出极强的数值鲁棒性。

6. 舍入模式对比（E.4.5）

表 11 显示了为什么循环检查点必须使用随机舍入（SR），而读取检查点适合最近舍入（RTN）。在 1K token 的短文本中，RTN 产生的误差小于 SR；然而当步数累加至 16K 以上时，RTN 的确定性偏置积累使其 KL 散度反超 SR；在 64K~128K 步处，RTN 甚至直接导致两个 GDN 模型的状态数值溢出变为非有限值（NaN/Inf），而 SR 全程收敛保持稳定。

7. 下游多基准评测细节与按长度分桶（E.4.6 & E.4.7）

表 12 给出了 MMLU-Pro、GPQA-Diamond、Math-L5 评测中各模型的生成 Token 长度及准确率衰减。图 14 全景揭示了所有模型与任务的长度分桶规律：在所有 48 个测试分桶中，解耦精度配置（NVFP4 读取 / FP8+SR 递推）有 41 个分桶与 FP32 完全处于同一置信区间内，坚实支撑了长链推理模型的精度稳定性。

全模型全任务按长度分桶的下游准确率表现
图 14：在所有 4 款模型和 3 个高难度基准上，按 FP32 参考生成长度四分位数绘制的准确率曲线。橙色线（共享 NVFP4）随长度骤降，蓝色线（ReplaySSM 双精度）成功挽救了长思维链任务的准确率崩溃。

附录 F：未来方向——基于衰减强度的检查点读跳过机制

在非刷新步中，输出公式为：

$$o_t = \gamma_h S_{\mathrm{ckpt}} q_t + \sum_{j=1}^h \gamma_{j+1:h} v_j (k_j^\top q_t)$$


对于衰减系数极强（$\alpha \ll 1$）的头，经过数步累加后 $\gamma_h$ 会极其微小。一旦 $\gamma_h$ 衰减至低于当前数据类型的有效位所能分辨的绝对阈值，该注意力头在当前步对检查点的访问即可完全置零跳过。在刷新步，其新检查点更新式也完全独立于旧检查点，退化为单纯的写出操作。理论分析表明，该优化能进一步将具有快速局部遗忘特性的 SSM 头的显存访问量削减至传统解码的 $50\%$ 以下。

附录 G：工业界落地与生态采用（Adoption）

ReplaySSM 的核心理念已被顶尖工业界大模型与推理引擎深度采纳：

  • GLM-5.3-Flash：在其生产推理栈中正式集成了 ReplaySSM，并成功扩展至 Kimi Delta Attention (KDA) 算子。
  • Kimi-K3：在其 KDA 架构的投机验证中直接采用了缓存草稿 token SSM 输入的策略，并在技术报告中引用了 ReplaySSM 的并行工作。
  • 开源生态集成：核心算子及系统架构已合并进入开源高性能推理生态，包括 vLLM、NVIDIA TensorRT-LLM、SGLang 以及 FlashInfer 算子库。

方法细节中所引用的参考文献汇总

索引编号 参考文献信息 引用段落与原文观点表述
[Dao & Gu, 2024] Transformers are SSMs: Generalized Models and Efficient Algorithms through Structured State Space Duality. arXiv:2405.21060, 2024. 引言与背景 2.1 节：引入 Mamba-2 状态空间模型递推公式，指出其将衰减限制为每头标量并通过 SSD 将块状计算组织为矩阵乘法；第 8 节：讨论块状 SSM 算法基础。
[Yang et al., 2025] Gated Delta Networks: Improving Mamba2 with Delta Rule. ICLR 2025, 2025. 引言与背景 2.1 节：引入 Gated DeltaNet (GDN) 架构及其基于 Delta 规则的信息更新机制；第 8 节：说明其在保持线性的同时提升状态表达能力。
[Kimi Team, 2025] Kimi Linear: An Expressive, Efficient Attention Architecture. arXiv:2510.26692, 2025. 引言与背景 2.1 节：引出 Kimi Delta Attention (KDA) 架构及其通道级（channelwise）衰减特性。
[Leviathan et al., 2023] Fast Inference from Transformers via Speculative Decoding. ICML 2023, 2023. 引言、第 2.2 节与第 8 节：介绍投机解码的基础范式，指明其一次前向验证多个候选 token 且在验证后必须丢弃拒绝后缀、回滚模型状态的核心约束。
[Kwon et al., 2023] Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP 2023, 2023. 引言与第 3.2 节：指出像 vLLM 这样的主流推理服务系统在处理 SSM 投机解码时，因无便宜回滚手段，只得为每个草稿 token 存储独立的状态快照，造成极高的显存带宽瓶颈。
[Gupta et al., 2015] Deep Learning with Limited Numerical Precision. ICML 2015, 2015. 第 2.2 节与附录 C.6：阐述定点与低精度量化中随机舍入（Stochastic Rounding）相比于最近舍入（RTN）具有数学期望无偏的统计特性。
[Croci et al., 2022] Stochastic Rounding: Implementation, Error Analysis and Applications. Royal Society Open Science, 2022. 第 2.2 节与附录 C.6：提供随机舍入在数值分析中的方差理论边界，为分析误差在时序递推链中的累积扩散建立理论模型。
[Dao et al., 2022] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022, 2022. 第 3.1 节与附录 C.2：指出混合架构中由于 FlashAttention 等内核经过极其深入的工程优化，使得在实测中 SSM 显存瓶颈超过注意力层的时间点推迟到了接近 100K 上下文。
[Wang et al., 2024a] The Mamba in the Llama: Distilling and Accelerating Hybrid Models. NeurIPS 2024, 2024. 第 3.2 节、第 8 节与附录 A：评测现有的 SSM 状态恢复方案，指出其将已提交 token 重放与验证融合，单步产生 $1\text{R}+2\text{W}$ 显存开销并在 GPU 上引入长达 $2T-1$ 的串行递推依赖链。
[Wu et al., 2024] Snakes and Ladders: Accelerating State Space Model Inference with Speculative Decoding. NeurIPS Workshop, 2024. 第 3.2 节、第 8 节与附录 A：对比其提出的两种恢复机制：Activation Replay（产生 $3\text{R}+3\text{W}$ 开销）与 Joint Attainment（增加填充行导致 $3\text{R}+4\text{W}$ 开销），在大并发下无法取得加速。
[Wu et al., 2026] STree: Speculative Tree Decoding for Hybrid State Space Models. NeurIPS 2026, 2026. 第 3.2 节、第 8 节与附录 A：分析其针对树状投机的重放方案，指出其即便剔除多余状态拷贝，算法本身的 $2\text{R}+1\text{W}$ 及串行寄存器常驻依然使其大 Batch 性能弱于简单快照。
[Bischof & Van Loan, 1987] The WY Representation for Products of Householder Matrices. SIAM J. Sci. Stat. Comput., 1987. 第 8 节：指出 Delta 规则模型保持块级并行度所依赖的 Householder 乘积矩阵表示法。
[Joffrain et al., 2006] Accumulating Householder Transformations, Revisited. ACM TOMS, 2006. 第 8 节与附录 B.3：说明用于并行化累积变换的 UT 变换数学源头。
[Yang et al., 2024] Parallelizing Linear Transformers with the Delta Rule over Sequence Length. NeurIPS 2024, 2024. 第 8 节与附录 B.3：说明利用 UT Transform 的上三角矩阵求解技术来并行化 Delta-Rule 训练的理论依据，ReplaySSM 借用该思想构建了草稿验证的三角求解系统。
[Pan et al., 2025] Marconi: Prefix Caching for the Era of Hybrid LLMs. MLSys 2025, 2025. 第 8 节：区分 Prefill 阶段的 Prefix Caching（复用历史以降低 Prefill 算力）与 ReplaySSM（降低 Decoding 阶段状态访存），证明两技术完全正交互补。
[Tianqi et al., 2025] Q-Mamba: Towards More Efficient Mamba Models via Post-Training Quantization. ACL 2025, 2025. 第 8 节：评析其提出的解耦尺度量化方法，指出其误差分析虽然注意到了衰减项，但在推理时依然让量化状态反馈进入循环，且仅评测了 2048 长度。
[NVIDIA, 2026] Nemotron 3 Ultra: Open, Efficient Mixture-of-Experts Hybrid Mamba-Transformer Model for Agentic Reasoning. arXiv:2606.15007, 2026. 第 8 节：论文作者参与的报告，探索了周期性检查点量化，但由于未切断输出与递推的精度依赖，仍无法在根本上阻断误差时序传播。