推测解码技术降低AI推理延迟详解

推测解码技术降低AI推理延迟详解

概述

使用大语言模型生成文本时常会遇到基本瓶颈。虽然GPU提供大规模计算能力，但由于自回归生成本质上是顺序执行的：每个令牌都需要完整的前向传递、重新加载权重并在每一步同步内存，导致大部分计算能力处于闲置状态。这种内存访问和逐步依赖的组合会增加延迟、降低硬件利用率并限制系统效率。

推测解码技术有助于突破这一瓶颈。通过同时预测和验证多个令牌，该技术缩短了结果生成路径，使AI推理更快、响应更迅速，在保持输出质量的同时显著降低延迟。本文将探讨推测解码的工作原理、适用场景以及如何在某中心GPU上部署先进的EAGLE-3技术。

什么是推测解码？

推测解码是一种推理优化技术，它将目标模型与轻量级草案机制配对，快速提出几个后续令牌。目标模型在单个前向传递中验证这些提议，接受与其自身预测匹配的最长前缀，并从中继续生成。与标准自回归解码（每次传递产生一个令牌）相比，该技术让系统能够一次生成多个令牌，在不影响准确性的情况下降低延迟并提高吞吐量。

尽管大语言模型能力强大，但常常挑战AI硬件的极限，使得在大规模优化用户体验变得困难。推测解码通过将部分工作卸载到资源密集度较低的模型来提供替代方案。

推测解码的工作方式类似于实验室的首席科学家依赖经验较少但效率高的助手处理常规实验。助手快速处理检查清单，而科学家专注于验证和进展，在必要时介入纠正或接管。

通过推测解码，轻量级助手模型提出多个可能的延续，较大模型批量验证它们。最终好处是减少顺序步骤的数量，缓解内存带宽瓶颈。关键的是，这种加速在保持输出质量的同时发生，因为验证机制会丢弃与基线模型自身可能生成结果不同的任何结果。

使用草案-目标和EAGLE-3的推测解码基础

本节阐述推测解码背后的核心概念，分解使其有效的机制。首先，Transformer前向传递展示如何并行处理序列。后续步骤包括草案生成、验证和使用草案-目标方法进行采样。这些基础共同提供了理解经典草案-目标方法和先进技术（如EAGLE-3）所需的背景。

什么是草案-目标方法的推测解码？

草案-目标方法是推测解码的经典实现，作为一个双模型系统运行。主要模型是大型高质量目标模型，其输出需要加速。与之协同工作的是一个更小更快的草案模型，通常是目标的蒸馏或简化版本。

回到实验室科学家的类比，将目标视为确保正确性的细致科学家，而草案则是提出科学家随后验证的可能性的快速助手。图1显示了这种合作关系的运作，草案模型快速生成四个草案令牌供目标模型验证，目标模型保留两个并自行生成一个额外令牌。

草案-目标方法的推测解码涉及以下步骤：

草案生成
更小更高效的机制生成候选令牌序列（通常为3到12个令牌）。通常，这采用在相同数据分布上训练的独立较小模型的形式。目标模型的输出通常作为草案模型训练的基本事实。

并行验证
目标模型在单个前向传递中同时处理输入序列和所有草案令牌，计算每个位置的概率分布。这种并行处理是关键效率增益，因为它利用了目标模型的全部计算能力，而不是在顺序生成期间让其利用率不足。由于KV缓存中原始前缀的值已经计算并存储，只有新的推测令牌在此验证传递期间产生计算成本。然后选择已验证令牌形成下一个生成步骤的新前缀。

拒绝采样
拒绝采样是在生成目标模型的概率分布后发生的决策阶段。

拒绝采样的关键方面是接受逻辑。如图2所示，该逻辑比较草案模型的提议概率P(Draft)与目标模型的实际概率P(Target)。

对于前两个令牌"Brown"和"Fox"，P(Target)高于P(Draft)，因此它们被接受。然而，对于"Hopped"，P(Target)显著低于P(Draft)，表明预测不可靠。

当诸如"Hopped"的令牌被接受逻辑拒绝时，它和草案中所有后续令牌都被丢弃。然后过程恢复到从最后一个接受的令牌"Fox"开始的标准自回归生成，以产生修正的令牌。

只有当草案令牌与目标模型会生成的内容匹配时，它才被接受。这种严格的逐令牌验证确保最终输出与目标模型会产生的结果相同，保证加速不会带来准确性损失。

接受的令牌数量与总生成数量的比率是接受率。更高的接受率等同于更显著的加速，在最坏情况下，如果所有草案令牌都被拒绝，则只生成单个目标模型令牌。

什么是EAGLE方法的推测解码？

EAGLE（Extrapolation Algorithm for Greater Language-Model Efficiency）是一种在特征级别操作的推测解码方法，从目标模型输出头之前的隐藏状态进行推断。与依赖独立草案模型提出令牌的草案-目标方法不同，EAGLE使用轻量级自回归预测头摄取目标模型隐藏状态的特征。这消除了训练和运行第二个模型的开销，同时仍允许目标模型每个前向传递验证多个令牌候选。

