一句话结论ReplaySSM 做的事情可以概括成一件事别再每步把 SSM 的递归状态写回显存。它把状态重新表示成「一个检查点 最近若干步的输入环形缓冲」只在缓冲写满时做一次 flush把缓冲折叠进检查点。结果是状态流量大致减半投机解码不再需要为每个草稿 token 存一份状态快照。论文是《ReplaySSM: Cache SSM Inputs, Not State》arXiv 2610.137782026-10-09 提交、10-12 公开作者来自 Princeton 的 Dao AI Lab 与 NVIDIAZe-Wei Liou、Harris Nover、Amey Naik、Andrii Skliar、Igor Shovkun、Ido Shahaf、Tri Dao。参考实现在 vLLM 的分支里代码在 GitHub。三个瓶颈其实来自同一个动作混合模型Nemotron 3、Qwen3.5–3.8、Kimi-K3、GLM-5.3-Flash 这类用多数 SSM 层搭配少量注意力层。Mamba-2 层的单步递归是S_t α_t · S_{t-1} Δ_t · (v_t k_tᵀ) # 更新 y_t S_t · q_t # 读出一次解码步骤要从 HBM 读整个状态矩阵 → 施加衰减和一次 rank-1 更新 → 算出输出 → 再写回去。算术量极小、字节搬运很多所以这个 kernel 是访存受限而不是算力受限。NVIDIA 博客给的两个数字很直观一个 Mamba-2 层的 FP32 状态每序列 4 MiB相当于 8000 个 token 的 FP8 KVNemotron-3-Super-NVFP4 的 SSM 层数量是注意力层的 5 倍。在 batch 256 下10 万 token 以内的上下文里 SSM 层是解码延迟的最大项。这个「读—改—写」的融合动作同时造成了三个问题状态流量窗口 L 步里读 L 次、写 L 次。投机解码无法廉价回滚注意力只要把 KV 指针往回挪被拒的草稿自然失效但 SSM 状态是把历史压成一个矩阵更新不可逆。vLLM 的现行做法是为每个草稿 token 存一份完整状态快照等于在一条本就访存受限的路径上每步多出 T 次写并且快照会占显存、压低并发。状态量化误差会进循环KV cache 被读但不写回自身所以容忍低精度量化后的 SSM 状态要作为下一步更新的输入舍入误差会沿递归累积。核心做法状态 检查点 输入缓冲关键观察是递归本身可以展开。把窗口内 h 步展开后S_t γ_h · S_ckpt Σ_j γ_{j1:h} · v_j k_jᵀ左边是「汇总路线」——加载上一步的摘要、用新输入更新它右边是「历史路线」——从最近输入窗口重建。常规解码一直默认走左边但这不是必须的。于是 ReplaySSM 的行为变成非 flush 步把当前的 (v, Δ, k) 追加进缓冲输入输出从检查点 缓冲直接算出不写回状态矩阵只写下几个向量。flush 步每 L 步一次把缓冲里的输入折叠进检查点写回一次完整状态清空缓冲。检查点每步还是要读算输出需要它但写只发生 1/L 次主导的状态流量从 8dn 量级降到 4dn 量级。日志里还有一个很实用的性质L 1 时缓冲恒空、每步都是 flushkernel 退化回基线可以逐位对齐验证。这让正确性测试有了锚点。注意 ReplaySSM 在数学上无损但在浮点上不是逐位一致——重建的 GEMM 跑在 Tensor Core 上fp32 用 TF32、bf16 用 bf16会重新结合递归差异在 ~1e-4TF32/ ~1e-3bf16量级。另一个副产品是output-only 解码从零状态出发o (v kᵀ) q可以写成v (kᵀ q)。左边先造出状态再读它右边只形成标量再缩放向量从不物化状态矩阵。每头 FLOPs 从 2(h1)dn 变成 2(h1)(dn)d n 128 时约小 64 倍而且 k、q 在同一组头之间共享k_jᵀ q可以每组只算一次。这就是为什么这个方法在投机解码的验证阶段价值最大——T 会成倍放大每步算术。投机解码回滚变成挪指针投机解码的流程变成把所有 T 个草稿 token 的输入直接追加进同一个缓冲验证后按接受前缀推进缓冲指针即可提交。指针之后的记录被忽略、会被下一次追加覆盖——和 KV cache 的语义一致。没有逐草稿的状态快照也没有状态恢复。验证本身也变成并行的所有草稿 query 都对同一个检查点和同一个缓冲窗口求值唯一的区别是因果掩码。kᵀq 变成一次 GEMM值的加权和变成掩码下的另一次 GEMM整段验证从「长度 T 的串行循环」变成 GEMM 一个 T×T 三角求解。这里有个容易踩的坑论文专门强调flush 要提前一个窗口触发。自然条件是h T LReplaySSM 用的是h 2T L。原因是若某步刚好落在 h T L − 1 且草稿全被接受下一步只剩一个空位草稿窗口会被截断成 1接受量就塌了。提前一个窗口能保证每步都至少有 T 个空位。GDN 为什么要缓存的东西不一样Mamba-2 缓冲 (v, k) 就够了。Gated DeltaNetQwen3.5 系多了个「擦除」修正项u β (v − S k) → S_t α_t S_{t-1} u_t k_tᵀ计算 u_t 需要 S_{t-1}所以缓存原始 v_t 没用——回放时还是得拿到每一个中间状态串行依赖又回来了。修法是缓存 u 而不是 v这样状态照样能展开成「衰减后的检查点 Σ u_j k_jᵀ」。另外 GDN 一步需要状态在 k 和 q 两个向量上的读出所以走的是「状态输出」路线而不是纯 output-only投机解码用分块 delta-rule把 T 个修正量并成一次三角求解(IA)⁻¹。数据与边界评测覆盖两个家族、四个模型Nemotron-3-Mamba-24B dense 在 1×H100120B 级 MoE 在 1×B300与 Qwen3.5-GDN4B / 122B 级。实现在 vLLM 上开启 CUDA GraphSSM 状态 FP32、缓冲向量 BF16。标准解码batch 256kernel 快 1.46×–1.80×整个 decode step 快 1.20×–1.47×Dao AI Lab 的博客口径是最多 1.48× 端到端大 MoE 上 1.43×。端到端增益小于 kernel 增益因为注意力、GEMM 等部分没有变化。投机解码batch 512GSM8K 提示T 4吞吐达到 vLLM 投机解码的 1.85× 和 2.10×接受长度持平对比 vLLM 标准解码则是 1.87×–1.96×。拆开看kernel 层面 2.28×–3.33×传导到 verify 前向是 1.23×–1.69×再到完整 decode step 是 1.20×–1.58×。并发方面在固定 HBM 预算下窗口 4基线预分配的状态快照把最大解码 batch 压到约 1/4ReplaySSM 捞回其中 3.0×–3.3×。缓冲长度是个倒 U 型。短缓冲 flush 太频繁写回成本高长缓冲每步读得更多最终把 kernel 推成算力受限。博客给的甜点是 L 16Dao AI Lab 博客中 Nemotron-3 用 8、Qwen3.5 用 16是他们各自的最优。什么时候没用batch size 1 时其他延迟来源占主导测试里 ReplaySSM 没有加速kernel 甚至可能更慢此时 kernel 是启动开销受限而非带宽受限。收益随 batch 增大而增大因为 SSM 在整步中的占比在上升。低精度读写分离是第三个收益。非 flush 步读状态只是为了算输出flush 步才更新检查点两者可以用不同精度递推保持高精度检查点另存一份低精度副本专供非 flush 读取。低精度副本永远不会被读回递推它的舍入误差只影响当前窗口的输出。数据上Mamba-2 非 flush kernel 读 NVFP4 状态比读 FP32 快 3.8×在 Qwen3.5-4B 上生成到 128K token 时NVFP4 读 FP32 递推把相对 FP32 的 KL 从 0.136 降到 0.0038。GPQA-Diamond 上共享 NVFP4 检查点最多掉 31 个百分点而 NVFP4 读 FP16/FP8 递推能保住四个模型中的三个在 FP32 的约 2 个百分点以内。与已有方案的取舍在 SSM 投机解码上此前已有工作不再为每个草稿 token 保存独立递归状态论文把它归到 Mamba-in-the-Llama 一类。ReplaySSM 的差别在于不只是省显存而是把输出计算与状态更新彻底解耦因此同时拿到「便宜回滚 并行验证 读写精度分离」三件事而不只是其中之一。工程侧最诚实的对照来自 SGLang 的移植 RFC#28511在 H20-3e 上 kernel 微基准 batch ≥ 16、L 8 时是 1.2×–1.5×batch 1 略慢但端到端在 Qwen3.5-35B-A3B 上 128 并发只有约2.3% 的 TPOT 改善。原因写得直白这个模型 MoE 占比高GDN 线性层在整步成本里只占一小块被 MoE 和全注意力层稀释了。RFC 里的理论天花板是 ~0.57× 流量L 8→ ~1.75×实测 1.2–1.5× 的差距来自带宽利用率不满 100%、每步仍要读检查点、重建算力和环形缓冲流量。结论一句话这是模型相关的优化不是普适加速。社区里的质疑也集中在这几点缓存输入本身要占额外显存收益高度依赖 batch 与模型结构小批量、强延迟场景收益有限与框架、检查点机制、流式推理的集成复杂度与维护成本是采用门槛。也有观点认为这本质上是系统/IO 优化而非新架构——但它把这类优化标准化并推进了主流推理栈这本身有价值。谁该用怎么试适合跑混合模型Mamba-2 / GDN / KDA 层的 Nemotron、Qwen3.5、Kimi-K3、GLM-5.3-Flash 一类、面向吞吐而非极限低延迟、batch 较大的服务或者正在为混合模型上投机解码、被逐草稿状态快照卡住并发的团队以及想给 SSM 状态上低精度量化但担心误差累积的场景。不适合batch 1 的本地单路推理、纯 Transformer 模型、MoE 占比极高导致线性层成本被稀释的模型。上手路径基本是配置项级别。vLLM 的 RFC#47572把它设计成 opt-in 路径默认全关关掉时解码路径与基线逐字节一致缓存作为额外张量挂在既有的 paged mamba state 上复用分页分配和 block table所以持续批处理、按请求分配/驱逐都沿用原有机制。关键 flag--use-replayssm# Mamba2/GDN 标准 cached decode--use-replayssm-spec# 投机 cached decode需先开投机--replayssm-buffer-len# 缓冲长度 L建议先试 16思路等价的最小伪代码Mamba-2 单头、单步# 非 flush 步只追加输入输出从 checkpoint 缓冲算出buf.append((v_t,delta_t,k_t))y_tsum_{jinbuf}a_j*(k_j q_t)*v_j \decay_prod(buf)*(S_ckpt q_t)# output-only不物化状态# flush 步折叠缓冲写回一次完整状态iflen(buf)L:S_ckptfold(S_ckpt,buf)buf.clear()自己验证时建议三层走先把 L 1 与基线逐位对齐此时数学上完全等价再开 L 1 看误差是否落在 TF32/bf16 容差内最后才做端到端准确率与吞吐对比。SGLang RFC 里那个「决策门」值得抄先测目标 batch 下投机 verify 前向的 HBM 带宽利用率——高于约 70%全量状态写摊薄能真换来时间低于约 30%重写只换来复杂度不要做。参考资料Liou et al.,ReplaySSM: Cache SSM Inputs, Not State, arXiv:2610.13778 — https://arxiv.org/abs/2610.13778Dao AI Lab,ReplaySSM算法与评测长文 — https://tridao.me/blog/2026/replayssm/NVIDIA Developer Blog,ReplaySSM: Rethinking State Storage for Faster Hybrid-Model Decoding— https://developer.nvidia.com/blog/replayssm-rethinking-state-storage-for-faster-hybrid-model-decodingReplaySSM 参考实现vLLM 分支 — https://github.com/Johnny-Liou/ReplaySSMvLLM RFC #47572: ReplaySSM cache SSM inputs instead of state — https://github.com/vllm-project/vllm/issues/47572SGLang RFC #28511: Porting ReplaySSM to SGLang — https://github.com/sgl-project/sglang/issues/28511vLLM Blog,Kimi K3 Performance Optimizations in vLLM含 ReplaySSM 生产落地 — https://vllm-project.github.io/2026/09/13/kimi-k3-performance-optimization.html创作说明本文在资料整理与成文中使用了 AI 辅助DeepSeek经作者审阅后发布作者对内容负责。