BloomBloom
发起研究

ReplaySSM：SSM 缓存输入换投机解码，批量不够大没收益

「省显存」的 SSM，为什么反而拿不到 Transformer 白拿的投机解码与量化？

在同一块 Blackwell Ultra B300 上，把一个 120B 混合模型的单步解码时间拆开看，到十万 token 这个量级，占最大一块的既不是 attention，也不是 MoE，而是那个号称最省显存的 SSM 层。NVIDIA 博客

这个结果反直觉，因为 SSM 的招牌卖点一直是省。它用一份固定大小的递归状态把整段历史压住，显存不随序列长度增长，推理时间线性爬升。Transformer 走的恰好相反，每一层都把历史 token 的 key 和 value 显式留在 KV cache 里，显存随上下文一路膨胀。按这个说法，长上下文本该是 SSM 的主场。

可真实服务里的账不是这么算的。SSM 层的墙钟时间反而最容易先到顶，而原因不是它算得慢。Mamba-2 每步更新只是一次逐元素的衰减加一次 rank-1 外积，算术强度大约 1 FLOP/byte，而 H100 的 Tensor Core 能跑到约 300 FLOP/byte。Tri Dao 博客 算术量相对搬运量几乎可以忽略，所以它根本不是算力受限，而是被显存带宽绑死的搬运受限。

顺带一提，rank-1 外积指的是一个列向量乘一个行向量，得到一个矩阵，而且这个矩阵的秩为 1。在整个更新式里，它是唯一真正写入新内容的那一项，前面那项只是把旧状态按比例缩一下。

问题就出在搬运上。Transformer 的 KV cache 是一份显式的历史张量，好处是这些张量可以反复被利用。投机解码里草稿被拒绝时，回滚只需要把 KV cache 的指针往回挪一格；把缓存量化到低位宽时，压小的是一份静止的数据。这两样红利，恰好都建立在「缓存里存的是可指名、可回退的历史」之上。

SSM 没有这份历史。它每步把当前 token 吸收进状态，下一步再整块读回来，算完又整块送回去。作者把代价归成三条既独立又纠缠的线：状态 I/O 远高于计算；状态是「输入已经被吸收进去」的有损摘要，没有可回退的指针；同一份低精度状态既算输出又喂下一步，舍入误差每步都重新进入递推。Tri Dao 博客

论文 Figure 1：conventional 解码与 ReplaySSM 的全局对照，下半部预先摆出三类挑战

图注：上半部左右对照两种解码：左边 conventional 每步都从显存载入整块状态、算完再写回；右边 ReplaySSM 把状态拆成 checkpoint 与一个缓冲，多数步只追加小记录，只有 flush 步才写回显存。下半部三块小图对应三条挑战：SSM 层在混合模型里到 100K token 仍是最大延迟块、状态回滚在常规投机解码里代价高、低精度 checkpoint 只服务输出而递推走高精度。ReplaySSM

这张图先只给轮廓，不解释机制。因为真正的主线问题在末尾：如果每一步都把状态写回显存，而写回的唯一用途只是让下一步原样读回来再算同一件事，那我们到底为什么要存它？论文给的答案是，我们一直存错了东西。真正该留下的不是每步的状态，而是产生这个状态的输入。

递推式决定了能缓存什么：Mamba-2 的加法，与 GDN 的加法加减法

要想改存输入，第一件事得弄清「输入」在这里指什么。答案不在工程习惯里，在每一类 SSM 自己的递推式里。

Mamba-2 的 SSD 层里，状态更新写成 $S_t = a_t S_{t-1} + v_t k_t^{\top}$。$a_t$ 是一个标量衰减因子，$v_t k_t^{\top}$ 是一次 rank-1 外积写入。Mamba-2 这个式子有一个漂亮之处：同一层既能写成线性时间的递归，也能写成二次时间的掩码注意力，两者是同一张张量收缩的两种顺序。论文里把输入序列 X 对应到注意力里的 V，把 B 对应到 K，把 C 对应到 Q。

