NVIDIA 详解推测解码：五条准则优化 LLM 推理加速

本文是 AI 模型协同设计系列的第三篇。文章探讨如何利用推测解码（speculative decoding）在保持精度的同时加速 LLM 推理，并给出五条准则，用于在帕累托前沿上选择草稿长度和草稿机制。

关于模型设计选择如何在不牺牲精度的前提下同时影响吞吐量和交互性，请参阅 AI Model Co-Design: Hardware-Friendly LLM Design（第 1 篇）。关于分组大小（每个 KV 头对应的查询头数）、头维度和序列长度如何影响稠密注意力的性能，请参阅 Co-Designing AI Model Attention for Fast, Interactive Long-Context Inference（第 2 篇）。

什么是推测解码？

推测解码是一种加速 LLM 推理中自回归解码阶段的技术，其做法是每次迭代预测多个 token。首先由一个小型草稿模型预测接下来可能出现的若干 token，然后用较大的目标模型通过一次前向传播并行验证这些 token。

这种方法减少了解码迭代的总次数，同时提高了目标模型的算术强度，且不需要更高的并发度。目标模型会依次接受提议的 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)$ 元组。

选择最优草稿长度

为简单起见，若忽略草稿模型的延迟，则在以下条件成立时推测能带来加速：

$\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 展示了对于一个具有代表性的 6144 x 6144 专家 GEMM 尺寸，在不同草稿长度下，每秒万亿次浮点运算（teraflops）如何随批大小变化。显然，草稿长度越大，GEMM 就能在越低的有效批大小下达到峰值性能。

值得注意的是，当 $D=7$ 时，进入计算受限状态所需的批大小只有 $D=0$ 时的八分之一。随着混合专家（MoE）模型变得越来越稀疏，且长上下文工作负载加大了 KV 容量压力，每个专家的有效并发度会下降，这使得较大的草稿长度在整个帕累托前沿上都颇具吸引力。

准则 1：增大推测解码的草稿长度，把 GEMM 推入计算受限区域，同时不增加 KV 缓存容量压力。

草稿长度与注意力性能

对于推理和智能体工作负载，在以吞吐量为导向的区域，注意力往往主导执行时间。如 Co-Designing AI Model Attention for Fast, Interactive Long-Context Inference 中所述，解码注意力的算术强度约为 $2 \times G$，其中 $G$ 是共享同一个 KV 头的查询头数量。

推测会把它提高到 $2 \times G \times (1 + D)$，因为被推测的 token 复用同一份 KV 缓存。有效的注意力 GEMM-$M$ 为 $G \times (1 + D)$。在当前的 GPU 设备上，注意力内核在 GEMM-$M=128$ 时能获得良好的硬件利用率，因此 $D = \frac{128}{G} – 1$ 就是最优草稿长度。

图 4 展示了在 $G = 8$ 和 $G = 32$、KV 序列长度为 32K 和 128K 的情况下，归一化的注意力吞吐量如何随 $D$ 增大而变化。$G = 32$ 的变体在较低的 $D$ 值处达到吞吐量饱和。超过吞吐量饱和点后，注意力不再受 DRAM 带宽限制，其运行时间随 $D$ 增长。由于 $AL$ 随 $D$ 呈次线性增长，把 $D$ 增大到这一点之后，很可能会拖慢以注意力为主的工作负载。

准则 2：当注意力主导解码时间时，选择 $D = \frac{128}{G} – 1$。

注意力的运行时间还取决于分块（tile）大小。图 5 显示，当 $G \times (1 + D)$ 跨过 128 的倍数（即被测注意力内核的软件分块大小）时，运行时间呈阶梯式上升。如果 $G \times (1 + D)$ 落在两个分块边界之间，最后一个分块只被部分利用，但开销仍与一个完整分块大致相当。

准则 3：如果选择 $D > \frac{128}{G} – 1$，应优先选择使 $G \times (1 + D)$ 为 128 的倍数的取值，以避免分块利用不足。

准则 1 相对于准则 2 和 3 的权重，取决于在偏好的工作点上，FFN 与注意力各占运行时间的百分比。随着 D 增大，通信也会增加验证成本，不过计算与通信的重叠可以缓解这部分开销。

帕累托曲线最右端的草稿长度

向帕累托曲线最右端移动时，$B$ 变得非常小。此时，无论是计算内核还是通信内核，固定的内核启动和后处理开销都占据主导。这些固定成本不会随被验证 token 数量显著增长，因此验证开销基本不随草稿 token 数量变化。

虽然随着草稿长度增加，MoE 确实会激活更多专家，但模型分片策略与 Grouped 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 才能使加速比最大化。

多年来人们提出了大量技术，它们在训练、参数和运行时成本上各有取舍。外部草稿（external drafting）使用一个独立的小型 LLM，而 MTP、EAGLE-3、DFlash 和 DSpark 则使用辅助层，并结合来自目标模型的信息来预测 token。后缀（suffix）和 n-gram 方法不依赖模型，而是复用 token 流中已经出现过的模式。

表 1 比较了主要的草稿方法，包括它们如何生成 token、推理时的内存占用以及草稿开销。

表 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$ 时达到 6 的 $AL$，而 4B 草稿模型的 AL 则超过 5。随着 $D$ 增大，MTP 和 DFlash 的 AL 趋于平稳。N-gram 在该工作负载上的接受率较低，更适合具有重复 token 模式的工作负载。

更高的 $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$ 和草稿机制。

1. 增大推测解码的草稿长度，把 GEMM 推入计算受限区域，同时不增加 KV 缓存容量压力。
2. 当注意力主导解码时间时，选择 $D = \frac{128}{G} – 1$ 作为起点。
3. 对于较大的草稿长度，优先选择使 $G \times (1 + D)$ 为 128 的倍数的取值，以与注意力内核的分块大小对齐。
4. 在极低延迟下，使用快速的草稿机制，并且只有当 $AL$ 的增益足以抵偿新增的草稿成本时，才增大 $D$。
5. 通过在 $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 的帮助。