论文
JetSpec:通过并行树绘图打破 推测解码 的扩展天花板
JetSpec: Breaking the Scaling Ceiling of Speculative Decoding with Parallel Tree Drafting
摘要
推测解码 (SD) 通过起草多个词元并并行验证它们来加速自回归 大语言模型 (LLM),但它面临着扩展限制:只有当接受度保持较高且起草开销保持较低时,增加草稿生成预算才能提高速度。这个上限很难打破,因为之前基于头部的 SD 方法面临着因果效率的困境。自回归绘图器生成对具有更高接受长度的树 推测解码 有效的路径条件候选,但它们的绘图成本随着树深度而增加。双向块扩散绘图器一次性生成所有位置,但其与分支无关的边缘可能会形成单独合理但相互不一致的树,从而浪费预算并降低接受度。我们提出了 JetSpec,一种基于头部的 SD 框架,它将单前向起草效率与分支因果条件作用相结合。 JetSpec 在来自冻结目标模型的融合隐藏状态上训练因果并行 草稿预测头,生成其分数与目标模型的自回归分解一致的候选树。这使得 JetSpec 能够将更大的草稿生成预算转换为更长的可接受的前缀和更高的端到端加速。在密集模型和 MoE Qwen3 模型上的数学、编码和聊天基准测试中,JetSpec 始终优于双向头和基于树的 SD 基准。在 H100 GPU 上,JetSpec 在 MATH-500 上实现了高达 9.64 倍的加速,在开放式会话工作负载上实现了 4.58 倍的加速,并通过实际服务负载下的 vLLM 集成证明了进一步的延迟增益。我们的代码和模型可在 https://github.com/hao-ai-lab/JetSpec 获取。