JetFlow：通过并行树草稿打破推测解码的缩放天花板

# JetFlow：利用并行树草稿打破推测解码的扩展天花板
来源：https://arxiv.org/html/2606.18394
Lanxiang Hu¹，Zhaoxiang Feng¹，Yulun Wu²，Haoran Yuan³，Yujie Zhao¹，Yu-Yang Qian⁴，Bojun Wang⁵，Daxin Jiang⁵，Yibo Zhu⁵，Tajana Rosing¹，Hao Zhang¹
¹UC San Diego，²浙江大学，³UIUC，⁴南京大学，⁵StepFun

###### 摘要
推测解码（Speculative Decoding, SD）通过草拟多个token并并行验证来加速自回归大语言模型（LLMs），但它面临一个扩展限制：增加草稿预算只能在接受率保持较高且草稿开销较低时提升速度。这一天花板难以突破，因为先前的基于头的SD方法面临因果-效率困境。自回归草稿生成器产生路径条件候选，这对于树型推测解码具有更高的接受长度是有效的，但其草稿成本随树深度增加。双向块扩散草稿生成器一次生成所有位置，但其与分支无关的边缘分布可能形成单独合理但相互不一致的树，浪费预算并降低接受率。我们提出JetFlow，一个基于头的SD框架，将单次前向草稿效率与分支级因果条件相结合。JetFlow在冻结目标模型的融合隐藏状态上训练一个因果并行草稿头，生成得分与目标模型自回归分解对齐的候选树。这使得JetFlow能够将更大的草稿预算转化为更长的接受前缀和更高的端到端加速。在密集型和MoE Qwen3模型的数学、编码和聊天基准测试中，JetFlow一致优于双向头和基于树的SD基线。在H100 GPU上，JetFlow在MATH-500上实现高达9.64×的加速，在开放式对话工作负载上实现4.58×的加速，并通过在真实服务负载下集成vLLM进一步展示了延迟收益。我们的代码和模型可在https://github.com/hao-ai-lab/JetFlow获取。

## 1 引言

