阶跃开源JetSpec,大模型推测解码提速近10倍
随着 Agent 应用爆发大模型推理效率问题成为 AI 行业焦点。当基础能力逐渐够用后效率就成了新的主战场。推理效率决定了模型能被怎样调用、是否普遍用得起、能铺多广。近期除了 DeepSeek 开源 DSpark 推理加速框架阶跃星辰StepFun也联合加州大学圣地亚哥分校、浙江大学、伊利诺伊大学、南京大学团队开源了 JetSpec。JetSpec 在 H100 上把 MATH-500 的解码速度拉到 9.64 倍开放式对话也能到 4.58 倍。这套面向推测解码Speculative DecodingSD的新框架用一次前向传播同时产出一棵带因果条件的候选树把推测解码长期存在的扩展天花板撕开了一道口子。解码的天花板大模型自回归解码每次只产出一个 token下一个 token 必须等上一个算完。串行特性让长输出场景的延迟居高不下模型越大、输出越长等待越痛苦。业界有几条路子缓解这个问题预训练或中期训练阶段引入多 token 预测MTP或者把预训练模型改造成扩散 LLMdLLM代价是不同程度的质量损失和较高的适配成本。推测解码是当下最实用的解法之一只用一个轻量级草稿模型先猜 N 个 token再让目标模型一次性并行验证接受其中最长的一致前缀不动模型本身质量无损。轻量对齐往往不到 10 亿 token 就能训出不错的草稿头部署也比维护独立草稿模型的头部方法省事。理论上草稿猜得准、猜得便宜加速比就能往上走。而实际并没那么简单。SD 的加速比有一个公式化的天花板。设 α 为平均接受率c 为草稿模型单步开销和目标模型单步开销的比值N 为草稿长度期望加速比大致是N 越大分子越大分母里的 Nc 也越大。α 稍微往下掉或者 c 没压到足够低扩展 N 的收益就被吃掉。换句话说盲目堆草稿预算不一定换得来更多加速。Figure 2 把这件事画得很直白。左图是典型 SD 设定c0.05草稿长度超过某个区间曲线就躺平甚至下滑再堆预算也是白搭。右图把 c 压到 0.0005 的极端情形曲线能一路爬到 20 倍加速α 越高爬得越狠。想把草稿预算加大换来更高加速必须同时把 c 压低、把 α 抬高两头都不能放。过往方案各做一头谁也没同时拿下。一头是基于头部的自回归草稿器代表是 EAGLE 系列。它走路径条件化的路子每个候选 token 都建立在前序选中的 token 之上分支之间因果清晰接受率 α 高适合做树形推测解码。代价是草稿开销随树深增长树越深自回归步数越多c 压不下来。树一深开销就把并行验证省下的时间吃回去了。另一头是双向块扩散草稿器代表是 DFlash。一次前向就把多个位置全预测出来c 极低。每个位置的预测是分支无关的边缘分布不依赖这条分支上具体选了哪些祖先 token。问题就出在组合环节单看每个位置都合理拼起来可能互相打架。这就是因果与效率两难causality-efficiency dilemma。自回归草稿器有因果没效率块扩散草稿器有效率没因果。草稿预算一旦加大自回归那头被开销拖死块扩散那头被不一致的树结构浪费掉扩展天花板就这么卡住了。要破局得想办法一次前向产出整棵树同时让每个节点都看到自己分支上的祖先。JetSpec 走的就是这条路。一次前向的解法JetSpec 的核心是一个因果并行草稿头causal parallel draft head。它复用冻结目标模型的多层隐藏状态作为输入特征一次前向就输出整棵候选树所有节点的 logits。和 DFlash 一样是一次前向c 压得很低。和 EAGLE 一样保留分支因果α 拉得上去。关键在树形因果注意力掩码tree-causal attention mask。每个树节点只能看到原始前缀和它自己这条分支上的祖先节点看不到后代也看不到兄弟分支。所有节点可以在一次前向里并行计算每条分支内部依然保持自回归式的依赖结构。从概率分解的角度看JetSpec 的草稿分布按分支做因式分解每个节点的分布都以本分支的具体祖先 token 为条件和目标模型的自回归因式分解对齐。DFlash 那种分支无关边缘分布带来的不一致问题就根除了。训练上JetSpec 用前向 KL 散度forward KL做软标签蒸馏让草稿分布保留目标模型对多个合理续写的相对偏好。训练数据用 780K 条 Nemotron Post-Training Dataset V2 的样本加 20K 条 CodeAlpaca覆盖数学、编程、对话。构造候选树时JetSpec 每个解码步从上一步验证接受的隐藏状态出发设好最大树深 N、分支宽度 W、总节点预算 B一次前向拿到所有深度的 logits每层取 top-W 候选再用累积草稿对数概率做分支打分贪心扩展直到预算填满。目标模型再用树注意力并行验证整棵树按推测解码的接受规则取最长一致前缀。速度跑出来了在 Qwen3-8B稠密和 Qwen3-30B-A3BMoE上实测覆盖数学、编程、对话三类基准。树预算 256 token 的端到端加速比对比JetSpec 在 MATH-500 上跑到 9.64 倍GSM8K 7.82 倍AIME25 8.78 倍HumanEval 7.12 倍MBPP 6.73 倍LCB 7.67 倍MT-Bench 4.58 倍。DDTree 最高在 MATH-500 上 8.78 倍DFlash 在 MATH-500 上 6.12 倍。EAGLE-3 在树模式下受训练失配限制大预算收益甚微。MoE 模型上同样成立。Table 5 是 Qwen3-30B-A3B 的结果JetSpec 在 MATH-500 上 9.45 倍加速、τ10.65AIME25 9.35 倍全面压过 DDTree。因果树草稿的思路不挑模型架构。EAGLE-3 受自回归草稿开销拖累预算 16 时加速比普遍在 2 倍出头JetSpec 同预算下能到 4 到 6 倍。预算从 16 加到 32DFlash 多数基准饱和或下滑JetSpec 还能小幅往上走说明多出来的预算被有效转化。服务场景的结论也有点意思。批量大小 1 时树预算从 16 加到 128吞吐从 443.3 TPS 涨到 968.2 TPS加速比从 3.09 倍到 6.75 倍。批量加大收益缩水批量 16 时预算 256 只剩 4.51 倍批量 32 时 2.85 倍。树预算要跟着服务负载走低负载时大预算换吞吐高负载时小到中等预算反而更划算验证开销和计算压力会把大预算的收益吃掉。B200 这类高端卡上收益更明显。团队目前只评估了静态预算策略动态按负载调预算留给后续工作。JetSpec 用一个树形因果掩码把并行草稿的低开销和分支因果的高接受率同时拿到手。推理效率的天花板正在被一步步攻破智能普惠的时代正在到来。参考资料https://jetspec-project.github.io/jetspec-web/https://huggingface.co/JetSpechttps://arxiv.org/pdf/2606.18394https://github.com/hao-ai-lab/JetSpec