本文是 AI 模型协同设计系列文章的第三篇。文章探讨了如何利用推测解码在保持准确性的同时加速 LLM 推理,并提供了五条指导原则,用于在帕累托前沿上选择草稿长度和草稿机制。
关于模型设计选择如何在不牺牲准确性的情况下影响吞吐量和交互性的讨论,请参阅 AI 模型协同设计:硬件友好的 LLM 设计(第一部分)。关于组大小(每个 KV 头的查询头数)、头维度和序列长度如何影响密集注意力性能的解释,请参阅为快速、交互式长上下文推理协同设计 AI 模型注意力(第二部分)。
什么是推测解码?
推测解码是一种通过每次迭代预测多个 token 来加速 LLM 推理自回归解码阶段的技术。一个小型草稿模型首先预测几个可能的下一个 token。然后,这些 token 通过一次对较大目标模型的单次传递进行并行验证。

图 1. 推测解码将 LLM 解码推向更计算密集的区域,将吞吐量-交互性帕累托前沿向上并向右移动。随着交互性向曲线右侧增加,批大小减小,GEMM 在运行时间中所占比例比注意力更大这种方法减少了总解码迭代次数,同时增加了目标模型的算术强度,而无需更高的并发性。目标模型按顺序接受提议的 token,直到遇到第一个不匹配。然后下一个预测周期从该位置恢复。由于只保留目标模型接受的 token,推测解码产生与标准解码相同的输出序列,除非故意放宽接受标准。图 2 说明了这种解码流程。
草稿长度((D))是每个目标迭代提议的 token 数量。接受长度((AL))是每个目标迭代产生(即接受)的 token 数量。(AL) 的范围从 (1) 到 ((1 + D)),因为目标模型除了接受的草稿 token 外,总是可以产生一个新的真实 token。
推测解码的加速可以量化为目标模型顺序生成 (AL) 个 token 所需时间与并行验证 D 个 token 所需时间的比率,同时还要考虑起草 (D) 个 token 的延迟开销:
(\mathrm{speedup} = \frac{T_{\mathrm{verif}}(B) \times \mathrm{AL}}{T_{\mathrm{verif}}\left(B \times (1 + D)\right) + T_{\mathrm{draft}}(B, D)})
其中 (B) 是批大小,(T_{\mathrm{verif}}(x)) 是 (x) 个 token 的目标验证时间,(T_{\mathrm{draft}}(b, y)) 是为批大小 (b) 生成长度为 (y) 的草稿所需的时间。显然,为了最大化加速,有必要找到最优的 (\left(D, AL, T_{\mathrm{draft}}\right)) 元组。

(D)
图 2. 推测解码的草稿-然后-验证流程。较小的草稿模型提出
令牌,目标模型在一次传递中一起验证它们## 选择最优草稿长度
为简化起见,如果我们忽略草稿模型延迟,推测在以下情况下提供加速:
(\frac{T_{\mathrm{verif}}\left(B \times (1 + D)\right)}{T_{\mathrm{verif}}(B)} < AL)
在验证期间,计算随 ((1 + D)) 缩放,但内存访问保持不变。因此,目标是增加 (D) 直到 (T_{\mathrm{verif}}) 保持恒定,通常直到验证从内存受限转变为计算受限。显然,这个最优 (D) 值取决于 (B),并预计在帕累托前沿上变化。考虑图 1 时记住这一点。
草稿长度和线性层性能
通过推测,每个目标线性层 GEMM 的 GEMM-(M) 从 (M) 增长到 (M \times (1 + D)),其中 (M) 是无推测时的 GEMM-(M)。图 3 显示了在代表性专家 GEMM 大小为 6144 x 6144 时,不同草稿长度下每秒 teraflops 如何随批次缩放。显然,更高的草稿长度允许 GEMM 在较低的有效批次大小下达到峰值性能。
值得注意的是,当 (D=7) 时,与 (D=0) 相比,需要八分之一的批次大小才能成为计算受限。随着专家混合 (MoE) 模型变得更稀疏,长上下文工作负载增加 KV 容量压力,每个专家的有效并发性降低,使得更大的草稿长度在帕累托前沿上具有吸引力。

