论文导读
加州大学圣地亚哥分校、南京大学、Meta、阶跃星辰等机构提出JetSpec++,在大模型核验当前文字时提前生成下一段草稿。论文的GB200双GPU配置下,单请求峰值达到每秒1361个Token,相对自回归解码最高加速7.45倍,同时由目标模型核验保留原输出分布。
草稿写得快,等验稿仍会浪费时间
推测解码的办法很直观:让较便宜的模型先拟一段文字,再让目标大模型一次检查多个Token。猜对的连续部分保留下来,猜错的位置由目标模型纠正,就能减少逐字运行大模型的次数。
但草稿和核验如果严格排队,目标模型验完以后还要等下一份草稿。JetSpec++把这两件事重叠:这一段尚未验完,下一段已经开始写。论文先用时间线对照串行与流水线,变化是等待关系,而不是换一个更小模型冒充原模型。
图:常规先草拟后核验,与JetSpec++重叠草拟和核验的时间关系示意。
提前写也有代价。当前草稿最终会接受几个Token、下一段从哪里接起,都尚未确定。新方法必须预想若干可能的验证结果,并为这些结果准备下一段,而不是默认整份草稿都会通过。
预想接续位置,再用空档把草稿改一遍
JetSpec++预测目标模型尚未完成的内部状态,以及可能的接续Token,再据此准备不同后续草稿。真实核验返回后,如果结果命中缓存,就直接拿出对应草稿;没有命中,则回到常规推测步骤。
在目标模型核验的时间窗口里,草稿模型还能迭代修正自己的并行预测。前一轮预测为后一轮提供上下文,逐步减少块内各位置互不知情带来的错误。多做几轮是否划算,取决于这些计算能否藏进核验窗口。
图:预测内部状态、准备不同接续结果与缓存命中后取出草稿的流程。
最终被提交的Token仍经过目标模型的接受规则,论文因此强调输出分布保持不变。这个性质属于解码算法,不代表草稿永远正确,也不能把省去关键路径等待写成不消耗草稿计算资源。
1361是峰值,比较时要带上硬件和任务
作者在七个基准上比较接受长度,也就是每次核验最终提交多少Token。温度0、核验预算16时,JetSpec++最好配置的平均接受长度为8.94,DSpark为6.89。接受得更多,才有机会摊薄每次大模型核验的成本。
服务实验另行测量实际速度,使用GB200、单请求和CUDA图,将草稿模型放在第二张GPU。论文报告最高每秒1361个Token,相对自回归解码最高7.45倍、相对DSpark最高1.54倍。峰值随目标模型和任务变化,不能当成七项基准平均提速。
图:GB200上不同目标模型、任务和方法的吞吐、GPU数量及每轮时间。
这也不是“只改代码就免费得到七倍性能”:额外GPU、草稿模型训练、通信和主机调度都属于运行条件。作者说明,融合计算与通信优化仍有改进空间,计算赶不上核验窗口时,流水线能隐藏的成本也会减少。
论文还对重复输出做了截断处理,避免模型陷入重复时草稿过于好猜,虚高接受长度与速度。对推测解码而言,Token吞吐之外,是否保住原模型输出、如何处理容易重复的样本,同样决定结果能否比较。
参考资料
https://arxiv.org/abs/2610.19792
https://arxiv.org/html/2610.19792