利用推测解码加速大语言模型推理的协同设计方法

本文是AI模型协同设计系列的第三篇，探讨如何在保持准确性的同时，利用推测解码技术加速大语言模型推理，并提供了五条选择草稿长度和草稿机制的指导原则。

什么是推测解码

推测解码是一种通过每次迭代预测多个Token来加速大语言模型自回归解码阶段的技术。首先由一个小型草稿模型预测若干可能的下一个Token，再通过目标大模型的单次前向传播并行验证这些Token。

这种方法在不提高并发度的前提下，减少了解码迭代总次数，同时提升了目标模型的算术强度。目标模型按顺序接受所提议的Token，直到遇到第一个不匹配项，然后从该位置重新开始预测。由于只保留目标模型接受的Token，推测解码在不刻意放宽接受标准的情况下，输出序列与标准解码完全一致。图2展示了这一解码流程。

草稿长度（D）是每次目标模型迭代中提议的Token数量，接受长度（AL）是每次迭代实际产出（即被接受）的Token数量。AL的范围为1到（1+D），因为目标模型在接受草稿Token的同时，始终能额外生成一个新的真值Token。

推测解码的加速比可以量化为：目标模型顺序生成AL个Token所需时间，与并行验证D个Token所需时间（包含草稿生成延迟）之比：

$$\\mathrm{speedup} = \\frac{T_{\\mathrm{verif}}(B) \imes \\mathrm{AL}}{T_{\\mathrm{verif}}\\left(B \imes (1 + D)\\right) + T_{\\mathrm{draft}}(B, D)}$$

其中B为批次大小，$T_{\\mathrm{verif}}(x)$为x个Token的目标验证时间，$T_{\\mathrm{draft}}(b, y)$为批次大小b下生成长度y草稿所需时间。显然，要最大化加速比，需要找到最优的$(D, AL, T_{\\mathrm{draft}})$组合。

选择最优草稿长度

为简化分析，忽略草稿模型延迟，当以下条件成立时推测解码可带来加速：

在验证阶段，计算量随（1+D）线性增长，但内存访问量保持不变。因此目标是增大D直到$T_{\\mathrm{verif}}$保持不变，通常是验证从内存带宽瓶颈转变为计算瓶颈的临界点。D的最优值取决于B，在Pareto前沿的不同位置也会有所不同。

草稿长度与线性层性能

引入推测解码后，每个目标线性层GEMM的M维度从M增长为$M \imes (1 + D)$，其中M是无推测时的GEMM-M值。图3展示了在6144×6144的代表性专家GEMM规模下，不同草稿长度时每秒万亿次浮点运算随批次大小的变化。可以看出，更大的草稿长度能让GEMM在更低的有效批次大小下达到峰值性能。

值得注意的是，当D=7时，达到计算瓶颈所需的批次大小仅为D=0时的八分之一。随着混合专家（MoE）模型愈加稀疏、长上下文工作负载加剧KV容量压力，每个专家的有效并发度下降，使得更大的草稿长度在Pareto前沿具有更广泛的吸引力。

指导原则1：增大推测解码的草稿长度，将GEMM推入计算瓶颈区域，同时避免增加KV缓存容量压力。

草稿长度与注意力性能

对于推理和智能体工作负载，在吞吐量导向区域，注意力机制往往主导执行时间。解码注意力的算术强度约为 \imes G$，其中G为共享一个KV头的查询头数量（详见系列第二篇）。

由于推测Token会复用同一KV缓存，推测解码将算术强度提升至 \imes G \imes (1 + D)$，有效注意力GEMM-M为$G \imes (1 + D)$。在当前GPU设备上，注意力内核在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在注意力密集型工作负载中很可能适得其反。

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

注意力运行时间还受到tile大小的影响。图5显示，当$G \imes (1 + D)$超过128（所测试注意力内核的软件tile大小）的整数倍时，运行时间呈阶梯式增长。若$G \imes (1 + D)$落在两个tile边界之间，最后一个tile只被部分使用，但代价与完整tile相当。

指导原则3：若选择$D > \\frac{128}{G} - 1$，优先选取使$G \imes (1 + D)$为128整数倍的D值，以避免tile利用率不足。

指导原则1与指导原则2、3的相对权重，取决于在首选工作点下FFN与注意力各占运行时间的比例。此外，通信也会随D增大而增加验证开销，但计算与通信的重叠可以缓解这一问题。

Pareto曲线远端的草稿长度

