用投机解码协同设计 AI 模型：NVIDIA 提出的五个草稿长度与机制选择指南

NVIDIA Developer Blog 发表技术文章，这是 NVIDIA 关于 AI 模型协同设计系列文章的第三篇。文章聚焦如何使用投机解码（speculative decoding）在保持准确率的同时加速 LLM 推理，并提出了选择草稿长度（draft length）和草稿机制（draft mechanism）的五个指南。核心思想是：投机解码通过小型草稿模型预测多个 token，再用目标模型并行验证，将解码推向更计算密集（compute-bound）的区域，从而在不增加并发的情况下提升吞吐量和交互性。本文基于 NVIDIA 技术博客，系统解读投机解码的原理、指南和工程实践。

背景：LLM 推理的瓶颈

自回归解码的挑战

LLM 推理的核心瓶颈是自回归解码：

• 每次迭代生成一个 token
• 顺序执行，无法并行
• 内存访问密集（memory-bound）
• 硬件利用率低
• 吞吐量受限

吞吐量与交互性的权衡

LLM 推理面临吞吐量-交互性帕累托前沿（Pareto frontier）：

• 提升吞吐量需要更大批处理
• 更大批处理增加延迟
• 交互性（低延迟）与吞吐量（高批量）难以兼得
• 传统方法通过提高并发来改善

投机解码的路线

投机解码通过预测多个 token 改变这一局面：

• 小型草稿模型先预测几个可能的下一个 token
• 目标模型用一次前向传递并行验证
• 减少总解码迭代次数
• 提高目标模型的算术强度
• 不需要更高并发
• 将帕累托前沿向右上推移

投机解码原理

核心流程

投机解码的 draft-then-verify 流程：

1. 小型草稿模型预测 D 个可能的下一 token
2. 目标模型在单次传递中并行验证这些 token
3. 目标模型按顺序接受提议的 token，直到遇到第一个不匹配
4. 下一轮预测从该位置恢复

关键术语

• 草稿长度（D）：每次目标迭代提议的 token 数量
• 接受长度（AL）：每次目标迭代产生的（被接受的）token 数量
• AL 范围从 1 到 (1+D)：因为目标模型总能额外产生一个真实 token

输出一致性

• 只有被目标模型接受的 token 被保留
• 除非故意放宽接受标准，投机解码产生与标准解码相同的输出序列
• 保持准确性

加速比公式

投机解码的加速比可以量化为：

speedup = T_verif(B) × AL / (T_verif(B × (1+D)) + T_draft(B, D))

其中 B 是批大小，T_verif(x) 是目标模型验证 x 个 token 的时间，T_draft(b, y) 是草稿模型生成 y 个 token 草稿的时间。

最大化加速比需要找到最优的 (D, AL, T_draft) 组合。

五个设计指南

指南 1：增加草稿长度以推动 GEMM 进入计算密集区

• 投机使每个目标线性层的 GEMM-M 从 M 增长到 M × (1+D)
• 更高的草稿长度让 GEMM 在更低的有效批大小下达到峰值性能
• 示例：D=7 时，与 D=0 相比只需 1/8 的批大小即可变为计算密集
• 验证时计算量随 (1+D) 扩展，但内存访问保持不变
• 对 MoE 模型：专家变得更稀疏、长上下文增加 KV 容量压力时，每个专家的有效并发减少，更大的草稿长度在帕累托前沿上更有吸引力
• 核心：增加草稿长度将 GEMM 推入计算密集区，而不增加 KV 缓存容量压力

指南 2：注意力主导解码时选择 D = 128/G - 1

• 对推理和 agent 工作负载，注意力往往在吞吐量区域主导执行时间
• 解码注意力的算术强度约为 2 × G（G 是共享一个 KV 头的查询头数）
• 投机将其增加到 2 × G × (1+D)，因为被推测的 token 复用相同 KV 缓存
• 有效注意力 GEMM-M 为 G × (1+D)
• 当前 GPU 上，注意力内核在 GEMM-M=128 时达到良好硬件利用率
• 因此 D = 128/G - 1 是最优草稿长度
• 注意：G=32 的变体在更小的 D 值即达到吞吐量饱和；超过饱和点后，注意力不再受 DRAM 带宽限制，运行时间随 D 线性增长，而 AL 随 D 亚线性增长，继续增加 D 会拖慢注意力密集型工作负载

