← 文章 / AI技术
NVIDIA 开发者博客 4小时前 · 2026-09-03 02:08:15 · 3 阅读

使用投机解码协同设计AI模型,加速LLM推理

本文是 AI 模型协同设计系列的第三篇。文章探讨了如何利用 投机解码(speculative decoding)在保持精度的前提下加速 LLM 推理,并在帕累托前沿上给出了跨不同草稿长度和草稿机制的五条选型指南。

如需了解模型设计选择如何在牺牲精度的情况下同时影响吞吐量和交互性,请参阅 AI 模型协同设计:硬件友好的 LLM 设计(第一篇)。如需了解分组大小(每个 KV head 的查询 head 数)、head 维度以及序列长度如何塑造密集注意力(dense attention)的性能,请参阅 协同设计 AI 模型注意力以实现快速交互长上下文推理(第二篇)。

什么是投机解码?

投机解码是一种通过每次迭代预测多个 token 来加速 LLM 推理自回归解码阶段的 techniques。小型草稿模型首先预测若干可能的后续 token,然后通过一次前向传播利用大型目标模型对这些 token 进行并行验证。

图 1. 推测解码使 LLM 解码向更计算密集的区域偏移,推动吞吐量-交互性帕累托前沿向右上方移动。沿着曲线向右,交互性提升、批次大小减小,GEMM 在运行时间中的占比超过注意力计算。

该方法在不增加并发量的前提下,减少目标模型的解码迭代次数,同时提升其算术强度。目标模型按序接受提议的 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)\) 三元组。
推测解码流程:展示预填充、草稿生成、验证、接受的 token 以及第一个被拒绝的 token。
图 2. 推测解码的草稿优先后验证流程。较小的草稿模型提出 \(D\) 个 token,目标模型在一次前向传播中统一验证

选择最优草稿长度

为简化讨论,若忽略草稿模型的延迟,推测解码在以下条件成立时可提供加速:

\(\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 展示了在草稿长度各异时,针对典型 expert GEMM 尺寸 6144 × 6144,TFLOPS 如何随批处理大小变化。显然更高的草稿长度使 GEMM 能在更低的实际批处理大小下达到峰值性能。

值得注意的是,当 $D=7$ 时,与 $D=0$ 相比,只需 1/8 的 batch size 就会进入计算瓶颈区。随着 [mixture-of-experts (MoE)](https://www.nvidia.com/en-us/glossary/mixture-of-experts/) 模型变得越来越稀疏,长上下文负载对 KV capacity 的压力持续增加,每个 expert 的有效并发数随之下降,使得更大的 draft length 在 Pareto 前沿上更具吸引力。 随着 D 增大,GEMM 吞吐量随有效 batch size 增长,且在较小的 batch size 下即可达到高利用率。
图 3. 不同 draft length 下,典型 expert GEMM(GEMM-N = 6,144,GEMM-K = 6,144)的归一化 TFLOPS(以 $D=0$、batch size=1 为基准)与 batch size 的关系

准则 1:增大推测解码的 draft length,在不增加 KV cache 容量压力的前提下,将 GEMM 推向计算瓶颈区。

Draft length 与 attention 性能

对于推理和 agentic 工作负载,attention 通常在面向吞吐量的区域主导执行时间。Decode attention 的算术强度约为 $2 \times G$,其中 $G$ 是共享同一个 KV head 的 query head 数量,详见 Co-Designing AI Model Attention for Fast, Interactive Long-Context Inference

推测解码将其提升至 $2 \times G \times (1 + D)$,因为推测的 token 复用了相同的 KV cache。此时有效的 attention GEMM-$M$ 为 $G \times (1 + D)$。在当前 GPU 设备上,attention kernel 在 GEMM-$M=128$ 时即可实现良好的硬件利用率,因此 $D = \frac{128}{G} – 1$ 即为最优 draft length。

Figure 4 展示了当 KV 序列长度分别为 32K 和 128K、G = 8 和 G = 32 时,归一化注意力吞吐量随 \(D\) 增大的变化情况。\(G = 32\) 的版本在更小的 \(D\) 值处便达到了吞吐量饱和。越过饱和点后,注意力计算不再受 DRAM 带宽限制,其运行时间随 \(D\) 增长。由于 \(AL\) 随 \(D\) 亚线性增长,继续增大 \(D\) 很可能会拖慢注意力密集型工作负载。

Attention TFLOPs versus D for G=8 and G=32, with G=32 reaching high utilization sooner.
Figure 4. 归一化注意力 TFLOPS(相对 \(D=0\))与 \(D\) 的关系,在 \(G=8\)\(G=32\) 下的对比。更大的 \(G\) 在更短的草稿长度下即可达到高利用率

准则 2:当注意力占据解码阶段主要时间时,选择 \(D = \frac{128}{G} – 1\)。

注意力的运行时间还取决于 tile 大小。Figure 5 表明,当 \(G \times (1 + D)\) 跨越 128 的倍数时(即被测试注意力 kernel 的软件 tile 大小),运行时间呈阶梯式增长。如果 \(G \times (1 + D)\) 落在两个 tile 边界之间,最后一个 tile 仅被部分利用,但开销仍接近完整 tile。

随着 G × (1+D) 跨越 128 的倍数,Attention 运行时呈阶梯式上升。
图 5. 在 \(G = 8\)\(G = 32\) 两种情况下,规范化 Attention 运行时(相对于 \(D=0\))与 \(D\) 的关系。随着 \(G \times (1 + D)\) 跨越 128 的倍数,运行时呈阶梯式增长

准则 3:若选定 \(D > \frac{128}{G} – 1\),建议取使 \(G \times (1 + D)\) 为 128 倍数的 \(D\) 值,以避免 tile 未充分利用。

准则 1 相对于准则 2 和 3 的重要性,取决于在优选工作点下 FFN 运行时在整个运行时中所占的比例。随着 D 增大,通信开销也会相应增加,尽管计算与通信的并行可以缓解这部分额外开销。

Pareto 曲线最右侧的 draft 长度

越靠近 Pareto 曲线最右侧,\(B\) 就越小。此时,无论是计算还是通信 kernel,固定的 kernel 设置和后处理开销都占主导地位。这些固定成本并不会随验证 token 数量的增加而显著增长,因此验证开销基本不随 draft token 数量变化。

虽然 MoE 模型会随 draft 长度增加而激活更多专家,但结合模型分片策略和 Grouped GEMM 等高效 kernel,可以将这部分开销控制在较低水平。因此,只要接受率保持较高,较大的 draft 长度在低延迟场景下反而更有利。

在极低延迟场景下,顺序执行的 kernel launch 次数决定了整体 workload 延迟。由于 kernel launch 次数随层数线性增长,对于层结构与目标模型相近的自回归 draft 模型,加速比可近似为:

\(\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\)

换言之,增大 \(D\) 只有在 \(AL\) 的提升足以抵消草稿开销时才有效。

指南 4:在极低延迟场景下,仅当 \(AL\) 的提升足以弥补额外草稿成本时,才增大 \(D\)。

选择草稿机制

确定 \(D\) 即决定了推测的 token 数量。接下来需决定如何生成这些 token 以实现最大加速。

多年来提出了多种技术方案,各自在训练、参数和运行时成本上有所不同。外部草稿方法采用独立的小型 LLM,而 MTP、EAGLE-3、DFlash 和 DSpark 则利用辅助层结合目标模型的信息来预测 token。后缀和 n-gram 方法无需模型,而是直接复用 token 流中已出现的模式。

表 1 对比了主要草稿方法,包括它们的 token 生成方式、服务时内存占用以及草稿开销。

 外部草稿模型EAGLE-3MTPDFlashDSpark后缀 / n-gram
草稿架构小型 LLM解码层 + 线性投影解码层 + 线性投影解码层 + 线性 KV 融合解码层 + 线性 KV 融合 + 轻量级马尔可夫头字符串匹配,无模型
输入Token ID上一 token 嵌入 + 目标隐状态目标最终隐状态融合的 target 隐状态用作草稿 KV融合的 target 隐状态用作草稿 KVToken 流
生成方式\(D\) 次自回归前向传播模块运行 \(D\) 次模块运行 \(D\) 次一次并行步骤生成 \(D\) 个 token一次并行步骤生成 \(D\) 个 token,后接轻量级串行校正
一次查表
发布方式独立模型发布独立的后训练头作为基础模型 checkpoint 的一部分独立的后训练头独立的后训练头与目标模型无关
训练成本从头训练:1T-10T+ tokens
蒸馏:100B-400B
适配:10M-1B
目标模型训练后额外 1-10B tokens通常与目标模型在预训练阶段同步训练目标模型训练后额外 1-10B tokens目标模型训练后额外 1-10B tokens
激活内存开销草稿模型权重 + 完整 KV cache权重 + 小型 KV cache权重 + 小型 KV cache权重 + 从目标模型构建的小型 KV cache权重 + 从目标模型构建的小型 KV cache
每\(D\)个token的投机成本完整的草稿模型前向计算1 层,串行1 层,串行约 5 层,一次性约 5 层一次性 + 轻量级 Markov 头,串行O(1) 查表
最佳适用场景LPU 和 LPU+GPU
不推荐在 GPU 上使用
不推荐——\(AL\) 低于 MTPGPU 上大模型小模型,batch size=1小模型,batch size=1高重复性 workload
表1. 常见草稿机制在输入方式、生成方法、训练成本、推理时内存和投机成本方面的对比

为量化\(AL\)与\(O_d\)之间的权衡,首先观察\(AL\)如何随\(D\)缩放。

图6展示了在SPEED-Bench上,随着D增加时AL的变化情况,目标模型为 Qwen 3.5 122B A10B。SPEED-Bench 是 NVIDIA 开发的投机解码基准,旨在反映真实的生产 workload,涵盖编程、摘要等多个任务领域,并包含不同输入序列长度的多种划分。我们推荐用 SPEED-Bench 进行\(AL\)比较。

在 32K 数据切分中,Qwen 3.5 35B A3B 在 $D = 9$ 时达到 AL = 6,而 4B 草稿模型 AL 高于 5。随着 $D$ 增大,MTP 和 DFlash 的 AL 趋于平稳。N-gram 在该负载上的接受率较低,更适合具有重复 token 模式的场景。
Line graph showing AL versus D for several draft mechanisms. External drafts keep improving, while MTP and DFlash level off earlier.
图 6. $AL$$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\) 和草稿延迟之间取得平衡,准确的基准测试至关重要。在推测解码的语境下,这意味着需要使用真实 prompt 进行测试,并覆盖广泛的任务领域。我们建议使用 SPEED-Bench 测量 \(AL\),并用 NVIDIA TensorRT LLM 等高性能推理框架来量化草稿开销。

除了推理性能,草稿训练的范围和成本同样需要考量。MTP 必须与目标模型联合训练,而 EAGLE、DFlash 和 DSpark 则可以直接添加到最终模型检查点中。同样,在外部草稿模型的构建上也存在一系列选项:从头训练或从目标模型蒸馏能获得最高的 \(AL\),而通过跨模型适配技术将现有模型适配到目标,则能在牺牲一定 \(AL\) 的前提下显著降低训练成本。

微调目标模型会改变其输出分布和隐藏表示。因此,专为某个目标检查点训练的草稿模型,即使目标模型性能提升也可能不再被接受。更改目标后,应在代表性工作负载上重新测量 \(AL\),所需的具体适配程度取决于草稿机制的类型。

MTP 作为目标模型的一部分,应在微调期间继续训练,或在专门的后续阶段重新对齐。像 EAGLE-3、DFlash 和 DSpark 这类依附于目标的草稿机制会使用目标的隐藏状态,需要针对更新后的检查点进行额外训练以完成适配。

外部草稿模型不依赖目标的隐藏状态,但仍需近似目标的输出分布,可能仍需微调或蒸馏。后缀匹配和 n-gram 方法没有可学习的草稿器,因此无需重新训练,不过其效果仍取决于部署工作负载中的重复程度。

准则 5: 结合 \(AL\)、草稿延迟以及训练和部署成本,选择能为你的工作负载和硬件带来最佳解码加速比的草稿机制。

开始协同设计推测解码方案

将以下五条准则作为检查清单,在 Pareto 前沿上选择 \(D\) 和草稿机制。

  1. 增大推测解码的草稿长度,将GEMM运算推入计算瓶颈区,同时不增加KV cache容量压力。
  2. 当注意力机制主导解码时间时,以 \(D = \frac{128}{G} – 1\) 作为起点。
  3. 对于较大的草稿长度,优先选择使 \(G \times (1 + D)\) 为128倍数的值,以对齐attention kernel的tile尺寸。
  4. 在极低延迟场景下,采用快速草稿机制,仅在AL的提升足以抵消额外草稿成本时才增大 \(D\)。
  5. 根据工作负载和硬件,权衡AL与草稿开销来选择草稿机制。在真实服务条件下对两者进行基准测试,并综合考虑训练与部署成本。

对于后训练草稿模型,EAGLE-3、DFlash和DSpark的即用训练示例已提供在 NVIDIA/Model-Optimizer 中。以NVIDIA Nemotron 3.5 Lightning为例:先 微调DSpark,再量化为 FP8NVFP4。以这些示例为起点,然后根据你自己的模型、工作负载和硬件验证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的帮助。

Line graph showing AL versus D for several draft mechanisms. External drafts keep improving, while MTP and DFlash level off earlier.
原始来源: NVIDIA 开发者博客

评论 (0)