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

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

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

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

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

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$ 的窗口中：

标准解码加速比对比
图 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)$$

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 中维护两套不同精度的检查点：

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）模式运行：

2. 硬件与软件配置

3. 基线与量化设置

4. 评估数据集与协议

实验结果

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

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

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

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

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

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 实现了数倍的延迟压缩。

附录 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$$

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

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

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

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

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

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$ 条缓冲记录与当前输入的单步计算：

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$：

附录 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：

附录 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 的核心理念已被顶尖工业界大模型与推理引擎深度采纳：

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

索引编号 参考文献信息 引用段落与原文观点表述
[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 节：论文作者参与的报告，探索了周期性检查点量化，但由于未切断输出与递推的精度依赖，仍无法在根本上阻断误差时序传播。