指南 3：选择避免 tile 利用不足的 D 值

• 注意力运行时间也取决于 tile 大小
• 当 G × (1+D) 跨越 128 的倍数（软件 tile 大小）时，运行时间阶梯式增加
• 如果 G × (1+D) 落在两个 tile 边界之间，最后一个 tile 只部分利用，但成本约等于完整 tile
• 指南：如果选择 D > 128/G - 1，优先选择 G × (1+D) 是 128 倍数的值，避免 tile 利用不足
• 指南 1 与指南 2/3 的相对权重取决于首选工作点上 FFN 与注意力在运行时间中的占比
• 通信也随 D 增加验证成本，但计算和通信重叠可以缓解开销

指南 4：极低延迟下仅在收益足够时增加 D

• 向帕累托曲线右端移动时，批大小 B 变得非常小
• 固定内核设置和后处理开销主导计算和通信内核
• 这些固定成本不随验证 token 数量显著增长
• 在极低延迟区域，顺序内核启动次数决定工作负载延迟
• 内核启动次数与层数线性相关，对类似目标结构自回归草稿模型：

  speedup = AL / (1 + ρD)
  

  其中 ρ = L_draft / L_target 是草稿深度比，草稿开销 O_d = ρD
• 核心：只有 AL 的增益足以抵消草稿开销时，增加 D 才有帮助
• MoE 模型中，模型分片策略和 Grouped GEMM 等高效内核可以保持开销较低，只要接受率高，更大草稿长度在低延迟区域有帮助

指南 5：选择草稿机制

决定草稿长度后，需要选择生成 token 的机制：

不同机制的权衡：

• 训练成本
• 参数开销
• 运行时成本
• 服务时内存
• 草稿开销

NVIDIA 提供 ready-to-run 训练示例：

• EAGLE-3、DFlash、DSpark
• 在 NVIDIA/Model-Optimizer 仓库中
• 包括在 Nemotron 3.5 Lightning 上演示的微调和量化工作流

工程实践

评测基准

• SPEED-Bench：用真实提示词和多样任务域基准测试接受长度
• 覆盖编码和摘要等任务域
• 评估草稿开销

工作流建议

1. 在 NVIDIA/Model-Optimizer 中探索训练示例和量化工作流
2. 用 SPEED-Bench 和真实提示词基准测试接受长度
3. 用 NVIDIA TensorRT LLM 量化草稿开销，实现高性能推理
4. 根据工作负载特点（推理 vs 注意力占比）应用五个指南

注意事项

• 本内容由 NVIDIA Nemotron 提供支持
• AI 生成内容可能不完整地总结信息
• 重要信息需要验证

总结

NVIDIA 的投机解码协同设计指南为 LLM 推理优化提供了系统的方法论。投机解码通过小型草稿模型预测多个 token、目标模型并行验证，减少了总解码迭代次数并提高算术强度，将吞吐量-交互性帕累托前沿向右上推移，且不改变输出序列（除非故意放宽接受标准）。五个指南构成完整的设计框架：指南 1（增加草稿长度将 GEMM 推入计算密集区，MoE 稀疏化和长上下文场景更有吸引力）；指南 2（注意力主导时选择 D = 128/G - 1，基于注意力内核在 GEMM-M=128 达到良好利用率的经验）；指南 3（避免 tile 利用不足，选择 G × (1+D) 为 128 倍数的 D）；指南 4（极低延迟区域，只有 AL 增益足以抵消草稿开销 ρD 时才增加 D）；指南 5（选择草稿机制，从外部草稿模型、EAGLE-3、MTP、DFlash/DSpark 到后缀/n-gram，各有训练、参数、运行时成本权衡）。工程实践包括使用 SPEED-Bench 评测接受长度、探索 NVIDIA/Model-Optimizer 的训练示例和量化工作流、用 TensorRT LLM 量化草稿开销。这一系列指南将投机解码从"调参技巧"提升为"可推导的设计方法"，为在不同硬件和工作负载下优化 LLM 推理提供了理论依据和实操路径。

来源：https://developer.nvidia.com/blog/co-designing-ai-models-using-speculative-decoding-for-faster-llm-inference/