这层对偶不只是数学欣赏。它顺带说清了缓存的对象：要重放 Mamba-2 的一步，需要的输入就是那几项，一个标量 $\alpha$、一个 $v$、一个 $k$，别的什么都不用留。

Gated DeltaNet 比 Mamba-2 多出了一项。它的递推是 $S_t = \alpha_t S_{t-1}(I - \beta_t k_t k_t^{\top}) + \beta_t v_t k_t^{\top}$。Gated DeltaNet 多出来的 $(I - \beta_t k_t k_t^{\top})$ 是一个 Householder 形式的修正项，作用是在 $k_t$ 这个方向上把旧记忆擦掉，再写进新值。这句话里「Householder」指的是用一次向量投影构造出的正交或收缩变换，这里是沿着 $k_t$ 方向做定向擦除。一句话概括，Mamba-2 是纯加法，GDN 是加法再减一笔。

Gated DeltaNet 的 H1、H2 混合架构与 Block Design 内部数据流

图注：三幅子图并排。左边 Gated DeltaNet-H1 把 Gated DeltaNet、MLP、SWA、MLP 逐层堆叠并带残差；中间 Gated DeltaNet-H2 是更宽的混合层栈；右边的 Block Design 展开内部细节，底部四个线性投影分别产生 q、k、v 与 α、β，q/k 经卷积与归一化，v 经卷积，α 与 β 单独走一路，再送入 Gated Delta Rule 并接归一化与输出投影。Gated DeltaNet

差别就卡在这个修正项上。它依赖上一步状态读出来的结果，所以照搬 Mamba-2 只缓存 $(v, k)$ 会失败。真要回放，就得先把上一步状态重建出来，串行依赖一点没省，等于白折腾。作者给出的改法是改存修正之后的量 $u$，GDN 的更新就退化成 Mamba-2 那种「衰减加外积」的干净形式。Tri Dao 博客

所以缓存什么从来不是工程口味，而是被递推式钉死的。先数清更新式里有几项、哪一项依赖旧状态，才知道能省下什么。Mamba-2 交出的答案是 $(\alpha, v, k)$，GDN 只能交 $(\alpha, u, k)$。既然每一步真正新增的只有这么几个小项，那把状态换成「一个旧 checkpoint 加一串小记录」看上去就顺理成章了。接下来要回答的是两个更硬的问题：这在数学上真的等价吗，省下的又到底是哪一笔？

把状态拆成「一个 checkpoint + 一段近期输入」，标准解码先快起来

论文的核心观察写在 §3 的一条展开式里。当前状态等价于一个较早 checkpoint 的衰减项，加上最近若干步 rank-1 更新之和：

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

这条式子真正值钱的地方在于，右边每一项都是输入侧就能拿到的东西，不需要先把中间状态物化出来。

数据库的做法能帮上忙。写事务时不会每次把整张表重写一遍，而是先往日志里追加一条小记录，隔一段时间把内存里的改动整体落盘成一个 checkpoint；真要恢复时，从最近一次 checkpoint 开始重放日志就行。ReplaySSM

ReplaySSM 的把戏和这个结构一一对应。被缓存的对象不再是每步的状态，而是一份 $S_{\mathrm{ckpt}}$ 加一个输入缓冲 $B$，$B$ 最多放 $L$ 条记录，论文全部实验里 $L$ 固定取 16。ReplaySSM 大多数步只做两件事：把当前 token 的 $(\alpha, v, k)$ 追加进 $B$，然后从 $S_{\mathrm{ckpt}}$ 和 $B$ 现算出现用状态去算输出，算完不写回。只有缓冲塞满时才触发一次 flush，把整块状态真正写回显存并清空缓冲。前者叫 non-flush 步，后者叫 flush 步。