现代自回归（AR）大语言模型（LLMs）在一个简单但代价高昂的瓶颈下运行：解码在很大程度上是串行的。这种顺序生成过程使得延迟成为在数学（Wei et al., 2022 (https://arxiv.org/html/2606.18394#bib.bib132)；Zhu et al., 2025 (https://arxiv.org/html/2606.18394#bib.bib133)）、编码（Deng and others, 2025 (https://arxiv.org/html/2606.18394#bib.bib128)；Jain et al., 2025 (https://arxiv.org/html/2606.18394#bib.bib139)）以及其他智能体推理任务（Patil et al., 2025 (https://arxiv.org/html/2606.18394#bib.bib136)；Yao et al., 2025 (https://arxiv.org/html/2606.18394#bib.bib137)；Lu et al., 2025 (https://arxiv.org/html/2606.18394#bib.bib138)）等应用中部署时的一个主要挑战，这些任务中模型经常生成长序列。现有方法通过两种方式应对这一挑战：一种是在预训练或中期训练中引入多token预测（MTP）（Gloeckle et al., 2024 (https://arxiv.org/html/2606.18394#bib.bib140)；DeepSeek-AI, 2024 (https://arxiv.org/html/2606.18394#bib.bib149), 2026 (https://arxiv.org/html/2606.18394#bib.bib150)；Huang et al., 2026 (https://arxiv.org/html/2606.18394#bib.bib151)；Xiao et al., 2026 (https://arxiv.org/html/2606.18394#bib.bib152)），另一种是在后训练中使用适应方法将预训练LLMs转化为扩散LLMs（dLLMs），但会带来不同程度的性能下降（Zhao et al., 2025 (https://arxiv.org/html/2606.18394#bib.bib147)；Wu et al., 2025 (https://arxiv.org/html/2606.18394#bib.bib148)）。相比之下，在推理时应用推测解码（SD）可以在保持质量的同时减少生成延迟（Leviathan et al., 2023 (https://arxiv.org/html/2606.18394#bib.bib141)；Chen et al., 2023 (https://arxiv.org/html/2606.18394#bib.bib156)）。在这些技术中，SD因其低适应成本和简洁性而脱颖而出。一方面，先前工作表明，通过轻量级草稿模型对齐（通常少于1B训练token）可以大幅提升SD性能（Liu et al., 2024 (https://arxiv.org/html/2606.18394#bib.bib164)；Zhou et al., 2024 (https://arxiv.org/html/2606.18394#bib.bib165)；Goel et al., 2024 (https://arxiv.org/html/2606.18394#bib.bib166)）。另一方面，基于头的方法（如Medusa和EAGLE）通过避免传统的独立草稿-目标模型设置来降低部署复杂性（Cai et al., 2024 (https://arxiv.org/html/2606.18394#bib.bib142)；Li et al., 2024b (https://arxiv.org/html/2606.18394#bib.bib143)）。更近期的基于头的方法，如EAGLE-3（Li et al., 2025 (https://arxiv.org/html/2606.18394#bib.bib144)）和DFlash（Chen et al., 2026 (https://arxiv.org/html/2606.18394#bib.bib145)），继续改进草稿质量并降低草稿成本以实现更好的端到端加速。尽管有这些进展，基于头的SD仍然面临因果-效率困境。自回归草稿生成器（如EAGLE）生成高质量的路径条件候选，但随着树深度增长需要顺序的草稿传递。双向块扩散草稿生成器（如DFlash）一次生成草稿，但其与分支无关的边缘分布可能形成单独合理但相互不一致的树，浪费预算并降低接受率。我们提出JetFlow，一种旨在打破这种因果-效率困境的SD范式：JetFlow训练一个因果并行草稿头，在一次前向传递中预测多个树节点，同时通过对隐藏状态的块级因果注意力保持分支级因果条件。这为每个分支提供了与目标模型自回归分解对齐的路径条件草稿分布，使树接受率能够更好地随额外草稿预算扩展。实验表明，JetFlow展现出显著的效率提升。在16个草稿token的低token预算情况下，JetFlow实现了有竞争力的端到端加速。在256个草稿token的高token预算情况下，JetFlow取得了强劲的增益，在MATH-500上达到τ=10.7和9.5×的加速，在MT-Bench上仍有效，达到τ=5.9和超过4×的加速。我们进一步将JetFlow集成到工业级服务引擎中，并在不同请求率下进行评估，在中小服务负载下一致优于基线，在高端GPU（如B200）上尤其表现出色。这些结果凸显了JetFlow在实际服务中减少解码延迟的潜力。

请参阅图1
图1：在H100 GPU上，跨数学、编码和聊天基准测试，相对于标准自回归解码的端到端解码加速。DFlash表示原始块并行草稿方法，DDTree是DFlash基于树的变体，JetFlow表示我们的方法。两者均使用算法1（https://arxiv.org/html/2606.18394#alg1）采用256个token的树预算。

简而言之，我们的工作做出了以下贡献：
- 我们提出了JetFlow，一种旨在联合优化推测解码的草稿成本和接受率的范式。JetFlow训练一个具有并行树草稿能力的因果预测头，在保持草稿token间路径条件依赖的同时实现低成本草稿生成。
- 我们开发并评估了树草稿算法，使推测解码能够更好地利用额外的解码计算。通过扩大草稿token预算，JetFlow实现了高达τ=10.7和超过9.5×的端到端加速。
- 我们将JetFlow集成到工业级服务引擎中，并在真实服务场景下进行评估。JetFlow在中小服务负载下一致优于基线，突出了其在实际部署中减少延迟的潜力。

## 2 JetFlow

### 2.1 背景

##### 推测解码中的扩展瓶颈。
在推测解码（SD）中，轻量级近似模型\(M_q\)首先提出\(N\)个草稿token，然后由目标模型\(M_p\)并行验证（Leviathan et al., 2023 (https://arxiv.org/html/2606.18394#bib.bib141)）。目标模型接受与其自身预测一致的最长前缀，并从第一个