第三版EAGLE-3在此基础上引入多层融合特征表示，将低、中、高级嵌入直接输入其草案头。它还使用上下文感知的动态草案树（继承自EAGLE-2）提出多个链式假设。然后这些候选令牌由目标模型使用并行树注意力进行验证，有效修剪无效分支并提高接受率和吞吐量。图3显示了此流程的运作。

什么是EAGLE头？

与草案-目标方法中使用独立的较小模型不同，EAGLE-3将轻量级草案组件附加到目标模型的内部层作为"EAGLE头"。EAGLE头通常由轻量级Transformer解码器层和最终线性层组成。它本质上是构成主模型的构建模块的微型简化版本。

这个EAGLE头不仅可以生成单个序列，还可以生成整个候选令牌树。此过程也是实例自适应的，头在构建树时评估自身置信度，如果置信度低于阈值则停止草案。这允许EAGLE头有效探索多个生成路径，为可预测文本生成较长分支，为复杂部分生成较短分支，所有这些都只需目标模型一次前向传递的运行时成本。

DeepSeek-R1中的多令牌预测是什么？

与EAGLE类似，多令牌预测是某机构多次迭代使用的推测技术，模型学习一次预测多个未来令牌，而不仅仅是紧接的下一个令牌。MTP使用多头方法，每个头充当令牌草案器。附加到模型的第一个头猜测第一个草案令牌，另一个猜测下一个，另一个猜测第三个，依此类推。然后主模型按顺序检查这些猜测并保留匹配的最长前缀。此方法自然消除了对独立草案模型的需求。

本质上，此技术类似于EAGLE风格的推测解码，两者都提出多个令牌进行验证。然而，它们在提案形成方式上有所不同：MTP使用专门的多令牌预测头，而EAGLE使用单个头推断内部特征状态以构建候选。

如何实现推测解码

可以使用某中心TensorRT模型优化器API将推测解码应用于自己的模型。按照以下步骤使用模型优化器推测解码模块将模型转换为使用EAGLE-3推测解码。

步骤1：加载原始Hugging Face模型

import transformers
import modelopt.torch.opt as mto
import modelopt.torch.speculative as mtsp
from modelopt.torch.speculative.config import EAGLE3_DEFAULT_CFG

mto.enable_huggingface_checkpointing()

# 加载原始HF模型
base_model = "meta-llama/Llama-3.2-1B"
model = transformers.AutoModelForCausalLM.from_pretrained(
    base_model, torch_dtype="auto", device_map="cuda")

步骤2：导入EAGLE-3的默认配置并使用mtsp转换

# 读取EAGLE3的默认配置
config = EAGLE3_DEFAULT_CFG["config"]

# 隐藏大小和词汇表大小必须匹配基础模型
config["eagle_architecture_config"].update(
    {
        "hidden_size": model.config.hidden_size,
        "vocab_size": model.config.vocab_size,
        "draft_vocab_size": model.config.vocab_size,
        "max_position_embeddings": model.config.max_position_embeddings,
    })

# 转换为eagle推测解码模型
mtsp.convert(model, [("eagle", config)])

查看实践教程，将此演示扩展为TensorRT模型优化器/examples/speculative_decoding GitHub存储库中的可部署端到端推测解码微调流水线。

推测解码如何影响推理延迟？

标准自回归生成中的核心延迟瓶颈是每个步骤的固定顺序成本。如果单次前向传递（加载权重和计算令牌）需要200毫秒，生成三个令牌将始终需要600毫秒（三个顺序步骤乘以200毫秒）。用户将此延迟体验为不同的累积等待期。

推测解码可以将这些多个等待期合并为一个。通过使用快速草案机制推测两个候选令牌，然后在单个250毫秒前向传递中验证它们所有，模型可以在250毫秒内生成三个令牌（两个推测加一个基础模型生成），而不是600毫秒。此概念如图4所示。

用户不是看到响应逐字出现，而是看到它以更快的多令牌块形式具体化。这在交互式应用程序（如聊天机器人）中尤其明显，较低的响应延迟创造了更流畅自然的对话。图5模拟了开启和关闭推测解码的假设聊天机器人。

开始使用推测解码

推测解码正在成为加速大语言模型推理的基本策略。从草案-目标生成和并行验证的基础到EAGLE-3等先进方法，这些方法解决了顺序令牌生成期间计算闲置的核心挑战。

随着工作负载扩展以及对更快响应时间和更好系统效率的需求增长，推测解码等技术将发挥越来越核心的作用。将这些方法与某中心TensorRT-LLM、SGLANG和vLLM等框架配对，确保开发人员可以在实际环境中部署性能更高、更实用且更具成本效益的模型。

准备开始？查看TensorRT模型优化器/examples/speculative_decoding GitHub存储库中的Jupyter笔记本教程，尝试将推测解码应用于自己的模型。
更多精彩内容 请关注我的个人公众号 公众号（办公AI智能小助手）或者 我的个人博客 https://blog.qife122.com/
对网络安全、黑客技术感兴趣的朋友可以关注我的安全公众号（网络安全技术点滴分享）