JetFlow：通过并行树起草打破推测解码的扩展天花板

背景

推测解码（Speculative Decoding, SD）是一种旨在加速自回归大型语言模型（LLM）推理的技术。其核心思想是利用一个小模型（草稿模型）快速生成多个候选 token，然后由一个大模型（目标模型）并行验证这些 token。如果候选 token 被接受，推理速度将显著提升；如果拒绝，则需回退并重新生成。

尽管推测解码在理论上能带来巨大的加速比，但在实际应用中面临着明显的“扩展天花板”（Scaling Ceiling）。当增加草稿预算（draft budget，即每次尝试生成的候选 token 数量）时，加速效果往往不再线性增长，甚至可能下降。这主要受限于两个关键因素：

1. 接受率（Acceptance Rate）：随着候选序列变长，小模型生成的 token 与大模型分布的一致性降低，导致接受率下降。
2. 起草开销（Drafting Overhead）：生成更多候选 token 需要更多的计算资源，如果验证带来的收益无法覆盖生成成本，整体效率反而降低。

现有的基于头部（head-based）的推测解码方法陷入了“因果性-效率困境”（causality-efficiency dilemma）：

• 自回归式草稿生成器：能够产生路径条件化的候选项，适合用于树状推测解码，能获得较长的接受长度。但其起草成本随树深度增加而线性增长，效率受限。
• 双向块扩散草稿生成器：能够在一次前向传播中生成所有位置的 token，效率极高。但由于其生成的边缘分布（marginals）与分支无关，容易形成“ individually plausible yet mutually inconsistent ”（单个看似合理但相互不一致）的树结构，导致预算浪费和接受率降低。

核心内容

为了解决上述困境，研究团队提出了 JetFlow，一种基于头部的推测解码框架。JetFlow 的核心创新在于结合了“单次前向传播的效率”与“分支级因果条件化”的优势。

1. 架构设计：因果并行草稿头

JetFlow 不再训练一个独立的草稿模型，而是在冻结的目标模型（Target Model）之上，训练一个因果并行草稿头（Causal Parallel Draft Head）。

• 输入：该草稿头作用于目标模型融合后的隐藏状态（fused hidden states）。
• 机制：它通过一次前向传播，并行地生成候选树。
• 对齐性：关键在于，JetFlow 生成的候选树的分数与目标模型的自回归因子分解（autoregressive factorization）保持一致。这意味着，虽然它是并行生成的，但其内部逻辑遵循自回归的因果约束。

2. 突破扩展天花板

传统方法中，增加草稿预算往往导致接受率急剧下降或开销过大。JetFlow 通过以下方式将更大的草稿预算转化为更长的接受前缀（accepted prefixes）：

• 一致性保证：由于草稿头的分数与目标模型对齐，生成的候选树在概率上是连贯的，避免了双向方法中常见的“逻辑冲突”问题。
• 效率优化：通过并行生成而非严格的自回归展开，JetFlow 大幅降低了生成深层树结构的计算开销。

3. 实验结果

研究团队在密集模型（Dense）和混合专家模型（MoE）的 Qwen3 系列模型上，针对数学推理、代码生成和聊天对话等多个基准进行了广泛测试。

• 性能对比：JetFlow 在各项指标上均一致优于双向头部（bidirectional-head）和基于树的推测解码基线方法。
• 加速比：
  • 在 H100 GPU 上，JetFlow 在 MATH-500（数学基准）测试中实现了高达 9.64倍 的加速。
  • 在开放式对话工作负载中，实现了 4.58倍 的加速。
• 实际部署：通过与 vLLM 集成，在真实的推理服务负载下，JetFlow 进一步展示了延迟降低的优势。

关键要点

• 解决因果-效率困境：JetFlow 成功平衡了自回归草稿的因果一致性与双向草稿的并行效率，消除了以往方法中“要么慢但准，要么快但不准”的两难局面。
• 冻结模型上的轻量级训练：无需重新训练庞大的目标模型，仅在冻结的目标模型隐藏状态上训练一个轻量级的并行草稿头，降低了部署和维护成本。
• 分数对齐机制：通过确保草稿生成的分数与目标模型的自回归因子分解对齐，JetFlow 能够生成高质量、逻辑一致的候选树，从而显著提高 token 接受率。
• 显著的端到端加速：在 Qwen3 模型上，JetFlow 不仅提升了数学和代码任务的速度，在自然语言对话场景中也表现出卓越的加速能力（最高 9.64x）。
• 开源与集成：代码和模型已开源，并且证明了与主流推理引擎 vLLM 的良好兼容性，便于工业界落地。

意义与影响

JetFlow 的提出标志着推测解码技术在工程实践和理论优化上的重要进步。

1. 重新定义推测解码的扩展性：以往认为推测解码存在固有的扩展天花板，即增加并行度必然牺牲质量或增加开销。JetFlow 证明了通过合理的架构设计（因果并行头），可以打破这一限制，使更大的草稿预算真正转化为推理速度的提升。
2. 降低 LLM 部署成本：对于依赖 Qwen3 等大模型的企业而言，JetFlow 提供的近 10 倍加速意味着在同等硬件条件下可以处理更高的并发请求，或在同等请求量下大幅降低 GPU 算力成本。
3. 推动 MoE 模型的高效推理：研究特别提到了在 MoE（混合专家）模型上的表现，这对于当前主流的稀疏大模型架构具有直接的商业价值，因为 MoE 模型的推理开销通常更高，加速潜力更大。
4. 方法论的普适性：虽然实验基于 Qwen3，但 JetFlow 提出的“冻结目标模型+因果并行草稿头”的范式具有通用性，未来可应用于其他大型语言模型，推动整个行业向更高效、更经济的 LLM 服务迈进。