(D=0)
图 3. 归一化每秒 teraflops(相对于
,批次大小=1)与不同草稿长度下的批次大小,对于代表性专家 GEMM,GEMM-N = 6,144 和 GEMM-K = 6,144指南 1:增加推测解码草稿长度,将 GEMM 推入计算受限区域,而不增加 KV 缓存容量压力。
草稿长度和注意力性能
对于推理和代理工作负载,注意力往往在吞吐量导向区域主导执行时间。解码注意力具有约 (2 \times G) 的算术强度,其中 (G) 是共享一个 KV 头的查询头数量,如为快速、交互式长上下文推理共同设计 AI 模型注意力中所述。
推测将其增加到 (2 \times G \times (1 + D)),因为推测的令牌重用相同的 KV 缓存。有效注意力 GEMM-(M) 是 (G \times (1 + D))。在当前 GPU 设备上,注意力内核在 GEMM-(M=128) 时实现良好的硬件利用率,从而使 (D = \frac{128}{G} – 1) 成为最优草稿长度。
图 4 展示了在 32K 和 128K KV 序列长度下,当 (G = 8) 和 (G = 32) 时,归一化注意力吞吐量随 (D) 增加的变化。(G = 32) 的变体在较低的 (D) 值处就达到了吞吐量饱和。超过吞吐量饱和点后,注意力不再受 DRAM 带宽限制,其运行时间随 (D) 扩展。由于 (AL) 随 (D) 次线性增长,将 (D) 增加到该点以上很可能会拖慢注意力密集型工作负载。

(D=0)
图 4. 归一化注意力每秒万亿次浮点运算(相对于*)与*(D)
对于(G=8)
和(G=32)
。更大的(G)
在更小的草稿长度下达到高利用率
准则 2: 当注意力主导解码时间时,选择 (D = \frac{128}{G} – 1)。
注意力运行时间还取决于瓦片大小。图 5 显示,当 (G \times (1 + D)) 跨过 128 的倍数时,运行时间呈阶梯式增加,128 是所基准测试的注意力内核的软件瓦片大小。如果 (G \times (1 + D)) 落在两个瓦片边界之间,最后一个瓦片仅被部分利用,但成本仍与完整瓦片大致相同。

