阶跃开源JetSpec，大模型推测解码提速近10倍

发布日期：2026-07-01 07:07:18 浏览次数： 1508

作者：算泥

微信搜一搜，关注“算泥”

算泥社区是集 “AI 大模型开发服务 + 算法 + 算力” 于一体的开源生态社区，欢迎关注！

随着 Agent 应用爆发，大模型推理效率问题成为 AI 行业焦点。当基础能力逐渐够用后，效率就成了新的主战场。

推理效率决定了模型能被怎样调用、是否普遍用得起、能铺多广。

近期，除了 DeepSeek 开源 DSpark 推理加速框架，阶跃星辰（StepFun）也联合加州大学圣地亚哥分校、浙江大学、伊利诺伊大学、南京大学团队开源了 JetSpec。

JetSpec 在 H100 上把 MATH-500 的解码速度拉到 9.64 倍，开放式对话也能到 4.58 倍。

这套面向推测解码（Speculative Decoding，SD）的新框架，用一次前向传播同时产出一棵带因果条件的候选树，把推测解码长期存在的扩展天花板撕开了一道口子。

解码的天花板

大模型自回归解码，每次只产出一个 token，下一个 token 必须等上一个算完。串行特性让长输出场景的延迟居高不下，模型越大、输出越长，等待越痛苦。

业界有几条路子缓解这个问题，预训练或中期训练阶段引入多 token 预测（MTP），或者把预训练模型改造成扩散 LLM（dLLM），代价是不同程度的质量损失和较高的适配成本。

推测解码是当下最实用的解法之一，只用一个轻量级草稿模型先猜 N 个 token，再让目标模型一次性并行验证，接受其中最长的一致前缀，不动模型本身，质量无损。轻量对齐往往不到 10 亿 token 就能训出不错的草稿头，部署也比维护独立草稿模型的头部方法省事。

理论上，草稿猜得准、猜得便宜，加速比就能往上走。

而实际并没那么简单。SD 的加速比有一个公式化的天花板。设 α 为平均接受率，c 为草稿模型单步开销和目标模型单步开销的比值，N 为草稿长度，期望加速比大致是：

N 越大，分子越大，分母里的 Nc 也越大。α 稍微往下掉，或者 c 没压到足够低，扩展 N 的收益就被吃掉。换句话说，盲目堆草稿预算，不一定换得来更多加速。

Figure 2 把这件事画得很直白。左图是典型 SD 设定（c=0.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-A3B（MoE）上实测，覆盖数学、编程、对话三类基准。

树预算 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 在树模式下受训练失配限制，大预算收益甚微。