推理加速方法 JetSpec 用并行草稿树,在数学评测中实现 9.64 倍加速 | AGIDaily
AGI Daily
资讯日报词条论文人物新创

推理加速方法 JetSpec 用并行草稿树,在数学评测中实现 9.64 倍加速

2026-06-29 / 共 4,191 字

Pexels

草稿预算增加后,加速收益为何会停止增长

推测解码(speculative decoding, SD)可以加快 LLM 推理,但让草稿模型每轮提出更多候选 token,并不总能继续提高速度。

UCSD Hao AI Lab 在 2026 年 6 月发表的 JetSpec 论文中,将这种随草稿预算增加而出现的收益限制称为 scaling ceiling(扩展上限)。除硬件条件外,草稿算法本身也会形成瓶颈。

它的含义非常具体:在推测解码框架中,增大草稿预算(draft budget)——也就是让草稿头每轮多猜几个 token——只有在这两个条件同时满足时才有效:草稿接受率(acceptance rate)足够高,且草稿生成成本足够低。一旦这两个条件中任何一个开始劣化,增大预算不但不会加速,反而可能成为拖累。

Medusa、EAGLE、DFlash 等方法尝试用不同草稿机制改善这一问题。在本文所讨论的评测配置下,部分既有方法在数学推理任务上的加速约为 8 倍多,在开放式对话上约为 3–4 倍;这些是具体实验结果,并非所有模型通用的速度上限。

JetSpec 在 Qwen3-8B 的一组评测中报告了:MATH-500 上 9.64 倍端到端加速(对比 DDTree 的 8.78 倍),GSM8K 上 7.82 倍,HumanEval 上 7.12 倍,开放对话 MT-Bench 上 4.58 倍。实现这些结果的方法是——因果并行树起草(causal parallel tree drafting)。

草稿生成越慢、接受率越低,验证收益就越小

要理解 JetSpec 的贡献,需要先回到推测解码的基本逻辑。

现代自回归 LLM 的推理瓶颈在于"串行":每生成一个 token,都需要跑一次完整的前向传播,后一个 token 依赖前一个 token 的结果,不能像已知序列那样直接并行生成。推测解码的解决方案是引入一个"草稿模型"(draft model)——用较小的代价一次性猜多个后续 token,然后让目标模型在一次前向传播中并行验证这些猜测。验证通过的 token 直接采纳,第一个被拒绝的位置由目标模型自己补上。

这个方法之所以"无损"(lossless),是因为验证阶段保证了输出分布严格等于目标模型的自回归分布。

但草稿模型的设计存在一个根本性的困境,Hao AI Lab 称之为 causality–efficiency dilemma(因果性-效率二难):

  • 自回归草稿头(如 EAGLE 系列):保持了目标模型的因果条件依赖——每个草稿 token 都基于前缀和同一路径上已生成的 token——因此草稿质量高、接受率好。但代价是草稿过程本身也是串行的:树的每一层深度都需要一次单独的前向传播,草稿成本随树深线性增长。

  • 块扩散草稿头(如 DFlash):可以在一次前向传播中并行生成整个 token 块,草稿成本极低。但问题在于块内每个位置是独立评分的,不同分支之间缺乏因果条件依赖——就像让九个互不沟通的人分别续写同一段话的后九句,每个人只能看到原始前缀,不知道其他人写了什么。结果是一个个单独看似乎合理、但放在一起就互相矛盾的 token 树。

DDTree 尝试在 DFlash 的基础上引入最佳优先树扩展,一定程度上改善了树结构,但并未从根本上解决"扩散草稿头缺乏因果条件依赖"的问题。论文数据显示,DDTree 确实优于原始 DFlash,但在 MATH-500 上仍然只有 8.78 倍加速,在继续增大预算时,加速收益也会受到限制。

论文用一个简洁的公式描述了推测解码加速的核心:

加速比 = 平均接受长度 ÷ 草稿成本系数

其中草稿成本系数 = 1 +(每轮草稿生成总耗时 ÷ 目标模型单步耗时)。这是简化的成本模型;实际服务还需计入验证、调度和缓存管理等开销。它说明:增大预算通常能提升接受长度,但如果草稿头的每次前向传播变得太贵(自回归头的问题),或者接受长度增长赶不上预算增长(扩散头的问题),加速比就会停滞甚至倒退。