(D=0)
图 5.* 归一化注意力运行时间(相对于**(D)*
)与(G = 8)
对于(G = 32)
和(G \times (1 + D))
。运行时间呈阶梯式增加,当
跨过 128 的倍数准则 3: 如果选择 (D > \frac{128}{G} – 1),优先选择 (G \times (1 + D)) 为 128 的倍数的值,以避免瓦片利用率不足。
准则 1 相对于准则 2 和 3 的相对权重取决于在首选操作点上,运行时间中花费在 FFN 上的百分比与花费在注意力上的百分比之比。随着 D 增加,通信也会增加验证成本,尽管计算和通信的重叠可以减轻这种开销。
帕累托曲线最右端的草稿长度
向帕累托曲线的最右端移动时,(B) 变得非常小。此时,固定的内核设置和后处理开销在计算和通信内核中都占主导地位。这些固定成本不会随验证令牌数量的增加而显著增长,使得验证开销在很大程度上与草稿令牌数量无关。
虽然随着草稿长度的增加,MoE 确实会激活更多专家,但模型分片策略和高效内核(如分组 GEMM)的结合可以将这种开销保持在较低水平。因此,只要接受率保持较高,较大的草稿长度可以在低延迟区域提供帮助。
在极低延迟下,顺序内核启动的数量决定了工作负载延迟。由于内核启动次数随层数线性扩展,对于与目标模型具有相似层结构的自回归草稿模型,加速比可以近似为:
(\mathrm{speedup} = \frac{L_{\mathrm{target}} \times \mathrm{AL}}{L_{\mathrm{target}} + D \times L_{\mathrm{draft}}})
其中 (L_{\mathrm{target}}) 和 (L_{\mathrm{draft}}) 分别是目标模型和草稿模型中的层数。定义常数 draft_depth_ratio,(\rho = \frac{L_{\mathrm{draft}}}{L_{\mathrm{target}}}):
(\mathrm{speedup} = \frac{\mathrm{AL}}{1 + \rho D}) 以及草稿开销,(O_d = \rho D)
换句话说,只有当 (AL) 的增益大到足以抵消草稿开销时,增大 (D) 才有帮助。
准则 4:在极低延迟下,只有当 (AL) 的增益足以证明增加的草稿成本合理时,才增大 (D)。
选择草稿机制
选择 (D) 告诉你需要推测多少个 token。接下来,决定如何生成这些 token 以最大化加速。
多年来已经提出了大量技术,它们在训练、参数和运行时成本上各有不同的权衡。外部草稿建议使用一个小型独立 LLM,而 MTP、EAGLE-3、DFlash 和 DSpark 则使用辅助层结合目标模型的信息来预测 token。后缀和 n-gram 方法无需模型,而是复用 token 流中已经出现过的模式。
表 1 比较了主要的草稿方法,包括它们如何生成 token、其服务时内存以及草稿开销。
外部草稿模型 | EAGLE-3 | MTP | DFlash | DSpark | 后缀 / n-gram | |
草稿架构 | 小型 LLM | 解码器层 + 线性投影 | 解码器层 + 线性投影 | 解码器层 + 线性 KV 融合 | 解码器层 + 线性 KV 融合 + 轻量级马尔可夫头 | 字符串匹配,无模型 |
输入 | Token ID | 前一个 token 的嵌入 + 目标隐藏状态 | 目标最终隐藏状态 | 融合后的目标隐藏状态用作草稿 KV | 融合后的目标隐藏状态用作草稿 KV | Token 流 |
生成方法 | (D) 次自回归前向传播 | 运行模块 (D) 次 | 运行模块 (D) 次 | 在一个并行步骤中生成 (D) 个 token | 在一个并行步骤中生成 (D) 个 token,随后进行轻量级顺序校正 | 一次查找 |
发布方式 | 独立模型发布 | 独立的后训练头 | 基础模型检查点的一部分 | 独立的后训练头 | 独立的后训练头 | 独立于目标模型 |
训练成本 | 从头开始:1T-10T+ token 蒸馏:100B-400B 适配:10M-1B | 目标训练后额外 1-10B token | 通常在预训练期间与目标一起训练 | 目标训练后额外 1-10B token | 目标训练后额外 1-10B token | 无 |
激活内存成本 | 草稿权重 + 完整 KV 缓存 | 权重 + 小型 KV 缓存 | 权重 + 小型 KV 缓存 | 权重 + 从目标构建的小型 KV 缓存 | 权重 + 从目标构建的小型 KV 缓存 | 无 |
每 (D) 个 token 的推测成本 | 完整草稿模型前向传播 | 1 层,串行 | 1 层,串行 | 约 5 层,一次 | 约 5 层一次 + 轻量级马尔可夫头,串行 | O(1) 查找 |
最适合 | LPU 和 LPU+GPU 不推荐用于 GPU | 不推荐——(AL) 低于 MTP | 最适合 GPU 上的较大模型 | 较小模型,批大小为 1 | 较小模型,批大小为 1 | 高重复性的工作负载 |
表 1. 常见草稿机制在输入、生成方法、训练成本、服务时内存和草稿成本方面的比较为了量化 (AL) 与 (O_d) 之间的权衡,首先观察 (AL) 如何随 (D) 缩放。
图 6 展示了在 SPEED-Bench 上,以 Qwen 3.5 122B A10B 为目标模型时,AL 如何随 D 的增加而变化。SPEED-Bench 是 NVIDIA 开发的推测解码基准,旨在代表真实的生产工作负载。它涵盖编码和摘要等多个任务领域,并针对不同输入序列长度提供多种划分。我们推荐使用 SPEED-Bench 进行 (AL) 比较。
在 32K 划分上,Qwen 3.5 35B A3B 在 (D = 9) 时达到 (AL) 为 6,而 4B 草稿模型达到的 AL 高于 5。MTP 和 DFlash 的 AL 随 (D) 增长而趋于平缓。N-gram 在此工作负载上的接受率较低,更适合具有重复 token 模式的工作负载。