举个具体的样子。设缓冲窗口 $L=16$，此刻显存里的 checkpoint 对应第 100 步的状态。接下来的第 101 到 115 步，每一步都只往缓冲追加一条记录，输出从这份 checkpoint 加现有记录现算，没有哪一步回写显存。到第 116 步，缓冲正好凑满 16 条，触发一次 flush，把状态一次性固化成一个新 checkpoint 并清空缓冲，新的窗口从第 117 步重新开始。十六步里，真正动到显存里那块状态的动作只有一次。

Baseline 解码每步双向读写显存，ReplaySSM 只在缓冲满时写回一次状态

图注：左右对照。左栏 Baseline Decode 画的是三个实心状态方块，底部的 HBM 长条与每一个方块之间都有双向箭头，表示每步都要读写整块状态。右栏 ReplaySSM 把状态画成虚线方块，表示按需重建；HBM 里只留下一个 checkpoint 和若干缓存输入 $v_j k_j$，箭头标注「从近期输入重建」，并在末步标出「写回状态」。底部图例区分实心状态、已缓存输入与按需重建的状态。Tri Dao 博客

把窗口拉长到 $L$ 步看总账，省下的那笔就很清楚了。两种做法都要读 $L$ 次状态，差别只在写：传统做法每步都写，一共 $L$ 次；ReplaySSM 只在 flush 那一步写一次。占主导的状态搬运量因此大致减半。ReplaySSM

这里还有一条顺手少算的账。算输出时，本来要先把状态物化出来再乘上 query，但这一路其实只需要 $(v k^{\top}) q$ 这个结果。把括号挪一下，先算标量 $k^{\top} q$，再拿它去缩放 $v$，答案一丝不变，而乘法的规模小得多。ReplaySSM 这条改写只在算输出时成立，没法用来更新状态，所以它天生就是给 non-flush 步用的。

这笔账在实测里兑现了。批量 256、$L=16$ 时，ReplaySSM 的 SSM kernel 在四个模型上加速 1.46 到 1.80 倍，整步解码加速 1.20 到 1.47 倍。ReplaySSM

但请注意另一头：批量 1 时几乎没有收益，只差几个百分点。原因不难想。解码是逐 token 推进的，批量越大，同一份状态搬运就能摊销到更多请求上，SSM 在整步里的占比也越高，省下的那点写才越来越显眼。批量 1 时整步本来就延迟受限，少写几次几乎看不出来。这条边界后面还会再出现一次。

标准解码只是第一笔账。真正有意思的是，一旦缓存里存的是可以重放的输入，那些原本只属于 Transformer 的红利，会不会也跟着回来？

投机解码的回滚，为什么能从「存 N 份快照」变成「挪一下指针」？

先说清投机解码在做什么。它让一个小而便宜的草稿模型一口气猜若干个 token，再让目标模型一次并行地验证全部猜测，猜对的整段接受，猜错的位置由目标模型补一个修正 token。原始论文证明这个过程对任意草稿模型都保持输出分布完全不变。Speculative Decoding

要跑起来，它需要一个能力。验证结束后，如果只接受了前 $a$ 个草稿，目标模型的状态必须能便宜地退回「接受段末尾」对应的那个前缀。

在 attention 这一侧，这件事几乎免费。KV cache 存的就是历史前缀本身，回滚只是把指针往回挪一格。放到 SSM 这一侧就麻烦了，状态是「输入已经被吸收进去」的有损摘要，没有可指的 token 历史。现行系统的通用做法是为每一个 draft token 多存一份完整状态快照，验证 T 个 token 就等于多 T 次状态写。Tri Dao 博客 这笔开销恰好压在本来就 memory-bound 的那条路径上。

Attention 靠挪 KV cache 指针回滚，SSM 把输入摘要进状态后无法撤销

图注：左右对照。左半是 Attention 侧，五个方块中前三个是被接受的草稿，后两个灰字方块是被拒绝的草稿，一条蓝色箭头标着 pointer 向左指向被拒位置，另一条虚线箭头标着 moved back，回滚只是指针回移。右半是 SSM 侧，五个虚线方块注着「原始输入已被丢弃」，一条 summarize 箭头把它们汇成一个实心状态方块，从状态指回输入的箭头被红叉截断，旁标 Irreversible。Tri Dao 博客

