使用投机解码协同设计AI模型,加速LLM推理
本文是 AI 模型协同设计系列的第三篇。文章探讨了如何利用 投机解码(speculative decoding)在保持精度的前提下加速 LLM 推理,并在帕累托前沿上给出了跨不同草稿长度和草稿机制的五条选型指南。
如需了解模型设计选择如何在牺牲精度的情况下同时影响吞吐量和交互性,请参阅 AI 模型协同设计:硬件友好的 LLM 设计(第一篇)。如需了解分组大小(每个 KV head 的查询 head 数)、head 维度以及序列长度如何塑造密集注意力(dense attention)的性能,请参阅 协同设计 AI 模型注意力以实现快速交互长上下文推理(第二篇)。
什么是投机解码?
投机解码是一种通过每次迭代预测多个 token 来加速 LLM 推理自回归解码阶段的 techniques。小型草稿模型首先预测若干可能的后续 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 展示了在草稿长度各异时,针对典型 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 前沿上更具吸引力。
准则 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\) 很可能会拖慢注意力密集型工作负载。
准则 2:当注意力占据解码阶段主要时间时,选择 \(D = \frac{128}{G} – 1\)。
注意力的运行时间还取决于 tile 大小。Figure 5 表明,当 \(G \times (1 + D)\) 跨越 128 的倍数时(即被测试注意力 kernel 的软件 tile 大小),运行时间呈阶梯式增长。如果 \(G \times (1 + D)\) 落在两个 tile 边界之间,最后一个 tile 仅被部分利用,但开销仍接近完整 tile。
准则 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-3 | MTP | DFlash | DSpark | 后缀 / n-gram | |
| 草稿架构 | 小型 LLM | 解码层 + 线性投影 | 解码层 + 线性投影 | 解码层 + 线性 KV 融合 | 解码层 + 线性 KV 融合 + 轻量级马尔可夫头 | 字符串匹配,无模型 |
| 输入 | Token ID | 上一 token 嵌入 + 目标隐状态 | 目标最终隐状态 | 融合的 target 隐状态用作草稿 KV | 融合的 target 隐状态用作草稿 KV | Token 流 |
| 生成方式 | \(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\) 低于 MTP | GPU 上大模型 | 小模型,batch size=1 | 小模型,batch size=1 | 高重复性 workload |
为量化\(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 模式的场景。
要在 \(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\) 和草稿机制。
- 增大推测解码的草稿长度,将GEMM运算推入计算瓶颈区,同时不增加KV cache容量压力。
- 当注意力机制主导解码时间时,以 \(D = \frac{128}{G} – 1\) 作为起点。
- 对于较大的草稿长度,优先选择使 \(G \times (1 + D)\) 为128倍数的值,以对齐attention kernel的tile尺寸。
- 在极低延迟场景下,采用快速草稿机制,仅在AL的提升足以抵消额外草稿成本时才增大 \(D\)。
- 根据工作负载和硬件,权衡AL与草稿开销来选择草稿机制。在真实服务条件下对两者进行基准测试,并综合考虑训练与部署成本。
对于后训练草稿模型,EAGLE-3、DFlash和DSpark的即用训练示例已提供在 NVIDIA/Model-Optimizer 中。以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的帮助。
