JetSpec: 使用并行树草稿突破推测解码的扩展上限

TL;DR JetSpec 通过因果并行草案头，将更大预算转为更长接受长度与更高加速，突破推测解码扩展瓶颈，最高达 9.64x 加速。

问题

问题背景

推测解码是加速自回归大语言模型推理的关键技术，通过并行验证多个候选 token 来减少串行解码步数。但现有方法面临扩展天花板：增加草案预算仅在接受率高且草案开销低时带来加速。

现有方法局限

主流 head-based 草案方法分为两类，均陷入因果-效率困境：

• 自回归草案(如 EAGLE、Medusa)：保持因果依赖，草案 token 条件于已生成的前缀，树推测解码的接受长度更高，但草案生成需逐步进行，计算开销随树深度线性增长。
• 双向块扩散草案(如 SpecDiff、BlockDiff)：一次前传生成所有草案位置，效率高，但其分支无关边际概率无法建模 token 间因果，导致任意分支上的 token 彼此孤立，易产生“个体合理、整体矛盾”的候选树，降低有效接受长度，浪费验证算力。

技术挑战与重要性

突破这一天花板难度在于：既要维持因果依赖以提高接受率，又要降低草案生成的计算开销。在批处理推理中，草案开销占比随 batch size 增大而显著上升，限制了大预算推测解码的实际收益。若不能根本性解决，即使堆叠算力，吞吐量提升也会快速饱和，直接制约高吞吐在线服务场景下的成本与延迟优化。业界对高性能 LLM 推理引擎（如 vLLM）的推测解码集成需求强烈，急需一种既能并行高效生成、又能保持分支因果一致的草案架构。

行业类比

如同 CPU 的分支预测单元必须在预测深度和吞吐量间权衡，推测解码也需设计一个轻量“预测头”，在单次前传中获得准确的 token 树。

核心洞察

• 并行树草拟同时实现了一次前向效率与分支因果条件化，从根本上化解了投机解码中因果-效率困境。与自回归草拟（成本随树深线性增长）和双向扩散草拟（形成互不一致的边际分布）不同，JetSpec 通过一个融合目标模型隐藏状态的因果并行草稿头，在单次前向中生成整棵候选树，且每个节点都基于其祖先路径条件化，保证了树内节点之间的因果一致性，从而在高 draft budget 下仍能维持高接受率。
• JetSpec 揭示了投机解码扩展上限可被有效突破的关键在于将 draft budget 转化为更长的接受前缀，而非单纯扩大树规模。传统方法在 budget 增加时往往遭遇接受率下降或草拟开销陡增，而 JetSpec 的对齐训练使草稿头的 logit 与目标模型自回归分解高度一致，使得更大的树预算能够线性地转化为更长的平均接受长度，在密集模型和 MoE 架构上均实现高达数倍的端到端加速，为投机解码的实用化提供了新的 scaling 范式。

方法

JetSpec 方法围绕三个核心环节：冻结目标模型提取隐藏状态 → 因果并行草稿头生成候选树 → 树形验证与接受。

输入与隐藏状态融合
给定前缀序列，首先用冻结的目标 LLM 执行一次前向传播，获取多层隐藏状态。JetSpec 将选定层的状态融合（fused hidden states），作为后续草稿头的条件信号。这一步几乎零额外开销，复用了目标模型已有的表示能力。

因果并行草稿头（Causal Parallel Draft Head）
草稿头是一个轻量级 Transformer 模块，接收融合后的隐藏状态，并在单次前向传播中并行生成多个未来位置的令牌分布。其关键设计在于 因果注意力掩码：每个待草稿位置只能关注到前缀和自身之前的草稿位置，从而保证序列的条件依赖关系（branch-wise causal conditioning）。但所有位置的令牌是同时预测的，避免了逐令牌自回归的串行开销。对于树形草稿，草稿头在每一层采用 top‑k 采样或分支扩展策略，产出多棵候选分支构成的候选树，树的每个节点携带对应位置的目标模型对齐分数。

训练目标
草稿头通过蒸馏损失训练，使其在每个草稿位置输出的分布逼近目标模型在该位置的自回归条件分布（即给定前缀和正确历史令牌时的概率）。训练数据来自目标模型自身的生成结果或真实文本，确保草稿头学会“模仿”目标模型的因果分解，而不仅仅是拟合独立边缘分布。

并行树验证
草稿树构建完成后，一次性送入目标模型进行并行验证。目标模型计算草稿树中所有节点的对数概率，并依据推测解码接受准则（如严格 top‑p 或贪婪匹配）选择最长的匹配前缀。未被接受的令牌则被丢弃，最终输出被接受的令牌序列。整个过程将草稿预算（树的总节点数）转化为实际加速收益，预算利用率高。

与同类方法的差异：相比自回归草稿器（如 Medusa）随树深度线性增长的草稿开销，JetSpec 凭借并行草稿保持固定前向次数；相比双向块扩散草稿器（如 SpecDiff）分支相互割裂、易产生不一致树的问题，JetSpec 的因果并行设计保证了候选树内部路径的条件一致性，从而在高预算下仍维持高接受率，突破传统投机解码的扩展天花板。

实验

实验设计