一次并行起草,保留候选分支的前缀依赖

JetSpec 的核心直觉可以概括为一句话:能不能既保持一次并行前向传播的效率,又恢复因果条件依赖?

答案是因果并行树起草(causal parallel tree drafting)。JetSpec 训练了一个轻量级因果并行草稿头,挂载在冻结的目标 LLM 之上。这个草稿头有三个关键设计:

第一,复用目标模型的多层隐藏特征。 草稿头不是从零开始计算,而是读取目标模型中间多层的隐藏特征(fused hidden states)。复用这些特征可以减少重复计算,控制草稿生成的开销;具体耗时仍取决于草稿预算和硬件配置。

第二,用注意力掩码限制每个节点可读取的上下文。 草稿头在并行生成多个树节点时,不是让它们独立计算,而是施加一个树因果注意力掩码(tree-causal attention mask):每个节点可以读取原始前缀和自身路径上的祖先节点,但不能读取后继节点或无关的兄弟分支。这样,各个分支的预测能保留各自的上下文依赖。

第三,块级监督与 forward-KL 蒸馏。 训练时,JetSpec 将草稿序列划分为多个块,每个块的第一个 token 作为锚点(anchor),剩余位置替换为掩码 token 嵌入。草稿头并行预测所有被掩码的未来 token,每个位置只关注前缀和同一块内的早期位置。论文的消融实验表明,forward-KL 风格的蒸馏略优于 SFT,因为它保留了多个可能延续上的概率质量,而 mode-seeking 的 reverse-KL 不太适合预算受限的树起草场景。

最终效果是:JetSpec 在一次草稿头前向传播中,生成一整棵内部因果一致的候选树——树的每一层有多个分支,每个分支内部保持因果条件依赖,不同分支共享前缀——然后目标模型一次验证整个树,采纳最长的匹配前缀。

论文还报告了平均接受长度:在 MATH-500 上以 256 预算运行 greedy 解码时,JetSpec 平均每个轮次接受 τ = 10.76 个 token——这意味着每跑一次目标模型前向传播(加上一次便宜的草稿头传播),就能推进超过 10 个 token。

H100 上的加速随任务变化,数学与代码任务收益较高

JetSpec 在 Qwen3-8B 上使用 NVIDIA H100 GPU、256 草稿预算、greedy 解码的完整评测数据如下:

基准测试 JetSpec DDTree 提升幅度
MATH-500 9.64× 8.78× +9.8%
AIME25 8.78× — —
GSM8K 7.82× 7.04× +11.1%
HumanEval 7.12× 6.31× +12.8%
LCB 7.67× — —
MBPP 6.73× — —
MT-Bench(开放对话) 4.58× — —

几个值得注意的观察:

一是 JetSpec 在推理密集型任务上优势最大。MATH-500 的 9.64 倍加速位居所有基准之首,AIME25(8.78 倍)、GSM8K(7.82 倍)紧随其后。这与 JetSpec 的设计直觉一致:数学推理的 token 分布更可预测,草稿接受率更高,因果条件依赖带来的质量提升也更显著。

二是开放对话上的加速为 4.58 倍,低于数学和代码任务。在这些实验中,对话候选更难获得较长的接受前缀,说明收益会随任务分布而变化。

三是 JetSpec 在 MoE 架构上也保持优势。论文在 Qwen3-30B-A3B(MoE)目标模型上使用相同训练配方评估,JetSpec 持续优于 DDTree,说明因果并行树起草不局限于某一种模型架构。

团队提供 vLLM 分支,B200 单请求测试约 1000 TPS

团队还测试了包含调度与缓存管理开销的服务实现。

Hao AI Lab 提供了支持 JetSpec 的 vLLM 分支,并用它进行服务测试。在单张 NVIDIA B200 GPU 上运行 Qwen3-8B 加上 JetSpec 草稿头,采用 batch 1、草稿预算 127,团队在 MATH-500 上实现了平均约 1000 tokens per second(TPS)。这一结果对应单请求场景,不能直接等同于多请求并发吞吐量。原报道列举的标准自回归解码约 100–150 TPS,也需要在相同服务配置下比较。

团队还评估了不同请求速率(request rate)下的表现:在中小负载场景下始终优于 baseline,且在更高端 GPU(如 B200)上增益尤其显著。但更大的草稿树也增加了验证和内存压力,高并发批处理下的收益可能下降;预算需要结合负载调整。

