使用预测解码共同设计 AI 模型，加快 LLM 推理速度

本文是关于 AI 模型协同设计的系列文章中的第三篇。它探讨了如何使用预测性解码在保持准确性的同时加速 LLM 推理，并提供了五个在整个帕累托前沿选择吃草长度和吃草机制的指南。

有关模型设计选择如何在不牺牲准确性的情况下影响吞吐量和交互性的讨论，请参阅 AI 模型协同设计：硬件友好型 LLM 设计 (第一部分) 。如需了解组大小 (每个 KV 头的查询头) 、头维度和序列长度如何影响密集注意力的表现，请参阅 共同设计用于快速交互式长上下文推理的 AI 模型注意力 (第二部分) 。

什么是预测解码？

预测解码是一种通过预测每次迭代的多个 token 来加速 LLM 推理自回归解码阶段的技术。一个小型草稿模型首先预测几个可能的下一个 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 显示了具有代表性的专家 GEMM 大小为 6144 x 6144 时，在各种草稿长度下，每秒 teraFLOPS 的规模。显然，较高的拔模长度允许 GEMM 在较低的有效批量大小下实现峰值性能。

值得注意的是，与 混合专家模型 (MoE) 变得越来越稀疏，长上下文工作负载增加了 KV 容量压力，每位专家的有效并发率也会降低，这使得更大的吃草长度在整个帕累托地区极具吸引力。

准则 1： 增加预测解码草稿长度，在不增加 KV 缓存容量压力的情况下，将 GEMM 推向计算受限的区域。

拔模长度和注意力表现

对于推理和代理式工作负载，在面向吞吐量的区域中，注意力往往占据执行时间的主导地位。解码注意力的算术强度约为 \(2 \times G\)，其中 \(G\) 是共享一个 KV 头的查询头数量，正如用于快速交互式长上下文推理的共同设计 AI 模型注意力中所述。

猜测会将其增加为 \(2 \times G \times (1 + D)\)，因为预测的token会重复使用相同的 KV 缓存。GEMM-\(M\) 的有效注意力是 \(G \times (1 + D)\)。在当前 GPU 设备上，注意力内核在 GEMM-\(M=128\) 下实现了良好的硬件利用率，从而使 \(D = \frac{128}{G} – 1\) 成为最佳吃草长度。

图 4 显示了在 32K 和 128K KV 序列长度中，\(G = 8\) 和 \(G = 32\) 的 \(D\) 增加时的标准化注意力吞吐量。\(G = 32\) 变体达到吞吐量饱和时的值较低，为 \(D\)。除了吞吐量饱和点之外，注意力不再受 DRAM 带宽的限制，其运行时随 \(D\) 进行扩展。由于 \(AL\) 随 \(D\) 呈次线性增长，因此将 \(D\) 增加到此水平之后可能会减缓注意力密集型工作负载的运行速度。

准则 2：当注意力控制解码时间时，选择 \(D = \frac{128}{G} – 1\)。

注意力运行时也取决于图块大小。图 5 显示，当 \(G \times (1 + D)\) 与基准测试注意力核函数的软件图块大小 ( 128) 相交时，运行时间以步长增加。如果 \(G \times (1 + D)\) 位于两个图块边界之间，则最后一个图块仅被部分使用，但其成本仍与完整图块差不多。

准则 3： 如果您选择 \(D > \frac{128}{G} – 1\)，首选 \(G \times (1 + D)\) 是 128 倍数的值，以避免图块利用率不足。

与准则 2 和准则 3 相比，准则 1 的相对权重取决于 FFN 中运行时间占首选操作点注意力的百分比。通信也会随着 D 值的增加而增加验证成本，尽管计算和通信重叠可以减少开销。

帕累托曲线最右侧的拔模长度

朝着帕累托曲线的最右侧移动，\(B\) 会变得非常小。在这里，固定的内核设置和后处理开销在计算和通信内核方面占据主导地位。这些固定成本不会随着已验证token的数量而显著增加，因此验证开销在很大程度上与草稿token的数量保持不变。

虽然 MoE 确实会看到越来越多的专家随着草稿长度的增加而激活，但通过结合使用模型分片策略和高效内核 (如分组 GEMM) ，可以保持较低的开销。因此，在低延迟区域，较大的拔模长度会有所帮助，只要接受度仍然很高。

在极低延迟下，连续内核启动次数决定了工作负载延迟。由于核函数启动时随层数量呈线性扩展，因此对于层结构与目标相似的自回归草稿模型，加速可近似为：

\(\mathrm{speedup} = \frac{L_{\mathrm{target}} \times \mathrm{AL}}{L_{\mathrm{target}} + D \times L_{\mathrm{draft}}}\)

其中 \(L_{\mathrm{target}}\) 和 \(L_{\mathrm{draft}}\) 分别是目标模型和草稿模型中的层数量。定义常量 draft_depth_ratio，\(\rho = \frac{L_{\mathrm{draft}}}{L_{\mathrm{target}}}\)：

\(\mathrm{speedup} = \frac{\mathrm{AL}}{1 + \rho D}\) 和草稿开销，\(O_d = \rho D\)

换言之，只有当 \(AL\) 中的增益足以抵消草稿开销时，增加 \(D\) 才会有所帮助。

准则 4： 在极低延迟下，仅增加 \(D\)，而 \(AL\) 中的增益证明增加的吃水成本是合理的。