在Pareto曲线的极右端，B非常小，固定的内核启动和后处理开销主导了计算和通信内核的执行时间。这些固定开销不会随验证Token数量的增加而显著增长，使得验证开销基本与草稿Token数量无关。

MoE模型随草稿长度增加会激活更多专家，但结合模型分片策略和分组GEMM等高效内核，可以将这部分开销控制在较低水平。因此，在接受率保持较高的前提下，更大的草稿长度在低延迟区域同样有益。

在极低延迟场景下，顺序内核启动次数决定了工作负载延迟。由于内核启动次数与层数成正比，对于层结构与目标模型相似的自回归草稿模型，加速比可近似为：

$$\\mathrm{speedup} = \\frac{L_{\\mathrm{target}} \imes \\mathrm{AL}}{L_{\\mathrm{target}} + D \imes 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。

选择草稿机制

确定D之后，下一步是决定如何生成这些Token以最大化加速比。

多年来已有多种技术被提出，各有不同的训练、参数和运行时成本权衡。外部草稿方法使用独立的小型大语言模型；MTP、EAGLE-3、DFlash和DSpark则利用辅助层结合目标模型的信息来预测Token；后缀和n-gram方法不依赖模型，而是复用Token流中已出现的模式。

表1对比了主要草稿方法的Token生成方式、服务时内存占用及草稿开销。

为量化AL与草稿延迟之间的权衡，首先观察AL随D的变化规律。

图6展示了以Qwen 3.5 122B A10B为目标模型，在SPEED-Bench上AL随D变化的情况。SPEED-Bench是NVIDIA开发的推测解码基准测试，旨在代表真实生产工作负载，涵盖编程、摘要等多种任务领域，并在不同输入序列长度下设有多个分组。推荐使用SPEED-Bench进行AL对比测试。

在32K分组上，Qwen 3.5 35B A3B在D=9时AL达到6，而4B草稿模型的AL超过5。MTP和DFlash的AL随D增长趋于平稳。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、草稿延迟以及训练和部署成本，选择在目标工作负载和硬件上能带来最佳解码加速的草稿机制。

开始推测解码协同设计

以下五条指导原则可作为在Pareto前沿选择D和草稿机制的检查清单：

增大推测解码草稿长度，将GEMM推入计算瓶颈区域，同时避免增加KV缓存容量压力。

当注意力主导解码时间时，以$D = \\frac{128}{G} - 1$作为起点。

对于更大的草稿长度，优先选取使$G \imes (1 + D)$为128整数倍的值，以与注意力内核tile大小对齐。

在极低延迟场景下，使用快速草稿机制，仅在AL增益能够证明增加草稿开销合理时才继续增大D。

通过平衡AL与草稿开销来选择草稿机制，在真实服务条件下对二者进行基准测试，并综合考量训练和部署成本。

对于训练后添加的草稿器，NVIDIA/Model-Optimizer中提供了EAGLE-3、DFlash和DSpark的即用型训练示例。可参考NVIDIA Nemotron 3.5 Lightning的实践经验：先微调DSpark，再量化为FP8或NVFP4。以这些示例为起点，在自己的模型、工作负载和硬件上验证AL和端到端速度。

Q&A

Q1：推测解码是如何加速大语言模型推理的？

A：推测解码通过让小型草稿模型先预测多个Token，再由目标大模型并行验证，从而减少解码迭代总次数。目标模型按顺序接受草稿Token，遇到不匹配时停止，并从该位置重新预测。由于只保留目标模型接受的Token，最终输出与标准解码完全一致，但整体推理速度得到提升。

Q2：草稿长度D应该怎么选？

A：草稿长度D的选择取决于工作负载类型和硬件特性。当线性层计算（GEMM）是瓶颈时，应增大D以将计算推入计算瓶颈区域；当注意力机制主导时，推荐$D = \\frac{128}{G} - 1$（G为每个KV头对应的查询头数）。在极低延迟场景下，只在AL增益能抵消草稿开销时才增大D。整体上需在Pareto前沿不同位置选取不同的D值。

Q3：SPEED-Bench是什么，为什么推荐用它测试推测解码？

A：SPEED-Bench是NVIDIA专为推测解码开发的基准测试工具，目标是模拟真实生产工作负载。它覆盖编程、摘要等多种任务类型，并在不同输入序列长度下设有多个测试分组。相比单一任务的测试，SPEED-Bench能更全面地反映推测解码在实际部署中的接受长度（AL）表现，因此被推荐用于不同草稿方法之间的AL对比测试。