GitHub 仓库已公开发布了 JetSpec 草稿模型权重,覆盖 Qwen3-8B、Qwen3-30B-A3B、Qwen3.6-35B-A3B、gpt-oss-20b、Gemma 4 26B A4B IT 和 Step 3.7 Flash。仓库提供 benchmark 脚本,用于对比 DFlash baseline 和 JetSpec 的多种树解码算法。

无损验证保留输出分布,部署收益仍需同条件实测

推测解码是 LLM 推理加速的核心技术路线之一。与量化、剪枝、蒸馏等"有损"手段不同,推测解码承诺无损加速——输出分布严格等价于原始模型。这使数学推理、代码生成、金融分析等场景可以在不主动改变目标模型输出分布的前提下尝试加速,但并不保证目标模型本身不会出错。

但也正是这个承诺,使得推测解码的加速上限天然受制于草稿质量与草稿成本之间的 trade-off。此前的研究分别尝试了 Medusa 的多头预测、EAGLE 的特征复用和 DFlash 的块并行,在草稿成本与接受长度之间寻找更好的平衡。

JetSpec 通过调整草稿过程的计算图,在并行生成候选的同时保留分支的因果依赖。在所测配置下,这让较大的草稿预算能继续换来更长的接受前缀,但仍需付出验证成本。

后续可以继续研究更强硬件和更大草稿预算的组合。B200 上约 1000 TPS 的结果展示了这一方向的潜力,但硬件升级能带来多少额外收益,还需在不同任务和并发条件下验证。

工程团队可以从项目提供的 vLLM 分支、配套草稿权重和脚本开始测试,并检查与现有模型及服务版本的兼容性。MATH-500 上的 9.64 倍加速有明确实验条件,迁移到其他 GPU、模型或负载时,需要重新测量。

JetSpec 的实验说明,改进草稿结构可以提高特定配置下的加速收益;推测解码仍有优化空间,收益上限则由算法与服务条件共同决定。


论文:JetSpec: Breaking the Scaling Ceiling of Speculative Decoding with Parallel Tree Drafting (arXiv:2606.18394)

代码与模型:https://github.com/hao-ai-lab/JetSpec

来源

  • JetSpec: Breaking the Scaling Ceiling of Speculative Decoding with Parallel Tree Drafting · arXiv
  • JetSpec: Breaking the Scaling Ceiling of Speculative Decoding with Parallel Tree Drafting | Hao AI Lab @ UCSD · Hao AI Lab
  • JetSpec GitHub Repository · Hao AI Lab
  • A speculative decoding method called 'JetSpec' has been developed that can speed up AI by up to 9.64 times · GIGAZINE

评论

登录后即可参与评论。

加载评论中…

相关文章

  • 扩散语言模型为何能加快文本生成?从 LLaDA 到 DiffusionGemma

    扩散语言模型为何能加快文本生成?从 LLaDA 到 DiffusionGemma

    2026-06-10

  • 腾讯混元 Hy4 preview 模型卡中的基准测试结果图表

    腾讯开源混元 Hy4:支持百万 token 上下文,并参与优化推理系统

    2026-08-31

  • 智谱官方示意图,对比 Infra Agent 只接收端到端结果的稀疏反馈,与由正确性、系统行为和性能验证组成的密集反馈回路

    智谱用密集反馈让 GLM 持续优化推理栈

    2026-09-17

  • 猎豹奔跑的抽象视觉图,MIT Technology Review 用于 Subquadratic 稀疏注意力报道
AGI Daily跟踪世界 AI 进展
邮件栏目主题RSS
定制投稿招聘广告
关于MCPTermsContact

Subquadratic 用稀疏注意力降低长文本计算量,公开验证仍有限

2026-06-21

  • 屏幕上的代码细节,象征 token 级推理控制与算法优化

    让AI选择快答或长推理,ADaPT将70亿参数模型的平均输出缩短约三分之一

    2026-07-02

  • StreamKL 概览:传统实现保存完整注意力矩阵;StreamKL 按块在线计算 KL 散度,避免在 HBM 中保存这些矩阵

    StreamKL 融合注意力蒸馏计算,前向算子最高加速 43 倍

    2026-06-24