💻 computer science

JetFlow: Breaking the Scaling Ceiling of Speculative Decoding with Parallel Tree Drafting

JetFlow 是一种新颖的投机解码框架，它通过训练一个因果并行草稿头来生成分支一致的标记树，从而克服了现有方法的扩展限制，在不降低接受率的情况下，在多种大语言模型工作负载上实现了显著的加速（高达 9.64 倍）。

原作者： Lanxiang Hu, Zhaoxiang Feng, Yulun Wu, Haoran Yuan, Yujie Zhao, Yu-Yang Qian, Bojun Wang, Daxin Jiang, Yibo Zhu, Tajana Rosing, Hao Zhang

发布于 2026-06-19

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Lanxiang Hu, Zhaoxiang Feng, Yulun Wu, Haoran Yuan, Yujie Zhao, Yu-Yang Qian, Bojun Wang, Daxin Jiang, Yibo Zhu, Tajana Rosing, Hao Zhang

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图写一个长篇故事，但你有一个非常严格的编辑（AI 模型），他只允许你一次写一个词。在你写第二个词之前，编辑必须阅读并检查第一个词，然后给你绿灯。在写第三个词之前，他要检查第二个，以此类推。这种“一次一个词”的规则使得写作变得非常缓慢，即使编辑极其聪明。

**投机采样（Speculative Decoding）**是一种加速这种过程的技巧。与其等待编辑检查每一个词，不如雇佣一个快速、廉价的助手（“草拟者”）来猜测接下来的几个词。然后你将这些猜测一次性展示给编辑。如果编辑同意这些猜测，你就能在写一个词的时间内写下好几个词。如果编辑不同意，你就丢弃错误的猜测并重新开始。

问题：“速度 vs. 准确度”陷阱

论文解释说，以往的方法遇到了瓶颈。它们面临着一个两难困境：

“谨慎型”助手： 有些助手非常小心。他们猜下一个词，然后根据他们的第一个猜测来猜下一个词，以此类推。这使得他们的猜测非常准确（编辑经常接受它们），但速度很慢，因为他们必须循序渐进地思考。
“快速型”助手： 另一些助手速度极快。他们不考虑词与词之间的联系，直接大声喊出一整串单词列表。这很快，但列表往往逻辑不通（例如，“猫……飞了……到……月球……昨天”）。由于这些词无法符合故事的流畅度，编辑不得不拒绝大部分内容，从而浪费了助手的速度。

论文称之为**“因果性-效率困境”（Causality-Efficiency Dilemma）**。你通常必须在“快速但不准确”和“准确但缓慢”之间做出选择。

解决方案：JETFLOW

作者创建了一个名为 JETFLOW 的新系统，它打破了这个陷阱。把 JETFLOW 想象成一个能够并行思考，但仍能遵循故事逻辑的超级助手。

它是如何工作的，使用一个简单的类比：

树的比喻： 想象故事是一棵树。树干是你已经写好的文本。你想长出新的分枝（未来的词）。
- 旧的“快速型”助手会随机生长分枝。一个分枝可能说“那只猫”，另一个说“那只狗”，第三个说“月球”。它们不知道哪条路径是真实的，所以浪费时间生长死掉的分枝。
- 旧的“谨慎型”助手会长出一个分枝，检查一下，然后再长下一个。这很安全，但很慢。
- JETFLOW 会观察树干，并瞬间勾勒出许多不同的可能分枝。但神奇之处在于：它确保每一条分枝都遵循故事的规则。如果一个分枝以“那只猫”开始，那么该特定分枝上的下一个词必须是猫会做的事情。它不会把不同的分枝混淆。

JETFLOW 如何实现

一次通过，多条路径： JETFLOW 使用一个特殊的“头”（AI 的一部分）来观察主编辑的隐藏想法。只需一眼，它就能预测出一整棵由可能的下一个词组成的树。
尊重流向： 它使用一种特殊的“掩码”（就像一套交通规则），强制要求助手只能观察其特定路径上之前的词。这防止了助手窥探未来或混淆不同的故事线。
结果： 因为助手的猜测在逻辑上与故事流保持一致，主编辑（目标模型）可以一次性接受更长的字符串。

