文章
通过推测加速声纳
推测解码通过使用快速且小的草案模型生成由较大的目标模型验证的完成候选，来加速大型语言模型(LLM)的生成速度。
在这种方案下，与其让昂贵的目标模型生成一个单一的标记，不如在一步中发出多个标记。这里我们介绍在Perplexity上应用各种推测解码实施细节，以减少Sonar模型上的标记间延迟。
推测解码
推测解码利用自然语言的结构和变压器的自回归特性来加速标记生成。尽管较大的模型如Llama-70B携带更多的知识，但在一些简单任务上它们与较小的模型如Llama-1B表现相似。这种重叠表明某些序列通过成本较低的模型生成更佳，而较复杂的问题则留给较大的模型。挑战在于判断哪种完成更好，以及较小模型的生成质量是否与较大模型相当。
幸运的是，LLM是自回归的变压器：当给定一系列标记时，它们会输出下一个标记的概率分布。此外，从输入序列中标记相关的中间特征中导出的logits也表明模型发出这些确切标记的可能性。这一特性使得推测成为可能：如果一段标记序列从输入前缀开始由较小模型生成，则可以把它通过较大的模型以确定它与目标模型的契合度。每个候选的前缀都根据概率得分，并选择超过接受阈值的最长一个。作为额外好处，目标模型也会免费提供一个后续标记：如果草案模型生成n个标记，则在一步中最多可以发出n + 1。
在推理时间，推测采样过程大致上可以被分为4个阶段：
预填充：目标和草案模型都必须在输入序列上运行以填充KV缓存条目。虽然一些方案如Medusa使用简单的稠密层进行预测，但在这篇文章中，我们重点关注需要自己KV缓存的基于变压器的草案。
草案生成：草案模型迭代生成一定数量的固定标记。草案序列可以是线性的，也可以让模型在给定深度（EAGLE, Medusa）内探索类似树的结构。在这里，我们重点关注线性序列。
接受：目标模型运行草案序列，构建与每个草案标记对应的logits。确定最长可接受序列的长度。
目标生成：由于目标生成了logits，在序列的不匹配位置或尾端logits对应于一个后续标记。这些logits可以被采样以提供目标的一个稳健标记，完成序列。
多种方法可用于实现推测解码。在此文章中，我们将专注于我们用来通过内部1B模型加速Sonar模型的方案，以及我们正在开发的用于加速DeepSeek规模模型的预测机制。
目标-草案
通过将现有的小型LLM作为草案模型与目标模型结合以生成候选序列，可以实现推测解码。在生产中，我们加速了使用在相同数据集上微调的Llama-1B模型的Sonar。虽然这种方法无需从头训练一个草案，但小模型仍然占用大量的KV缓存容量，并引入了轻微的预填充开销，增加了TTFT。
在这种方案下，解码器仅对解码专用批次进行推测，在预填充期间或混合的预填充-解码批次中通过标准采样生成标记。在预填充阶段，目标logits会立即被采样以在草案的KV缓存中也预填充新生成的标记。草案尚未被采样，但所产生的logits会被带到解码阶段。
在解码中，草案模型前进，在每个阶段采样最高的标记。在达到所需的草案长度后，标记会通过目标模型生成基于此的logits，采样器确定所接受的序列长度。接受通过比较来自草案和目标的全概率分布来确定。由于目标总是在草案序列后输出一组logits，因此可以对其进行采样以产生额外的输出。由于草案模型尚未看到被接受的标记，它会被重新运行以填充其相应的KV缓存条目，为下一个解码步骤做准备，并再次带入logits。
EAGLE
EAGLE是一种推测解码方案，通过类似树的有可能草案标记遍历生成多个草案序列。使用对每个节点考虑Top-K候选者而不是在线性序列中跟随评分最高的标记的固定（EAGLE）或动态形状（EAGLE-2）树进行探索。对序列进行评分，选择最长的合适一个继续前进，同时也附加一个来自目标的额外标记。
为实现更准确的预测，EAGLE草案模型不仅基于标记预测，还使用目标模型的特征（最后一层的隐藏状态）。EAGLE的缺点是需要训练自定义的小草案模型，其需要足够准确以在低延迟的预算内部生成合适的候选者。通常，一个草案模型是与原始模型的解码器层相同的单变压器层，通过其嵌入和lm_head投影与目标紧密耦合。由于这需要更少的KV缓存容量，EAGLE的内存占用更低。
为在目标模型中验证类似树的序列，必须使用自定义注意力掩码。不幸的是，使用自定义注意力掩码对整个序列显著减慢了现实输入长度的注意力（最多降低50%），抵消了一些通过推测能实现的加速。我们尚未因为这个原因将完整的树探索部署到生产中，而是专注于通过DeepSeek-V3技术报告提出的MTP样式方案进行单标记预测的特殊案例。
MTP
此方案类似于草案-目标解码，除了用于预测的隐藏状态与标记一起使用。与常规草案-目标推测相比，在预填充和解码阶段都需要做更多工作。草案模型使用标记和隐藏状态：标记t_{i+1}从对应于标记t_i的logitsL_i中采样，而它们又是从隐藏状态H_i导出的结果。因此，相对于目标输出的隐藏状态向量，输入的标记缓冲区必须向左移动一步。下面的图标记了用于训练的对应关系，以及推理期间的移动。
解码流程与草案-目标解码非常相似，除了同时携带隐藏状态和logits。我们的实现共享与所有相关的采样和logit处理logits，专门进行模型前向调用。当预测多个标记时，草案模型使用草案隐藏状态进行预测，同时根据其自身特征填充KV缓存条目。从长远看，这可能降低精度。随后，在为目标预测填充KV缓存条目时运行草案模型，我们在整个序列上运行它，以上司准确的目标隐藏状态作为输入。由于这些草案模型较小，处理其他标记的额外成本是可以忽略不计的。
训练MTP头
为了利用MTP功能，我们建立了基础设施，在Perplexity的数据集上训练附加到我们的微调模型的MTP头，在一个节点上运行8xH100设备。大约一天内，我们可以为从Llama-1B到Llama-70B和DeepSeek V2-Lite的模型构建头。对于较大的模型，我们依赖于在微调过程中构建的MTP头。
MTP训练的目标是将从目标隐藏状态中外推的草案隐藏状态和logits与目标的下一个标记logits和隐藏状态匹配起来。由于隐藏状态的推理昂贵，我们使用针对推理优化的目标模型实现预先计算它们，以在训练期间使用。然而，为验证推理MTP实现并确保量化或优化造成的数值差异不会阻碍结果，对于验证损失和精度估算，我们完全重新使用目标模型和草案模型的推理实现。
当从原始论文中使用的ShareGPT数据集扩展到更大的样本时，我们注意到EAGLE论文中概述和实现的MTP头架构未能为70B大小的模型训练成功。与包含大量较短序列的ShareGPT不同，我们在数量稍少但实质上更长的提示上进行训练。由于原始EAGLE头在结构上略有偏离典型的变压器，我们重新引入了一些被剥离的RMS归一化层。我们发现这不仅让训练收敛，而且还使头的精度提高了几个百分点。
不仅分层规范促进训练，重新引入规范也在数学上直观。MTP头重用目标模型的嵌入和logit投影，因为它们可以是相当大的（Llama 70B约为2 GB）。在训练过程中，这些是冻结的，期望是MTP层学习将预测嵌入与原模型在训练期间学习的投影层相同的向量空间内。通过删除规范，期望单个MLP学习与一个MLP后跟一个规范的功能相同，这阻碍了草案和目标模型的隐藏状态的匹配。
推测解码的推理
在推理引擎中，为了生成输入序列的标记，它们需要首先分组为合理大小的批次，然后必须在KV缓存中分配页面用于下一个标记。输入标记和KV页面信息随后被打包到一个缓冲区广播给所有运行该模型的并行排名。最后，元数据被复制到GPU内存中并执行模型以生成logits从中采样出下一个标记。
与某些通过一个胶合器在推理服务器之间松散耦合草案和目标的实现不同，我们的草案-目标对紧密耦合并协同生成。批次调度和KV页面分配在所有形式的推测解码中在模型之间共享：这统一了集成推理服务器的模型逻辑，因为它们全部暴露相同的接口。
Perplexity的推理运行时围绕FlashInfer构建，它确定需要构建以配置和调度注意力内核的元数据。给定组成批次的一些输入序列，对于预填充、解码或验证，必须在CPU侧完成工作以分配中间缓冲区并填充注意力中使用的某些常量缓冲区。此工作是批次调度和KV页面分配成本之外，还要隐藏延迟以最大化GPU使用率。
虽然我们在没有推测的情况下完全并行化了推理的CPU侧和GPU侧工作，但我们发现推测解码的CPU-GPU平衡更为复杂。主要挑战源于被接受的标记数量决定了后续运行的序列长度，导致难以避免的GPU到CPU同步点。我们尝试了不同的调度方案，以最好地隐藏CPU工作的延迟。
草案-目标调度
尽管比目标模型更小，当整个位LLM用作草案时，它仍然在GPU上引入相当的延迟，提供了一些隐藏昂贵CPU操作的空间。由于较小的模型没有从张量并行中受益，目标和草案各自跨的排名数量不匹配。在我们的实现中，草案模型仅在TP组的领导排名上运行。
如前所述，解码步骤将logits传递到下一个运行。这允许我们重叠草案模型的一次执行与CPU侧的批次调度工作。批次备好后，采样器和草案连续调用生成草案标记。与此同时，验证批次为目标模型组合并与并行工作者同步。目标logits被验证和采样以确定接受的序列长度。在这一点上，GPU到CPU同步是必要的，以便确定后续序列长度。由于草案模型仅在领导节点上运行，其批次被顺序设置并其执行被启动以用目标生成的额外标记填充其KV缓存条目。当前运行的草案生成的logits将用于在后续运行中采样第一个草案标记。最重要的是，在草案运行时，下一批可以被调度。
单标记的MTP调度
虽然运行时还没有提供类似鹰式草案树探索，我们实施了这种方案的特例，考虑一个由一个单一变压器解码器层尺寸的模型生成的草案标记线性序列。此方案可与DeepSeek R1的开源权重一起用于草案预测。预测单个标记的子案例有趣的是，大型MTP层实现了足够高的接受率以证明其开销是合理的。
MTP调度相对更加复杂，因为草案模型运行更快，隐藏了更少的CPU侧延迟。此外，草案与目标模型一起分片，需要进行批次信息的共享内存传输。运行开始于转移批次信息并从携带过来的logits中采样第一个标记，与之前方案类似。接下来，目标通过处理2 * D标记进行验证，其中D是解码批次大小。这对于在较慢的互连如InfiniBand上的混合专家模型（MoE）是理想的，因为批次平均分为两半。目标的隐藏状态传递到下一个草案运行，而logits传递给采样器进行验证。
通过在GPU上进行一有限量的额外工作，我们避免了草案序列接受后的CPU到GPU同步。在目标的输入标记被移位后，一个内核将后续目标标记插入它们的相应位置。草案然后在与目标相同的批次信息上被重新运行，填充KV缓存条目并为下一次运行构建logits和隐藏状态，对未被接受的标记进行多余工作。在这些情况下，由于草案模型的体积较小，未使用工作的延迟几乎是不可测的。与草案运行并行，序列长度在CPU上确定，并启动下一批的调度，而无需等待GPU工作完成。
在草案层的额外工作开销在注意力中是难以察觉的，然而在MLP层中问题更为明显。由于矩阵乘法指令沿标记数量维度按64的边界填充，如果翻倍不需要显著更多的区块，开销将被隐藏。对于较长的草案序列，开销更为昂贵，并且用于常规草案-目标型号的方案效果更好。