选择草稿机制

选择 \(D\) 会告诉您要预测的 token 数量。接下来，决定如何生成这些 token 以更大限度地提高速度。

多年来，人们提出了多种技术，在训练、参数和运行时成本方面进行了不同的权衡。外部起草建议使用小型独立 LLM，而 MTP、EAGLE-3、DFlash 和 DSpark 则使用辅助层与目标模型的信息相结合来预测 token。后缀和 n-gram 方法是无模型的，而是重复使用token流中已看到的模式。

表 1 比较了关键的绘图方法，包括它们如何生成token、服务器时间内存和草稿开销。

要量化 \(AL\) 和 \(O_d\) 之间的权衡，首先观察 \(AL\) 如何使用 \(D\) 进行扩展。

图 6 显示了以 Qwen 3.5 122B A10B 为目标时，AL 在 SPEED-Bench 上如何随着 D 值的增加而变化。SPEED-Bench 是 NVIDIA 开发的预测性解码基准测试，旨在表示逼真的生产工作负载。它涵盖编码和摘要等多个任务领域，并具有不同输入序列长度的各种拆分。我们推荐使用 SPEED-Bench 进行 \(AL\) 比较。

在 32K 拆分时，Qwen 3.5 35B A3B 在 \(D = 9\) 时达到 \(AL\) 的 6，而 4B 草稿达到 AL 5 以上。随着 \(D\) 的增长，MTP 和 DFlash AL 趋于平稳。N – Gram 对这种工作负载的接受度较低，更适合具有重复标记模式的工作负载。

较高的 \(AL\) 并不等于更高的加速。您还需要考虑生成草稿的成本。

如图 6 所示，在 \(D > 3\) 下，所有外部草稿的 \(AL\) 均高于其他草稿方法。Qwen 3.5 122B MTP 要小得多，具有 25 亿个总参数和不到 1.5 亿个活动参数。在 \(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 也可能是最佳选择。

为了权衡 SPEED-Bench 和草稿延迟，准确地对两者进行基准测试至关重要。在预测解码的背景下，这意味着使用逼真的提示进行基准测试，并涵盖广泛的任务领域。我们推荐使用用于 SPEED-Bench 测量的 NVIDIA TensorRT LLM 等高性能推理框架来量化吃水开销。

除推理性能外，草稿训练范围和成本也值得考虑。虽然 MTP 必须与目标模型协同训练，但可以将 EAGLE、DFlash 和 DSpark 添加到最终模型检查点中。同样，也存在一系列用于创建外部草稿模型的选项。从头开始训练和从目标中训练蒸馏可产生最高的 AL，而通过跨模型自适应技术使现有模型适应所需目标可以显著降低训练成本，而 \(AL\) 会成为代价。

微调目标模型可以改变其输出分布和隐藏表示。因此，即使目标模型得到改进，针对特定目标检查点训练的习得起草者也可能会失去接受度。更改目标后，重新测量代表性工作负载的 \(AL\)。所需的适应能力取决于草稿机制。

MTP 是目标模型的一部分，应在微调期间继续进行训练，或在专门的后续阶段进行重新调整。EAGLE-3、DFlash 和 DSpark 等目标附加绘图程序会使用目标的隐藏状态，并且需要通过额外的训练来适应更新后的检查点。

外部草稿模型不使用目标隐藏状态，但它们仍然接近目标的输出分布，可能需要微调或蒸馏。后缀和 n-gram 方法没有习得的绘图器，因此不需要重新训练，尽管它们的有效性仍然取决于所部署工作负载中出现的重复次数。

准则 5： 考虑到 \(AL\)、草稿延迟以及训练和部署成本，选择能够为您的工作负载和硬件提供最佳解码加速的草稿机制。

开始共同设计预测性解码

将以下总结的五个准则用作在帕累托边境选择 \(D\) 和草稿机制的检查清单。

1. 增加预测的解码草稿长度，在不增加 KV 缓存容量压力的情况下，将 GEMM 推向受计算限制的区域。
2. 当注意力控制解码时间时，选择 \(D = \frac{128}{G} – 1\) 作为起点。
3. 对于较大的草稿长度，首选 \(G \times (1 + D)\) 是 128 倍数的值，以便与注意力内核图块大小保持一致。
4. 在极低延迟下，使用快速草稿机制，仅增加 \(D\)，而 \(AL\) 中的收益证明增加的草稿成本是合理的。
5. 通过平衡 \(AL\) 和工作负载和硬件的草稿用度，选择草稿机制。在实际服务条件下进行基准测试，并考虑训练和部署成本。

对于经过后训练的起草者，可在NVIDIA/Model-Optimizer中为EAGLE-3、DFlash和DSpark提供现成的训练示例。了解我们如何为NVIDIA Nemotron 3.5 Lightning实现这一目标：微调DSpark，然后将其量化为FP8或NVFP4。以这些示例为起点，然后对您自己的模型、工作负载和硬件验证AL和端到端速度。

致谢

本文是 NVIDIA 跨团队协作的成果。我们非常感谢 Bhargava Gopireddy、Ritika Borkar、Tor Tsur、Andrii Skliar、Benjamin Chislett、Yaniv Galron、Talor Abramovich、Yoav Miron、Rabeeh Karimi Mahabadi、Roger Waleffe、Uni Karpas、Ran Zilberstein、Brian Pharris 和 Eduardo Alvarez 提供的帮助。