ReplaySSM 把这一步整个换掉。T 个草稿的记录一起追加到缓冲里，用同一份 checkpoint 和同一个窗口一次性算出 T 个输出；验证完拿到接受的长度 $a$，回滚就是把缓冲的有效长度直接设成 $a$，等价于从 T 往回退 $T-a$ 格，被拒的记录等下一次追加时覆盖掉。ReplaySSM 不需要任何状态快照，也不需要 restore。

拿 $T=4$、$a=2$ 举个例。草稿模型先猜出 4 个 token，验证后发现前 2 个被接受、后 2 个被拒。传统做法要在这里从 4 份状态快照里挑回第 2 份，再把它装回推理状态；ReplaySSM 只是把缓冲的有效长度从 4 设回 2，也就是往回退 $T-a$ 格，被拒的那两条记录留在原地，等下一次追加时自然盖掉。同一个动作，一边是复制整块状态，另一边是改一个整数。

实测把这条差别翻译成了吞吐。两个 120B 级模型、GSM8K prompt、$T=4$、批量 512 时，ReplaySSM 达到 vLLM 常规投机解码吞吐的 1.85 倍与 2.10 倍，接受长度持平。NVIDIA 博客 更能说明问题的是对照组：同一批量下，常规投机解码在 Qwen3.5-122B-A10B-NVFP4 上已经掉到标准解码之下，越投机越慢。另一个证据是验证 kernel 的延迟，$T=8$ 时常规做法飙到标准步的 4.66 倍，ReplaySSM 仍维持在 0.89 到 1.21 倍之间。ReplaySSM

GDN 上还要多一步。delta rule 的修正项带着串行依赖，即使缓存了输入，也没法直接把 T 个草稿一起算出，得先做一次 $T \times T$ 的三角求解拿到所有修正项 $u$，之后才全部是矩阵乘。Tri Dao 博客

到这里，前两条路径救的都是「搬运」。第三条路径要救的是质量：把状态压到 4 位，为什么在长文本上会塌？

把状态压到 4 位，为什么长文本不再「一步错、步步错」？

失败模式讲起来很直接。如果同一份低精度状态既算输出又喂下一步，那么每步舍入产生的误差都会重新进入递推，随序列增长被一步步放大，长生成的质量就会塌掉。

复印机是个好用的比方。每一次都从原件重新复印，误差不会积累；但如果每次都拿上一次的复印件再复印一次，几轮之后字迹就糊了。低精度状态直接喂递推，就是后一种复印。

ReplaySSM 的做法是让 checkpoint 存两份。一份高精度的 $S_{\mathrm{ckpt}}^{\mathrm{HP}}$ 专门承载递推，一份低精度的 $S_{\mathrm{ckpt}}^{\mathrm{LP}}$ 只用来在 non-flush 步算输出。每次 flush 都从高精度那一路重新算出状态，再独立压出一份低精度副本。ReplaySSM 低精度副本于是变成只读的，它的舍入误差没有机会被下一步重新放大。

共享量化状态的 KL 随长度上升，分离精度后紧贴 FP32 下限

图注：左图横轴是生成长度、纵轴是输出 KL，共享低精度状态的曲线随长度单调上升，而「低精度只读 checkpoint + 高精度递推」的曲线几乎平直地贴着 FP32 下限。右图是 GPQA-Diamond 准确率按参考长度分箱，共享低精度状态在最长两档掉得最多，分离精度后损失大幅收窄。ReplaySSM

数字很有说服力。在 Qwen3.5-4B 上做 PG-19 teacher forcing，长度拉到 128K token 时，「低精度只读、递推走 FP32」把输出 KL 从 0.136 压到 0.0038，约 36 倍改善。ReplaySSM 下游准确率上，共享低精度状态在 GPQA-Diamond 最长的两个长度区间里掉 20.9 与 28.2 个点；换成「低精度只读 + 高精度递推」后只掉 2.0 与 3.4 个点。ReplaySSM

