title: "NVIDIA 利用推测解码协同设计 AI 模型,加速 LLM 推理"
source_url: "https://developer.nvidia.com/blog/co-designing-ai-models-using-speculative-decoding-for-faster-llm-inference/"
author: "NVIDIA Technical Blog"
excerpt: "公共早报 本技术指南探讨如何通过推测解码优化 LLM 推理速度,提供五项工程准则,帮助在吞吐量与交互性的帕累托前沿上选择最优草稿长度与机制。"
这篇文章是关于 AI 模型协同设计系列的第三篇。它探讨了如何使用推测解码在保持准确性的同时加速 LLM 推理,并提供了在帕累托前沿上选择草稿长度和草稿机制的五项准则。
关于模型设计选择如何影响吞吐量和交互性而不牺牲准确性的讨论,请参阅 AI 模型协同设计:硬件友好的 LLM 设计(第一部分)。关于组大小(每个 KV 头的查询头数)、头维度和序列长度如何影响密集注意力的性能,请参阅协同设计 AI 模型注意力以实现快速、交互式长上下文推理(第二部分)。
什么是推测解码? {#what_is_speculative_decoding}
推测解码是一种通过预测每个迭代的多个 token 来加速 LLM 推理自回归解码阶段的技术。首先,一个小的草稿模型预测几个可能的下一个 token。然后,这些 token 通过一个 pass 对较大的目标模型进行并行验证。
图 1. 推测解码将 LLM 解码推向更计算受限的区域,将吞吐量和交互性的帕累托前沿向上和向右移动。随着交互性向右增加,批大小减小,GEMM 在运行时中所占份额比注意力更大
这种方法减少了解码迭代的总数,同时增加了目标模型的算术强度,而无需更高的并发性。目标模型按顺序接受 proposed tokens,直到遇到第一个不匹配。下一个预测周期然后从该位置恢复。因为只有被目标模型接受的 tokens 被保留,推测解码产生与标准解码相同的输出序列,除非 acceptance criteria 被有意放宽。图 2 说明了这种解码流程。
草稿长度((D))是每个目标迭代 proposed tokens 的数量。接受长度((AL))是每个目标迭代产生的 tokens 数量(即被接受的)。(AL) 范围从 (1) 到 (1 + D),因为目标模型总是可以除了接受的草稿 tokens 之外产生一个新的 ground-truth token。
推测解码的加速可以用目标模型顺序生成 (AL) 个 tokens 所花费的时间与并行验证 (D) 个 tokens 所花费的时间来量化,同时考虑生成 (D) 个 tokens 的延迟开销:
(\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) 个 tokens 的目标验证时间,(T_{\mathrm{draft}}(b, y)) 是为批大小 (b) 生成长度为 (y) 的草稿所需的时间。显然,要最大化加速,需要找到最优的 (\left(D, AL, T_{\mathrm{draft}}\right)) 元组。
图 2. 推测解码的「先草稿后验证」流程。一个较小的草稿模型提出 (D) 个 tokens,目标模型在一个 pass 中验证它们
选择最优草稿长度 {#selecting_the_optimal_draft_length}
为简单起见,如果我们忽略草稿模型延迟,当以下条件成立时,推测提供了加速:
(\frac{T_{\mathrm{verif}}\left(B \times (1 + D)\right)}{T_{\mathrm{verif}}(B)} \< AL)
在验证期间,计算量随 (1 + D) 缩放,但内存访问保持不变。因此,目标是增加 (D),直到 (T_{\mathrm{verif}}) 保持不变,通常直到验证从内存受限过渡到计算受限。显然,这个最优 (D) 值取决于 (B),并预计在帕累托前沿上有所不同。结合这一点考虑图 1。
草稿长度与线性层性能 {#draft_length_and_performance_of_linear_layers}
使用推测,每个目标线性层的 GEMM-(M) 从 (M) 增长到 (M \times (1 + D)),其中 (M) 是无推测时的 GEMM-(M)。图 3 显示了对于具有代表性的专家 GEMM 大小 6144 × 6144,在各种草稿长度下 teraflops per second 如何随批大小缩放。显然,更高的草稿长度允许 GEMM 在更低的有效批大小下达到峰值性能。
值得注意的是,当 (D=7) 时,与 (D=0) 相比,需要八分之一的批大小才能成为计算受限。随着混合专家(MoE) 模型变得更稀疏且长上下文工作负载增加 KV 容量压力,每个专家的有效并发性降低,这使得更大的草稿长度在整个帕累托前沿上具有吸引力。
图 3. 对于 GEMM-N = 6,144 和 GEMM-K = 6,144 的代表性专家 GEMM,不同草稿长度下相对于 (D=0)、批大小=1 的归一化 teraflops per second 与批大小的关系
准则 1:增加推测解码草稿长度,将 GEMM 推入计算受限区域,而不增加 KV 缓存容量压力。
草稿长度与注意力性能 {#draft_length_and_attention_performance}
对于推理和 agentic 工作负载,注意力往往在面向吞吐量的区域主导执行时间。解码注意力的算术强度约为 (2 \times G),其中 (G) 是共享一个 KV 头的查询头数量,如协同设计 AI 模型注意力以实现快速、交互式长上下文推理中所解释。
推测将其增加到 (2 \times G \times (1 + D)),因为推测的 tokens 重用相同的 KV 缓存。有效注意力 GEMM-(M) 为 (G \times (1 + D))。在当前 GPU 设备上,注意力 kernel 在 GEMM-(M=128) 时实现良好的硬件利用率,因此 (D = \frac{128}{G} -- 1) 是最优草稿长度。
图 4 显示了当 (G = 8) 和 (G = 32) 在 32K 和 128K KV 序列长度下随着 (D) 增加归一化注意力吞吐量的变化。(G = 32) 变体在更低的 (D) 值处达到吞吐量饱和。超过吞吐量饱和点后,注意力不再受 DRAM 带宽限制,其运行时间随 (D) 缩放。因为 (AL) 随 (D) 次线性增长,将 (D) 增加到超过此点可能会减慢注意力繁重的工作负载。
图 4.* 相对于 (D=0)* 的归一化注意力 teraflops per second 与 (G=8) *和 (G=32) 的 (D) 的关系。更大的 (G) 在更小的草稿长度处达到高利用率
准则 2:当注意力主导解码时间时,选择 (D = \frac{128}{G} -- 1)。
注意力运行时间也取决于 tile 大小。图 5 显示,当 (G \times (1 + D)) 越过 128 的倍数(benchmark 注意力 kernel 的软件 tile 大小)时,运行时间以 step 形式增加。如果 (G \times (1 + D)) 落在两个 tile 边界之间,最后一个 tile 仅部分利用但仍花费大约与完整 tile 相同的成本。
图 5. *相对于 *** (D=0)** 的归一化注意力运行时间与 (G = 8) 和 (G = 32) 的 (D) 的关系。当 *** (G \times (1 + D))*** 越过 128 的倍数时,运行时间以 step 形式增加**
准则 3:如果选择 (D \> \frac{128}{G} -- 1),则优先选择 (G \times (1 + D)) 是 128 倍数的值,以避免 tile 利用不足。
准则 1 相对于准则 2 和 3 的相对权重取决于在首选操作点上在 FFN 中与注意力相比所花费的运行时间百分比。通信也增加验证成本随着 D 增加,尽管计算和通信重叠可以缓解开销。
帕累托曲线最右端的草稿长度 {#draft_length_at_the_far_right_of_the_pareto_curve}
走向帕累托曲线最右端,(B) 变得非常小。在这里,固定 kernel 设置和后处理开销在计算和通信 kernel 上都占主导地位。这些固定成本不会随着已验证 tokens 数量的增加而显著增长,使验证开销在很大程度上与草稿 tokens 数量无关。
虽然 MoE 确实随着草稿长度增加看到更多专家被激活,但模型分片策略和高效 kernel(如分组 GEMM)的组合可以将此开销保持在低位。因此,在接受度保持较高的前提下,较大的草稿长度可以帮助低延迟区域。
在非常低的延迟下,顺序 kernel launches 的数量决定了工作负载延迟。因为 kernel launches 随层数线性缩放,对于具有与目标相似层结构的自回归草稿模型,加速可以近似为:
(\mathrm{speedup} = \frac{L_{\mathrm{target}} \times \mathrm{AL}}{L_{\mathrm{target}} + D \times L_{\mathrm{draft}}})
其中 (L_{\mathrm{target}}) 和 (L_{\mathrm{draft}}) 分别是目标和草稿模型中的层数。定义一个恒定的草稿深度比,(\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)。
选择草稿机制 {#selecting_a_draft_mechanism}
选择 (D) 告诉你应该推测多少个 tokens。接下来,决定如何生成这些 tokens 以最大化加速。
多年来已经提出了多种技术,各有不同的训练、参数和运行时间权衡。外部草稿建议使用小的独立 LLM,而 MTP、EAGLE-3、DFlash 和 DSpark 使用辅助层结合来自目标模型的信息来预测 tokens。Suffix 和 n-gram 方法是无模型的,而是重用 token 流中已经看到的模式。
表 1 比较了关键的草稿方法,包括它们如何生成 tokens、服务时的内存和草稿开销。
|-----------------------------------------|-------------------------------------------------------------------|-------------------------------------------------|----------------------------------------------------|-----------------------------------------------|------------------------------------------------------------------------------------------------|--------------------------------| | | 外部草稿模型 | EAGLE-3 | MTP | DFlash | DSpark | Suffix / n-gram | | 草稿架构 | 小型 LLM | 解码器层 + 线性投影 | 解码器层 + 线性投影 | 解码器层 + 线性 KV 融合 | 解码器层 + 线性 KV 融合 + 轻量级马尔可夫头 | 字符串匹配,无模型 | | 输入 | Token IDs | 先前 token embedding + 目标隐藏状态 | 目标最终隐藏状态 | 融合目标隐藏状态用作草稿 KV | 融合目标隐藏状态用作草稿 KV | Token 流 | | 生成方法 | (D) 次自回归前向传播 | 模块运行 (D) 次 | 模块运行 (D) 次 | 在一个并行步骤中生成 (D) 个 tokens | 在一个并行步骤中生成 (D) 个 tokens,然后进行轻量级顺序纠正 | 一次查找 | | 发布方法 | 独立模型发布 | 独立后训练头 | 基模型检查点的一部分 | 独立后训练头 | 独立后训练头 | 独立于目标 | | 训练成本 | 从头开始:1T-10T+ tokens 蒸馏:100B-400B 适配:10M-1B | 目标训练后额外的 1-10B tokens | 通常与目标一起在预训练期间训练 | 目标训练后额外的 1-10B tokens | 目标训练后额外的 1-10B tokens | 无 | | 激活内存成本 | 草稿权重 + 完整 KV 缓存 | 权重 + 小 KV 缓存 | 权重 + 小 KV 缓存 | 从目标构建的权重 + 小 KV 缓存 | 从目标构建的权重 + 小 KV 缓存 | 无 | | 每个 (D) tokens 的推测成本 | 完整草稿模型前向传播 | 1 层,串行 | 1 层,串行 | ~5 层,一次 | ~5 层一次 + 轻量级马尔可夫头,串行 | O(1) 查找 | | 最适合 | LPU 和 LPU+GPU 不推荐用于 GPU | 不推荐——比 MTP 更低的 (AL) | GPU 上较大模型的最佳选择 | 较小模型,批大小为 1 | 较小模型,批大小为 1 | 高重复性的工作负载 |
表 1. 常见草稿机制在输入、生成方法、训练成本、服务时内存和草稿成本方面的比较
为量化 (AL) 和 (O_d) 之间的权衡,首先观察 (AL) 如何随 (D) 缩放。
图 6 显示了在 SPEED-Bench 上,当使用 Qwen 3.5 122B A10B 作为目标时,(AL) 如何随着 (D) 增加而变化。SPEED-Bench 是 NVIDIA 开发的一个推测解码基准测试,旨在代表真实的生产工作负载。它涵盖多个任务域,如编码和摘要,并有各种不同输入序列长度的 split。我们推荐使用 SPEED-Bench 进行 (AL) 比较。
在 32K split 上,Qwen 3.5 35B A3B 在 (D = 9) 时达到 (AL) 为 6,而 4B 草稿达到 (AL) 高于 5。MTP 和 DFlash (AL) 随着 (D) 增长而趋于平稳。N-gram 在此工作负载上接受度较低,更适合具有重复 token 模式的工作负载。
图 6. 在 SPEED-Bench 吞吐量 32K split 上,以 Qwen 3.5 122B A10B 为目标时,不同草稿机制的 (AL) 与 (D) 的关系
更高的 (AL) 不等于更高的加速。你还需要考虑生成草稿的成本是多少。
如图 6 所示,所有外部草稿在 (D \> 3) 时比其他人达到更高的 (AL)。Qwen 3.5 122B MTP 要小得多,总参数为 2.5B,活跃参数少于 150M。在 (D = 3) 时,较大的外部草稿可能不值得增加的成本。在 (D = 21) 时,假设有高效的 serving 策略,其更高的 AL 可能使该成本值得。
同时,DFlash (AL) 相对较快地趋于平稳。然而,DFlash 和 DSpark 都以并行方式生成 (D) 个草稿 tokens,降低了 (O_d)。作为比较,在最小延迟和 (D=11) 时,单层 MTP head 需要 11 步,导致 (O_d{\mathrm{MTP}} = \frac{11}{L_{\mathrm{target}}}),而五层 DFlash head 在一次 pass 中生成草稿,得到 (O_d{\mathrm{DFlash}} = \frac{5}{L_{\mathrm{target}}})。
对于具有许多层的大型目标模型,两个开销都可以忽略不计,但随着 (L_{\mathrm{target}}) 减小,草稿开销变得显著。因此,对于低延迟的小模型,因此,即使 (AL) 较低,DFlash 或 DSpark 也可能是最优选择。
要在 (AL) 和草稿延迟之间导航权衡,必须对两者进行准确基准测试。在推测解码的背景下,这意味着使用真实 prompt 进行基准测试并覆盖广泛的任务域。我们推荐使用 SPEED-Bench 进行 (AL) 测量,并使用 NVIDIA TensorRT LLM 等高性能推理框架来量化草稿开销。
除了推理性能,草稿训练范围和成本也值得考虑。虽然 MTP 必须与目标模型共同训练,但 EAGLE、DFlash 和 DSpark 可以添加到最终模型检查点。同样,外部草稿模型创建也存在一系列选项。从头开始训练和从目标蒸馏产生最高的 (AL),而通过跨模型适配技术将现有模型适配到所需目标可以显著降低训练成本代价是 (AL) 降低。
微调目标模型可以改变其输出分布和隐藏表示。因此,为特定目标检查点训练的 learned drafters 即使在目标模型改进时也可能失去接受度。更改目标后,在代表性工作负载上重新测量 (AL)。所需的适配取决于草稿机制。
MTP 是目标模型的一部分,应在微调期间继续训练或在专用后续阶段重新对齐。附加在目标上的草稿(如 EAGLE-3、DFlash 和 DSpark)使用目标的隐藏状态,将需要使用额外训练来适应更新后的检查点。
外部草稿模型不消耗目标隐藏状态,但它们仍然近似目标的输出分布,可能需要微调或蒸馏。Suffix 和 n-gram 方法没有学习的 drafter,因此不需要重新训练,尽管它们的效果仍然取决于部署工作负载中出现多少重复。
准则 5:选择能在你的工作负载和硬件上获得最佳解码加速的草稿机制,同时考虑 (AL)、草稿延迟以及训练和部署成本。
开始协同设计推测解码
使用下面总结的五项准则作为在整个帕累托前沿上选择 (D) 和草稿机制的检查清单。
增加推测解码草稿长度,将 GEMM 推入计算受限区域,而不增加 KV 缓存容量压力。
当注意力主导解码时间时,选择 (D = \frac{128}{G} -- 1) 作为起点。
对于更大的草稿长度,优先选择 (G \times (1 + D)) 是 128 倍数的值,以与注意力 kernel tile 大小对齐。
在非常低的延迟下,使用快速草稿机制,仅当 (AL) 的增益证明增加的草稿成本合理时,才增加 (D)。
通过平衡工作负载和硬件的 (AL) 和草稿开销来选择草稿机制。在真实 serving 条件下对两者进行基准测试,并考虑训练和部署成本。
对于后训练草稿,在 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、Taylor Abramovich、Yoav Miron、Rabeeh Karimi Mahabadi、Roger Waleffe、Udi Karpas、Ran Zilberstein、Brian Pharris 和 Eduardo Alvarez 的帮助。