结果：有多快？

论文在强大的计算芯片（H100 GPU）上，针对数学问题、编程任务和聊天对话进行了测试。

数学问题： 在困难的数学测试中，JETFLOW 比标准的缓慢方法快了 9.6 倍。
聊天： 对于开放式对话，它快了 4.5 倍。
可扩展性： 允许助手进行更多“猜测”（预算）时，速度就越快。不像旧方法在被要求猜测太多词时会变得混乱或变慢，JETFLOW 能够保持高效且越来越快。

总结

JETFLOW 就像雇佣了一支助手团队，他们可以同时喊出不同的故事结局，但他们足够聪明，知道每个结局本身必须是合理的。这使得主编辑能够瞬间批准大块的文本，在不损失故事质量的前提下，打破了“一次一个词”的速度限制。

技术摘要：JETFLOW

1. 问题陈述

投机解码（Speculative Decoding, SD）通过并行草拟多个 Token 并对其进行验证，从而加速自回归大语言模型（LLMs）。然而，投机解码面临一个根本性的扩展瓶颈（scaling ceiling）：只有当**接受率（acceptance rate）保持高水平且草拟开销（drafting overhead）**保持低水平时，增加草拟预算（即草拟 Token 的数量）才能提升端到端的速度。

现有方法由于面临因果-效率困境（causality-efficiency dilemma），难以同时优化这两个因素：

自回归草拟器（例如 EAGLE）： 生成高质量、路径条件化的候选 Token，适用于树状投机解码，从而实现较高的接受长度。然而，其草拟成本随树深度的增加而线性增长，限制了可扩展性。
双向块扩散草拟器（例如 DFlash）： 在单次前向传递中生成所有草拟位置，实现了极低的单 Token 草拟成本。然而，它们生成的是与分支无关的边际分布（branch-agnostic marginals）。这些边际分布形成的树虽然单个 Token 可能合理，但在不同分支间却存在相互矛盾之处。这种不一致性浪费了草拟预算并降低了有效接受率，阻碍了系统随预算增加而有效扩展。

2. 方法论：JETFLOW

JETFLOW 是一种基于 Head 的投机解码框架，旨在通过结合单次前向传递的草拟效率与分支级因果条件化来打破因果-效率困境。

核心架构

因果并行草拟头（Causal Parallel Draft Head）： JETFLOW 在冻结的目标模型之上训练了一个轻量级的草拟头。与标准的块扩散头独立预测位置不同，JETFLOW 的 Head 在单次前向传递中预测多个树节点，同时保留因果依赖关系。
树因果注意力掩码（Tree-Causal Attention Mask）： 为了确保因果一致性，草拟头采用了专门的注意力掩码。树中的每个节点可以关注原始输入前缀及其在树内的祖先节点，但不能关注其后代或无关的兄弟分支。
- 这强制执行了分支分解： $q(\pi(v) | x) = \prod_{u \in \pi(v)} q(y_u | x, h^o_x, \pi_{<u})$ 。
- 这使草拟分布与目标模型的自回归分解保持一致，确保候选分支是基于其路径上的特定 Token 进行条件的。
特征融合（Feature Fusion）： 草拟头复用来自冻结目标模型的融合隐藏状态，将目标上下文特征注入草拟层的 KV Cache 中，以引导生成，而无需承担单独草拟模型的成本。

训练策略

数据： 草拟头在与目标模型对齐的序列上进行训练，这些序列可以来自训练语料库，或者更有效地来自由目标模型自身重新生成的序列。
损失函数： JETFLOW 使用**前向 KL 散度（Forward KL Divergence）**进行蒸馏。作者发现，前向 KL 的表现优于硬标签监督微调（SFT）和反向 KL（Reverse KL）。前向 KL 保留了目标模型在多个可能的延续路径上的软标签偏好，这对于树状扩展至关重要；而反向 KL 则过于追求模式收敛（mode-seeking），导致概率过于集中。
树构建： 在推理期间，草拟头并行生成所有树深度的 Logits。分支评分函数（默认为累积草拟对数概率）对节点进行排序，并通过一种最佳优先扩展算法构建候选树，直至达到特定预算（例如 256 个 Token）。