这里有个容易被忽略的细节：伤害集中在长生成上。Nemotron-3-Super 在 GPQA-Diamond 上最长那档参考长度的中位数约 61K token，而 MMLU-Pro 与 Math-L5 大约只有 8K，所以短任务上几乎看不出问题，误差要长到一定程度才现形。ReplaySSM 另外要提醒一句，36 倍这个数字只用在这个特定组合上，换模型、换精度、换长度都不能直接外推。

三条路径各自成立之后，剩下的问题就很工程了：它到底在真实模型和真实服务栈上跑起来了吗，边界又在哪里？

它到底跑起来了、能替代谁：四个生产模型与两条边界

论文的实测矩阵是四个生产混合模型乘两块 GPU，都不做张量并行。NVIDIA 博客

模型SSM 类型精度硬件SSM / attention 层数
Nemotron-3-Nano-4BMamba-2BF161×H10021 / 4
Nemotron-3-Super-120B-A12BMamba-2NVFP41×B30040 / 8
Qwen3.5-4BGDNBF161×H10024 / 8
Qwen3.5-122B-A10BGDNNVFP41×B30036 / 12

落地证据也不是伪代码。仓库里是四个真实的 Triton kernel，Mamba-2 与 GDN 各有一套标准解码与投机解码，配套一个长度 16 的缓冲窗口和两个互斥开关。ReplaySSM repo Mamba-2 标准解码那一块已经合入 vLLM 主线，其余部分在 vLLM、TensorRT-LLM、FlashInfer 上处于合并路径或已落地，SGLang 也在集成名单里；生产系统层面，据 NVIDIA 博客披露，GLM-5.3-Flash 与 Kimi-K3 已在各自服务栈中采用 ReplaySSM（厂商单方声明，尚无独立验证）。NVIDIA 博客

从厂商视角给这张矩阵收个口，也能看清它为什么被接住。NVIDIA 把它和 SSM 解码的三类挑战摆在同一张图里：状态访存流量对应输入缓冲加 checkpoint，投机回滚对应基于指针的回滚，量化误差对应分离精度的 checkpoint。三条挑战各自对应一条路径，这正是它被多个服务栈同时接纳的原因。NVIDIA 博客

不过它并非万能加速器，两条边界要说清楚。

第一条是批量门槛。这套方法的收益本来就来自状态写变少，而状态写只有在批量大、SSM 占比高时才成为瓶颈。批量 1 时几乎没有收益，这也是为什么它是 serving 侧的优化，而不是单请求延迟的优化。ReplaySSM

第二条是它只支持 chain，不做树形。同期工作 TreeWY 把现有的 SSM 投机验证方法分成四类，ReplaySSM 被归入「推迟物化 + rank-1 缓存」这一类，与它在验证窗口内共享同一个三角求解内核，差别只在于状态什么时候被物化；树形方向是 TreeWY 补的那一块。TreeWY 两者在 Qwen3.5-35B-A3B 上做过同 baseline 对照，各自归一化到自己的 store-all 基线后，ReplaySSM 的吞吐是 1.06 到 1.38 倍，TreeWY 是 0.94 到 1.19 倍。TreeWY TreeWY 作者把这几点差距归因为调度而非数学差异，并明确说明这只是一个尚未用原型验证过的假说。

所以 ReplaySSM 的价值不在某个单点数字上。它做的是把 SSM 解码重新拉回「值得做投机解码、值得做量化」的工程基线，让三条此前互相独立的优化路径第一次同时可用。它不是一个万能加速器，它只是终于把缓存里该放的东西放对了。

想深入可以读原文：ReplaySSM: Cache SSM Inputs, Not State — arXiv:2610.13778

参考文献

ReplaySSM：SSM 缓存输入换投机解码，批量不够大没收益