(AL)
图 6。**(D)
与
在 SPEED-Bench 上不同草稿机制的对比,吞吐量 32K 划分,以 Qwen 3.5 122B A10B 为目标模型更高的 (AL) 并不等于更高的加速比。你还需要考虑生成草稿的成本。
如图 6 所示,在 (D > 3) 时,所有外部草稿都达到比其他草稿方法更高的 (AL)。Qwen 3.5 122B MTP 要小得多,总参数为 2.5B,活跃参数少于 150M。在 (D = 3) 时,更大的外部草稿可能不值得增加的成本。在 (D = 21) 时,假设采用高效的服务策略,其更高的 AL 可能使该成本变得值得。
与此同时,DFlash 的 (AL) 相对较快地趋于平缓。然而,DFlash 和 DSpark 都并行生成 (D) 个草稿 token,从而降低了 (O_d)。作为对比,在最小延迟且 (D=11) 的情况下,单层 MTP 头需要 11 步,导致 (O_d^{\mathrm{MTP}} = \frac{11}{L_{\mathrm{target}}}),而五层 DFlash 头一次前向即可生成草稿,得到 (O_d^{\mathrm{DFlash}} = \frac{5}{L_{\mathrm{target}}})。
对于层数众多的大型目标模型,这两种开销都可以忽略不计,但随着 (L_\mathrm{target}) 减小,草稿开销变得显著。因此,对于低延迟的小型模型,即使 (AL) 较低,DFlash 或 DSpark 也可能成为最优选择。
为了在 (AL) 与草稿延迟之间进行权衡,准确地对两者进行基准测试至关重要。在推测解码的语境下,这意味着使用真实的提示词进行基准测试,并覆盖广泛的任务领域。我们推荐使用 SPEED-Bench 进行 (AL) 测量,并使用 NVIDIA TensorRT LLM 等高性能推理框架来量化草稿开销。
除了推理性能之外,草稿训练的范围和成本也值得考虑。MTP 必须与目标模型联合训练,而 EAGLE、DFlash 和 DSpark 可以添加到最终模型检查点中。同样,外部草稿模型的创建也存在一系列连续选项。从头训练和从目标模型蒸馏可获得最高的 AL,而通过跨模型适配技术将现有模型适配到所需目标,可以显著降低训练成本,代价是牺牲 (AL)。
微调目标模型会改变其输出分布和隐藏表示。因此,针对特定目标检查点训练的学习型草稿模型即使在目标模型改进后也可能失去接受率。更换目标模型后,需在代表性工作负载上重新测量 (AL)。所需的适配取决于草稿机制。
MTP 是目标模型的一部分,应在微调期间继续训练,或在专门的后续阶段重新对齐。目标附加型草稿模型(如 EAGLE-3、DFlash 和 DSpark)使用目标的隐藏状态,需要通过额外训练来适配更新后的检查点。
外部草稿模型不消耗目标隐藏状态,但它们仍然近似目标的输出分布,可能需要微调或蒸馏。后缀和 n-gram 方法没有学习型草稿模型,因此无需重新训练,尽管其有效性仍取决于部署工作负载中出现的重复程度。
准则 5: 根据 (AL)、草稿延迟以及训练和部署成本,选择能为你的工作负载和硬件带来最佳解码加速的草稿机制。
开始协同设计投机解码
使用下面总结的五条准则作为在帕累托前沿上选择 (D) 和草稿机制的检查清单。
- 增加投机解码草稿长度,将 GEMM 推入计算受限区域,同时不增加 KV 缓存容量压力。
- 当注意力主导解码时间时,选择 (D = \frac{128}{G} – 1) 作为起点。
- 对于更大的草稿长度,优先选择 (G \times (1 + D)) 为 128 的倍数的值,以与注意力内核瓦片大小对齐。
- 在极低延迟下,使用快速草稿机制,并且仅在 (AL) 的增益证明增加的草稿成本合理时才增加 (D)。
- 通过平衡 (AL) 和草稿开销来选择草稿机制,以适应你的工作负载和硬件。在真实服务条件下对两者进行基准测试,并考虑训练和部署成本。
对于后训练草稿模型,NVIDIA/Model-Optimizer 中提供了 EAGLE-3、DFlash 和 DSpark 的即用型训练示例。了解我们如何为 NVIDIA Nemotron 3.5 Lightning 完成此操作:微调 DSpark,然后将其量化为 FP8 或 NVFP4。将这些示例作为起点,然后在你自己的模型、工作负载和硬件上验证 AL 和端到端速度。
致谢
本文是 NVIDIA 的跨团队合作成果。我们感谢 Bhargava Gopireddy、Ritika Borkar、Dor Tsur、Andrii Skliar、Benjamin Chislett、Yaniv Galron、Talor Abramovich、Yoav Miron、Rabeeh Karimi Mahabadi、Roger Waleffe、Udi Karpas、Ran Zilberstein、Brian Pharris 和 Eduardo Alvarez 的帮助。
