SevenTnewS

AI推理加速

JetSpec:并行树草稿解码打破推测解码缩放天花板

JetSpec在目标模型的融合隐藏状态上训练一个因果并行草稿头,生成与自回归分解对齐的候选树。它在稠密和MoE Qwen3模型的数学、编码和对话任务中持续优于基线方法。

Emmanuel Fabrice Omgbwa Yasse

2026-07-10 · 阅读需 4 分钟

JetSpec:并行树草稿解码打破推测解码缩放天花板

推测解码已成为加速自回归大型语言模型的关键技术。其思路简单明了:并行草拟多个token,然后一次性验证,而非逐个生成token。但存在一个瓶颈:扩大草稿预算(即每一步草拟的token数量)历史上一直遭遇瓶颈。速度提升仅在接受率保持高位且草稿开销较低时才能实现。基于头的推测解码的两种主流方法各有缺陷,造成了JetSpec研究人员所称的“因果效率困境”。

自回归草稿器生成路径条件候选,在树推测解码中表现良好,实现更高的接受长度。缺点是草稿成本随树深度线性增长,限制了效率。而双向块扩散草稿器则能在单次前向传递中生成所有位置。然而,其分支无关边际分布可能产生个别合理但整体不一致的树,浪费预算并拖低接受率。

JetSpec(于2026年6月16日提交至arXiv的一篇论文中详述)提出一种基于头的推测解码框架,避开了这一权衡。其核心创新是一个因果并行草稿头,在冻结的目标模型的融合隐藏状态上进行训练。与早期方法不同,JetSpec生成的候选树得分与目标模型的自回归分解一致。这种对齐使框架能将更大的草稿预算转化为更长的已接受前缀和更高的端到端加速。

技术架构

JetSpec的工作原理是将轻量级草稿头附加到冻结的预训练大型语言模型上。草稿头从目标模型的多个层获取融合隐藏状态作为输入,并输出候选token树。关键设计选择是跨分支强制执行因果条件:草稿树中的每个token都根据通向它的路径进行条件化,镜像目标模型推理的自回归特性。这确保了树内部一致性,不同于块扩散方法的分支无关输出。

草稿头与一个拒绝不一致分支的验证机制联合训练。训练涉及最小化一个平衡草稿接受长度与计算成本的损失函数。训练完成后,草稿头可直接用于推理,除了隐藏状态融合所需的单次前向传递外,无额外开销。

基准测试结果

研究团队在Qwen3模型的稠密和混合专家(MoE)变体上评估JetSpec,涵盖三个基准:数学(MATH-500)、编码(HumanEval)和开放对话(MT-Bench)。在NVIDIA H100 GPU上,JetSpec在MATH-500上实现高达9.64倍加速,在开放对话工作负载上实现4.58倍速度提升。这些增益在稠密和MoE架构间保持一致,JetSpec在所有配置中均优于双向头和基于树的推测解码基线。

通过集成到流行的开源推理服务库vLLM,进一步展示了延迟优势。在实际服务负载下,与基线推测解码方法相比,JetSpec将尾部延迟降低高达60%,这清楚表明其在生产部署中的实际可行性。

对大型语言模型部署的影响

在不牺牲质量的前提下加速推理的能力,对于大型语言模型在实时应用(如聊天机器人、代码助手和交互式教学系统)中的落地至关重要。JetSpec解决了一个根本性瓶颈:草稿质量与计算效率之间的权衡。通过使草稿过程既快速又因果一致,它释放了更大草稿预算的潜力,从而直接转化为更快的生成速度。

该框架基于头的方法意味着它可以应用于现有大型语言模型而无需修改核心模型架构。这是一个实际优势:组织可以在不重新训练或替换已部署模型的情况下升级推理速度。

代码与可用性

JetSpec代码和模型在宽松许可下开源,地址为https://github.com/hao-ai-lab/JetSpec。项目页面带有额外的可视化和文档,地址为https://jetspec-project.github.io/jetspec-web/。截至论文提交时,该项目在Hugging Face上获得27个赞。

“我们的代码和模型可在该https URL获取,”论文作者注释道,强调他们对可重复性和社区采纳的承诺。

结论

JetSpec标志着推测解码领域迈出了有意义的一步,打破了先前方法面临的缩放天花板。通过因果并行草稿头解决因果效率困境,它在多样化的基准和模型架构上实现了显著的加速。随着对更快大型语言模型推理的需求持续增长,像JetSpec这样的框架将对于大规模部署先进人工智能系统至关重要。