一个反直觉的事实:SSM 层成了新的瓶颈
昨天 NVIDIA 发布了一篇技术博客 ReplaySSM,与普林斯顿 Dao AI Lab(就是 Tri Dao 的组)合作,论文题为 “Cache SSM Inputs, Not State”——缓存 SSM 输入,而非状态。
这个标题本身就在向 KV Cache 致敬,也在宣告一件事:混合架构模型的推理优化,已经进入了深水区。
先交代背景。今天的前沿模型几乎清一色是混合架构:大量 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 Cache,从而让解码变快。
但 ReplaySSM 开篇就给了一个反直觉的测量结果:在真实服务的 batch size 下,直到 100K token 上下文,SSM 层——而不是注意力层——才是解码延迟的最大组成部分。
那个承诺让解码变快的东西,恰恰成了慢的原因。这值得深挖。
为什么 SSM 层反而慢:带宽,不是算力
以 Nemotron 3 Ultra 使用的 Mamba-2 层为例,每个解码步执行这样的递推:
S_t = α_t · S_{t-1} + v_t k_tᵀ
o_t = S_t q_t
每一步都要:从 HBM 加载完整状态 S_{t-1},施加衰减和一次秩一更新,计算输出,再把 S_t 写回显存。注意这个操作的特征——读写的数据量巨大,浮点运算却很少。这是一个典型的访存受限(memory-bound)内核,被 HBM 带宽卡死,GPU 的算力大量闲置。
两个因素叠加放大了这个问题:
第一,层数比例失衡。 Nemotron 3 Super 中 SSM 层是注意力层的 5 倍(40 比 8)。SSM 单层略慢一点,乘以 5 倍就是统治性的开销。
第二,状态本身不小。 一层的 FP32 状态是每序列 4 MiB,而注意力层的 FP8 KV 条目是每 token 0.5 KiB。换算下来,一层 SSM 的状态等于注意力层 8K token 的 KV Cache。上下文在 8K 以内时,SSM 的「固定状态」甚至比 KV Cache 还大。
所以现实很讽刺:SSM 理论上解决的是长上下文的显存问题,但在中等上下文 + 大 batch 的典型服务负载下,它变成了带宽瓶颈。学术论文里画的那种「Mamba 解码延迟随上下文平稳不变」的曲线,在真实 serving 场景中要重新审视。
更麻烦的是:SSM 天然排斥两项加速技术
如果只是慢也就罢了,SSM 还让推理系统里最可靠的两项加速技术近乎失效。
投机解码(speculative decoding)的回滚难题。 投机解码让一个小模型起草 T 个 token,目标模型一次前向验证,再回滚到最后被接受的位置。注意力层回滚天然简单——KV Cache 移个指针就行。但 SSM 一旦把草稿 token 的信息「总结」进了状态,这个更新没有逆操作。vLLM 的 workaround 是给每个草稿 token 存一份完整状态快照,验证后挑匹配的那份——每步一次读、T 次写,全发生在访存受限的路径上,而且快照挤占了能装进显存的请求数。结果是:对 Transformer 稳定有效的投机解码,在 SSM 上几乎带不来收益。
状态量化的误差累积。 KV Cache 量化到 FP8 甚至 4-bit 已经是常规操作,因为一条量化后的 KV 条目只读不写,误差不会自我循环。但 SSM 状态既是本步输出的来源,又是下一步更新的输入——量化的舍入误差每一步都重新进入递推,随生成长度累积。这就是为什么 Qwen3.5 用 NVFP4 权重、FP8 KV Cache,SSM 状态却保守地留在 FP32。
ReplaySSM 的解法:把「状态」拆成「检查点 + 输入缓冲」
ReplaySSM 的核心观察很简单也很漂亮:上面那个递推式可以在 h 个 token 的窗口上展开——
S_t = γ_h · S_ckpt + Σ_{j=1..h} γ_{j+1:h} · v_j k_jᵀ
也就是说,状态可以等价地表示为「窗口前的检查点 + 最近若干条输入记录」。传统的逐步解码走左边,ReplaySSM 走右边。
具体机制:把 SSM 状态表示为一个检查点(checkpoint)加一个小的输入缓冲区(存最近 L=16 条 α、v、k 输入记录)。绝大多数「非冲刷步」只做一件事:把当前输入追加到缓冲区,用「检查点 + 缓冲区」直接计算输出——不写任何状态。当缓冲区满时,一次「冲刷步」(flush)把缓冲的输入折叠进检查点,写一次 HBM,清空缓冲。数值上与传统解码等价,只差浮点误差。
这一个拆分,同时解决了前面所有三个问题:
带宽减半。 传统解码在 L 个 token 上读状态 L 次、写状态 L 次;ReplaySSM 每步仍读检查点(输出需要它),但只在冲刷时写一次。非冲刷步只写几条向量而非整个矩阵,主导性的状态流量大约砍掉一半。batch 256 下,内核加速 1.46×–1.80×,整个解码步加速 1.20×–1.47×。
回滚变成移指针。 投机解码时,草稿 token 的输入直接追加进缓冲区;验证通过后推进缓冲区指针即可「提交」——和 KV Cache 的回滚逻辑完全同构。不再需要任何状态快照。冲刷只折叠已提交的记录,永远不用等接受长度。
量化误差不再循环。 既然输出计算和状态更新分离了,两者可以用不同精度:递推用高精度检查点(FP32,或带随机舍入的 FP16 就够),非冲刷步的读取用低精度副本(NVFP4/FP8)。低精度副本永远不被读回递推,舍入误差只影响当前窗口的输出,不跨步累积。在 Qwen3.5-4B 上生成 128K token,共享 NVFP4 检查点的 KL 散度是 0.136,NVFP4 读取 + FP32 递推只有 0.0038——误差缩小 35 倍且不随长度增长。GPQA-Diamond 上,共享 NVFP4 检查点最多掉 31 个百分点,而分离精度配置基本追平 FP32。
更实际的是落地速度:ReplaySSM 已合入 TensorRT-LLM、vLLM、SGLang、FlashInfer,并且已经在 GLM-5.3-Flash 和 Kimi-K3 的生产服务中运行。这不是一篇论文,是正在线上省钱的代码。
深层启示:系统研究的主战场正在转移
如果只把 ReplaySSM 当成一篇「内核优化」文章看,就错过了它真正的信号。
第一,模型架构的创新红利,最终要靠系统研究兑现。 混合架构是 2025–2026 年最重要的架构共识——它被认为兼得注意力的召回能力和 SSM 的解码效率。但 ReplaySSM 证明,如果不重做状态管理,混合架构在真实负载下的效率优势可能根本兑现不了。「设计一个新架构」和「让新架构在 serving 栈里真的快」之间的鸿沟,正由这类工作填补。Kimi-K3 论文把 ReplaySSM 引为 concurrent work,说明各家头部实验室都独立撞上了同一堵墙。
第二,优化范式在向「账本式思维」迁移。 KV Cache 的 PagedAttention、KV 复用、前缀缓存,本质上都是把「存什么、何时存、存多精」变成可独立调度的决策。ReplaySSM 把同样的思维带进了 SSM:状态不必是原子实体,可以拆成检查点和增量日志——这正是数据库和存储系统几十年前就走过的路(WAL、checkpoint/redo log)。推理系统正在重演数据库系统的历史,这个视角比任何单点技巧都更有预测力:接下来一定会出现更激进的跨请求状态共享、检查点分层存储(HBM/DDR/NVM)、以及针对 SSM 状态的服务质量调度。
第三,投机解码正在成为架构设计的硬约束。 当回滚能力由「状态表示方式」决定,任何无法廉价回滚的状态设计都会在投机解码普及的时代吃亏。可以预期未来的架构论文里,「rollback cost」会像参数量一样成为标准汇报项。KDA、GDN 这类 delta-rule 变体能写成秩一更新形式,恰好都能套用 ReplaySSM——这不是巧合,而是「输入可重放」正在成为新架构的隐性设计准则。
第四,对国产模型生态是个注脚。 论文致谢级别的信息是:这项研究的生产验证发生在 GLM-5.3-Flash(Z.ai 公开了基础设施报告,描述了把它加进 KDA 解码内核的过程)和 Kimi-K3 上。推理侧的工程能力——而不只是训练侧的模型能力——正在成为中国模型厂商在价格战里的真实护城河。每百万 token 的定价背后,是这类把带宽流量砍半的硬功夫。
结语
「Cache SSM Inputs, Not State」这个标题,放在更大的语境里读,是推理系统领域的一个宣言:状态不是必须的,可重放的输入才是。 数据库行业用几十年的实践证明了 log-structured 思维的威力,现在轮到推理栈了。
对从业者,两个可操作的建议:如果你的服务栈是 vLLM/SGLang/TensorRT-LLM 且跑混合架构模型,升级到包含 ReplaySSM 的版本可能是目前性价比最高的一步;如果你在做架构研究,把「状态的可回放性、可量化性、可共享性」写进设计目标,比再发明一个新的注意力变体更有机会成为下一个被引用的基础设施。
混合模型的时代到来了——但只有当它的状态被正确存储时,才算真正到来。
参考: