投机解码让一个便宜的草稿模型先猜出接下来几个token,大模型再一次性校验。猜错的直接丢掉。这套机制在纯注意力模型上一直很划算,但放到混合模型里,情况会变。

打开网易新闻 查看精彩图片

混合模型混了两类层。注意力层维护的缓存只增不减,丢掉一个猜测只需要把指针往回挪。Mamba层不一样,它每个请求只保留一份固定大小的状态,每生成一个token就覆盖一次。旧状态没了,想撤销一个被拒绝的猜测,服务器就得在每生成一个草稿token后保存一份完整状态副本。四个草稿,就是每层、每请求、每步四份副本。

以Nemotron-3-Super为例,它有40个Mamba层和8个注意力层,每个请求的Mamba状态是4 MiB。这个量级相当于一个注意力层存约8000个token的缓存。问题不在容量,在流量。

副本为什么这么贵

生成一个token的速度受限于GPU读写内存的速度。权重每步读一次,由整个批次共享;状态是每个请求各一份。批次越大,状态流量越大,很快就会压过权重。每一份副本,都等于多写一遍完整状态。

NVIDIA实测了vLLM在这些模型上的投机解码:批次超过128后就不再带来增益;在Qwen3.5-122B-A10B上,批次到512时反而低于普通解码。校验8个草稿token的耗时接近普通一步的五倍。

把状态换成日志

Mamba-2层的状态更新只有一条规则:先衰减旧状态,再加上新token的贡献。这条规则可以展开。今天的输出等于检查点衰减后的结果,加上日志里每个token的贡献、按时间远近各自衰减。于是这一层可以直接从检查点和日志算出输出,中间不必写状态。

状态只在日志写满时落一次盘。NVIDIA和Tri Dao把这套做法叫ReplaySSM。窗口取16个token,完整状态就从每步写一次变成每16步写一次。被拒绝的草稿永远不进日志。数学上没有变化,只是换了排列顺序。

什么时候会踩到

自己部署混合模型、跑投机解码、批次又比较大的场景。Nemotron 3、Qwen3.5、Qwen3-Next都混了Mamba式层与注意力层。这是一道内存流量墙,比算力闲置那道墙来得更早,因为批次里每个请求都要为自己的副本买单。

判断规则可以直接算:一旦状态流量占主导,带副本的一步所搬运的状态约为普通一步的(1 + T) / 2倍,T是草稿数。只有每步保留的token数超过这个值,投机才划算。按接受率0.6、4个草稿算,系数是0.92,大批次下会输;接受率0.8时系数1.34,仍然赢。接受率要读自己流量的服务指标,不能照搬基准测试——代码补全保留的草稿比开放式对话多,同一个模型可能落在线的两侧。

落地要动三处

  • 容量规划:按模型层数、头数、头维度、状态大小和字节数算出每请求状态量。Nemotron-3-Super是40层各4 MiB,约160 MiB每请求,每步都要读一遍。
  • 服务配置:上游vLLM为Mamba-2模型提供了两个默认关闭的开关,--use-replayssm--replayssm-buffer-len。窗口16是vLLM自己扫出来的默认值,Tri Dao的测试里Nemotron 3用8最好,8、16、32都值得试。状态保持FP32;低精度副本只适合读,一旦让共享的低精度状态喂给递推,Nemotron-3-Nano-4B在GPQA-Diamond上掉了31分。Qwen3.5用的是Gated DeltaNet层,相关内核在写作时还只在作者的vLLM分支里,不在上游,升级前先看发布说明。
  • 基准测试:固定提示词、输出长度和GPU,逐档扫批次,每个点跑两次,因为首次运行要预热缓存和编译内核。对比普通解码、投机、replay、replay加投机四种组合。

NVIDIA测得批次256时,仅replay就能让每个解码步快1.20到1.47倍。加上replay后,投机应该能越过批次128继续增益,而不是走平。最后用评测集复核,分数应在噪声范围内一致,因为数学是精确的。

什么时候不该用

批次很小时,权重主导每一步,状态副本只是零头,NVIDIA也没看到replay带来加速。纯注意力模型不需要,它的缓存本来就能用指针撤销。草稿很少被接受时,replay只能让撤销变便宜,救不了糟糕的猜测。换一种递推结构也要先确认更新能否展开成衰减加求和,Mamba-2和Gated DeltaNet可以,新架构不能想当然。至于把状态压到4比特,那是另一个决定。