3. 核心贡献

JETFLOW 框架： 一种全新的范式，通过训练因果并行草拟头，同时优化草拟成本和接受率。它实现了低成本、单次前向传递的树生成，同时保持了高接受率所需的路径条件依赖。
可扩展的树状草拟算法： 本文开发并评估了允许投机解码有效利用更大草拟 Token 预算的算法。通过维持因果一致性，JETFLOW 将增加的计算量转化为更长的接受前缀，而非浪费在不一致的分支上。
工业级集成： JETFLOW 已集成至 vLLM 服务引擎中。该实现包含了定制算子（例如带有树掩码分阶段处理的 SM90 Paged FlashAttention），以在真实的推理负载下高效处理树验证。

4. 实验结果

作者在 Qwen3-8B（稠密模型）和 Qwen3-30B-A3B（MoE 模型）上，针对数学（GSM8K, MATH-500, AIME25）、代码（HumanEval, MBPP, LiveCodeBench）和对话（MT-Bench）基准测试进行了评估。

低预算场景（16 个 Token）： JETFLOW 实现了具有竞争力的加速比（例如在 GSM8K 上约 4.8×），与 DFlash 相当，证明了短线性草拟足以提供高概率的延续。
高预算场景（256 个 Token）： 随着预算增加，JETFLOW 的表现显著优于基准模型。
- 在 MATH-500 上，JETFLOW 在 H100 GPU 上实现了 9.64× 的标准自回归解码加速（ $\tau \approx 10.7$ 个接受 Token）。
- 在 开放式对话（MT-Bench） 上，实现了 4.58× 的加速。
- 相比之下，双向基准模型（DFlash-T）由于分支无关树的不一致性，在高预算下表现出收益递减或饱和现象。
服务性能： 集成到 vLLM 后，JETFLOW 在中小规模服务负载下持续优于基准模型。最优预算取决于负载：在小 Batch Size 下，较大的预算（256）能产生最高的吞吐量（高达 6.75× 加速）；而在重负载下，较小的预算在减少每步开销方面更高效。
消融实验洞察：
- 因果性至关重要： 将因果头替换为扩散头会导致性能大幅下降（例如在 MATH-500 上从 8.36× 降至 5.46× 加速），因为扩散头无法构建一致的树。
- 损失函数： 前向 KL 蒸馏优于反向 KL，后者会导致约 40% 的相对性能下降。
- 数据： 由目标模型重新生成的序列提供了最佳的训练信号，尽管在原始语料库数据上训练仍能获得一致的加速效果。

5. 意义与主张

论文声称 JETFLOW 通过解决草拟效率与接受质量之间的权衡，成功打破了投机解码的扩展瓶颈。

打破瓶颈： 作者证明，以往的方法无法有效扩展，是因为它们要么为了因果性支付了高昂的顺序成本（EAGLE），要么为了并行效率牺牲了一致性（DFlash）。JETFLOW 证明，通过在前缀隐藏状态上进行块级因果条件化，可以实现既廉价又具备因果一致性的并行草拟。
实际影响： 通过支持使用更大的草拟预算（如 256 个 Token）而不降低接受率，JETFLOW 显著降低了解码延迟。结果表明，这种方法对于复杂的推理任务（数学、代码）特别有效，并可部署在现实世界的服务环境（如 vLLM）中，以在现代硬件（H100, B200）上实现显著的吞吐量增益。
研究范围： 论文承认，虽然静态预算策略是有效的，但动态服务时预算调度仍是未来研究的一个开放领域。其主要贡献在于转向因果并行草拟的架构转变，这使得此前被认为受限于因果-效率困境的投机解码得以实现规模化扩展。