JetSpec 的评估覆盖 密集模型与 MoE Qwen3 变体，在 数学 (MATH-500)、代码 和 开放域对话 三类 benchmark 上测试，涵盖低预算与高预算两种推测解码场景。基线包括自回归 drafter、双向块扩散 drafter 以及现有树解码方案。JetSpec 通过从冻结目标模型的融合隐状态训练一个 因果并行 draft head，实现单次前向生成候选树。系统级实验在 H100 GPU 上运行，并集成 vLLM 模拟真实服务负载。

关键发现

• 打破扩展瓶颈：JetSpec 成功将更大的 draft budget 转化为更长的接受前缀，在高低预算下均持续优于基线。
• 显著加速：在 MATH-500 上获得最高 9.64× 加速，开放域对话任务上 4.58× 加速，且通过 vLLM 集成进一步降低延迟。
• 因果条件的作用：与双向 head 相比，分支级因果条件避免了互不一致的 token 树，接受率更高；与自回归 drafter 相比，单次前向起草大幅降低开销，使得深树在预算增加时仍有效。

基线对比深度解读

自回归 drafter 虽能产出路径条件化的候选，但起草成本随树深度线性增长，高预算下不成比例地消耗算力。双向块扩散 drafter 一次生成所有位置，但其 分支无关的边缘概率 常生成单个合理但整体矛盾的树，浪费预算、降低接受。JetSpec 融合两者优势：单次前向的起草效率 + 分支因果条件。这使其在高预算下仍能将额外预算转化为更高接受长度，而双向 drafter 的接受率随预算增加出现饱和或下降。vLLM 集成结果进一步证明，在真实服务中 JetSpec 的低起草延迟能转化为端到端吞吐提升，而其他方法可能因起草开销过高或接受不足而优势有限。

行业影响

落地场景

JetSpec 面向大语言模型在线推理场景，尤其适合需要低延迟、高吞吐的实时生成任务。可落地的产品与业务包括：

• 对话式 AI 服务：智能客服、虚拟助理、语音助手的文本生成后端，对 <100ms 首 token 延迟敏感；
• AI 编程助手：如代码补全、对话式编程，需在用户输入间隙快速给出多 token 建议；
• 内容平台与教育：实时续写、数学解题、作文批改等交互式应用；
• 游戏 NPC 对话：需大量并发、严格控制推理预算的场景。

商业价值

• 降本：将相同吞吐下所需的 GPU 数量减少 2-4 倍，直接降低推理算力成本。在半精度 H100 上，MATH-500 任务取得 9.64× 加速比，日常对话负载加速达 4.58×，意味着同样硬件支撑更多请求。
• 体验提升：延迟显著降低，用户体验改善，可带来更高的交互完成率和用户留存。
• 增收：对于 API 型服务，吞吐提升可直接转化为收入增长，无需额外硬件投入。

与现有工作流的接口

JetSpec 的设计遵循即插即用原则：

• 目标模型保持冻结，不修改权重，只需额外训练一个轻量因果并行草稿头；
• 已适配主流推理引擎 vLLM，可作为投机解码插件直接加载，无需重构服务架构；
• 训练数据仅需将目标模型隐藏状态与文本对齐，资源需求远低于全模型微调。

集成步骤清晰：在 vLLM 部署中，将草稿头模型文件放置指定路径，通过配置文件启用 speculative decoding 组件，即可获得加速。

具体落地用例

1. 电商智能客服：头部电商平台每天面临数千万次用户咨询，部署基于 Qwen3-MoE 的对话代理。使用 JetSpec 后，单个 H100 节点可承载的并发会话量翻倍，同时回复延迟从 >500ms 降至 <200ms，直接减少支撑同等服务的 GPU 集群规模。
2. 教育科技公司：在线编程教学平台提供实时代码助教，学生输入提示后需秒级生成多行代码。在 MATH-500 和 HumanEval 类任务上 JetSpec 提速显著，可让平台在不增加硬件成本的前提下，为更多学生提供接近实时的反馈。

局限

• **训练依赖与模型特定性**：JetSpec 需要为每个冻结的目标模型额外训练一个因果并行草稿头，训练过程涉及数据准备和蒸馏损失调优，这增加了适配新模型或模型更新时的工程成本，相比无训练或参数共享的草稿方法（如自推测解码）通用性较低，可能限制在需要快速切换到不同模型或频繁微调的场景中的部署灵活性。
• **验证阶段树注意力开销**：随着草稿预算（如分支数和树深）增大，树注意力并行验证的计算复杂度上升，尤其在长序列生成或低端硬件上，验证延迟可能抵消部分草稿带来的加速。论文主要在 H100 上测试，对于计算资源受限的环境或超大模型（例如 70B+），加速效果可能缩水，且未讨论内存带宽瓶颈对加速上限的影响。
• **实验模型和任务覆盖面有限**：评估主要基于 Qwen3 系列的密集和 MoE 配置，未涉及其他主流架构（如 Llama、Mistral）或更大规模的模型，也未测试长文本生成、多轮对话等复杂场景下树草稿的质量衰退和预算利用率。因此，JetSpec 在不同模型家族上的泛化能力及其扩展至更大参数规模时的加速稳定性仍待验证。