数据中心/云端

用于降低 AI 推理延迟的预测性解码简介

使用大语言模型(LLM)生成文本时,通常会面临一个基本瓶颈。尽管 GPU 能够提供强大的计算能力,但由于自回归生成本质上是顺序进行的,大部分算力往往处于闲置状态:每个 token 的生成都需要完成一次完整的前向传播,重复加载权重,并在每一步中进行内存同步。这种频繁的内存访问与逐 token 的依赖关系共同导致了较高的延迟,降低了硬件利用率,进而限制了整体系统效率。

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

什么是预测解码?

预测解码是一种推理优化技术,通过将目标模型与轻量级的草稿机制结合,快速生成多个候选的下一个 token。目标模型在一次前向传播中对这些候选进行验证,采纳与自身预测一致的最长前缀,并在此基础上继续生成后续内容。相比标准的自回归解码方式(每次仅生成一个 token),该技术能够一次性产出多个 token,从而在保持准确性的同时有效降低延迟、提升吞吐量。

尽管大语言模型(LLM)功能强大,但其运行通常逼近AI硬件的极限,导致在大规模提升用户体验方面面临挑战。预测性解码通过将部分计算任务卸载至资源密集型模型,提供了一种可行的替代方案。

预测解码的工作原理类似于实验室中的首席科学家,由经验较少但效率较高的助手负责处理常规实验任务。助手能够快速完成基础检查,而科学家则专注于整体验证与进度把控,并在必要时介入纠正或接管工作。

借助预测性解码,轻量级助手模型会生成多个可能的后续内容,而更大的模型则对这些内容进行批量验证。这种方法能够减少连续生成的步骤,缓解内存带宽的瓶颈。关键在于,这一加速过程在不牺牲输出质量的前提下实现,因为验证机制会筛除那些与基准模型可能生成结果不一致的内容。

使用 draft-target 和 EAGLE-3 的预测性解码基础知识

本节介绍预测解码的核心概念,并解析其高效运作的关键机制。首先,通过 Transformer 的前向传播过程,展示序列如何实现并行处理;随后,逐步阐述草稿生成、验证以及基于草稿与目标模型的采样方法。这些基础内容共同构建了理解经典草稿-目标方法及 EAGLE-3 等先进技的必要背景。

预测性解码的 draft-target 方法是什么?

“draft-target”方法是预测性解码的一种经典实现,可作为双模型系统运行。其中,主模型是希望加速其输出的大型高质量目标模型,而另一个较小且运行更快的模型则作为草稿模型,通常为目标模型的蒸馏或简化版本,与主模型协同工作。

回到实验室科学家的类比,可将目标模型视为一位严谨求实的科学家,致力于确保结果的准确无误;而草稿模型则如同一位快速提出各种可能性的助手,由科学家进行验证与筛选。图1展示了这种协作机制的实际应用:草稿模型快速为目标模型生成四个候选 token,目标模型经过验证后保留其中两个,并自行生成一个额外的 token。

A gif showing an example where the input is “The  Quick”. From this input, the draft model proposes “Brown”, “Fox”, “Hopped”, “Over”. The input and draft are ingested by the target model, which verifies “Brown” and “Fox” before rejecting “Hopped” and subsequently everything after. “Jumped” is the target model’s own generation resulting from the forward pass.
图 1。用于预测解码的 draft-target 方法采用双模型系统

使用 draft-target 方法进行预测解码包含以下步骤:

草稿生成

一个更小且更高效的模型会生成一系列候选 token(通常为 3 到 12 个)。这种机制通常以在相同数据分布上训练的独立小型模型的形式实现,而目标模型的输出通常被用作该小型模型训练时的参考标签。

并行验证

目标模型在一次前向传递中同时处理输入序列和所有草稿 token,并计算每个位置的概率分布。这种并行处理方式显著提升了效率,因为它充分调动了目标模型的全部计算能力,避免了顺序生成过程中计算资源的闲置。由于原始前缀对应的键值(KV)已预先计算并缓存在内存中,因此在验证阶段,仅需对新生成的预测 token 进行计算。随后,被验证通过的 token 将被选为新的前缀,用于下一轮生成。

拒绝采样

拒绝采样是生成目标模型概率分布后,在决策阶段进行的过程。

拒绝采样的关键在于其接受逻辑。如图2所示,该逻辑通过比较草稿模型 P(草稿) 的提议概率与目标模型 P(目标) 的实际概率,来决定是否接受采样。

对于前两个词元“Brown”和“Fox”,目标概率 P(目标) 高于草稿概率 P(草稿),因此可以接受。然而,对于“Hopped”,P(目标) 显著低于 P(草稿),表明该预测结果不可靠。

当逻辑拒绝“Hopped”等 token 时,该 token 及其在草稿中后续的所有 token 都会被丢弃。随后,生成过程将从上一个被接受的 token “Fox” 恢复,转为标准的自回归生成,以产生修正后的 token。

A gif showing the verification phase within the target model. P(Target) and P(Draft) are compared for each token. “Brown” passes because P(Target) ≥ P(Draft). “Hopped” Fails because P(Target) ≤ P(Draft). As each following token is affected by previous generations, all draft tokens past “Hopped” are discarded. The final generation is thus the prefix plus “Brown Fox Jumped”, where “Brown and Fox” are accepted draft generations and “Jumped” a generation solely from the target model.
图 2。接受逻辑是并行验证期间拒绝采样的关键方面

只有当草稿模型生成的 token 与目标模型生成的 token 相匹配时,该 token 才会被接受。这种严格的逐 token 验证机制确保了最终输出与目标模型的输出完全一致,从而在不牺牲准确性的前提下实现推理加速。

接受 token 的数量与总生成 token 数量的比值即为接受率。接受率越高,加速效果越显著;而在最不利的情况下,若所有草稿 token 均被拒绝,则每次仅能生成一个目标模型 token。

EAGLE 的预测性解码方法是什么?

EAGLE,即提升语言模型推理效率的外推算法,是一种在特征层面运作的预测解码方法,通过目标模型输出层之前的隐藏状态进行推断。与依赖独立草稿模型生成候选 token 的 draft-target 方法不同,EAGLE 采用一个轻量级的自回归预测头,直接从目标模型的隐藏状态中提取特征。这种方法避免了训练和运行额外模型所带来的计算开销,同时仍能使目标模型在每次前向传播过程中验证多个候选 token。

EAGLE-3,第三版,在此基础上进一步优化,引入了目标模型的多层融合特征表示,并将低、中、高层的嵌入信息直接传递至绘图头。该版本沿用了 EAGLE-2 中上下文感知的动态草图树机制,用于生成多个链式假设。随后,目标模型通过并行树注意力机制对这些候选 token 进行验证,有效剪枝无效路径,从而提升接受率与系统吞吐量。图 3 展示了该流程的实际运行过程。

A gif showing that the lightweight EAGLE head is not a standalone model. It drafts tokens from feature outputs taken from the target model’s layers, generates prediction trees, then feeds this back into the model for verification.
图 3。EAGLE-3 起草机制从目标模型中生成候选标记树

什么是 EAGLE 头?

与草稿目标方法中采用独立的小型模型不同,EAGLE-3 改为将轻量级绘图组件直接附加到目标模型的内部层,形成所谓的“EAGLE 头”。该头部通常由若干轻量级 Transformer 解码器层和一个最终的线性层构成,本质上是主模型结构中核心组件的精简版本,可视为一个微型化但功能聚焦的子结构。

该 EAGLE 头不仅能生成单个序列,还能生成完整的候选 token 树。这一过程具有实例自适应特性,即头部在构建树的过程中会动态评估自身的置信度,并在置信度低于设定阈值时停止扩展。这使得 EAGLE 头能够高效探索多层级的生成路径:在文本可预测的部分生成较长的分支,而在复杂部分则生成较短的分支,整个过程仅需目标模型的一次前向传播,运行成本较低。

DeepSeek-R1 中的多 token 预测是什么?

与 EAGLE 类似,多 token 预测(MTP)是一种被 DeepSeek 多次迭代采用的预测技术,其核心思想是让模型一次性预测多个后续 token,而非仅预测下一个 token。MTP 采用多头结构,每个头充当一个 token 的起草者:第一个头负责预测首个草稿 token,第二个头预测下一个 token,依此类推。随后,主模型按顺序验证这些预测结果,并保留其中连续匹配的最长前缀。该方法的优势在于无需依赖独立的草稿模型,即可实现高效推理。

从本质上讲,这种技术与 EAGLE 的预测性解码类似,两者都会生成多个候选 token 进行验证。但提案的生成方式有所不同:MTP 采用专门的多 token 预测头,而 EAGLE 则利用单一头部,通过推断模型内部的特征状态来构建候选序列。

如何实现预测性解码

您可以使用 NVIDIA TensorRT-Model Optimizer API 将预测性解码应用于自有模型。请按照以下步骤,利用 Model Optimizer 的预测解码模块,将模型转换为支持 EAGLE-3 预测解码的格式。

第一步:加载原始的 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()

# Load original HF model
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 工具进行转换。

# Read Default Config for EAGLE3
config = EAGLE3_DEFAULT_CFG["config"]

# Hidden size and vocab size must match base model
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,
    }
)

# Convert Model for eagle speculative decoding
mtsp.convert(model, [("eagle", config)])

查看实战教程,了解如何在 TensorRT-Model-Optimizer/examples/speculative_decoding GitHub 库中将此演示扩展为可部署的端到端预测解码微调工作流。

预测性解码如何影响推理延迟?

在标准自回归生成过程中,核心的延迟瓶颈在于每一步骤必须按固定顺序执行。若单次前向传播(包括加载权重和计算 token)耗时 200 毫秒,则生成三个 token 将始终需要 600 毫秒(即三个连续步骤 × 每步 200 毫秒)。这种延迟会表现为用户感知中逐步累积的等待时间。

预测性解码可能将多个等待时间合并为一次。通过快速草稿机制预测两个候选 token,随后在一次 250 毫秒的前向传播中完成全部验证,模型可在 250 毫秒内生成三个 token(包含两次预测和一次基础模型生成),而无需耗时 600 毫秒。图 4 展示了这一原理。

A gif showing a base model (top) using standard autoregressive generation generating a single token in each 200 ms pass, taking 600 ms to generate three. A model with speculative decode (bottom) took slightly longer on one pass (250 ms), but generated three tokens in a single pass.
图 4。生成 (采用和不采用预测解码)

用户通常不会逐字阅读响应,而是以更快的多 token 块形式接收内容。这一点在聊天机器人等交互式应用中尤为明显:较低的响应延迟能够带来更流畅、更自然的对话体验。图5模拟了在启用与关闭预测性解码两种情况下,假设的聊天机器人对话场景。

A gif with side-by-side chatbot outputs labeled ‘Speculative Decoding Off’ (left) and ‘Speculative Decoding On’ (right). The chatbot on the right shows how speculative decoding reduces the time it takes to generate each token or batch of tokens, shortening the user's waiting period. This makes the chatbot feel more responsive, fluid, and natural to interact with.
图 5。启用预测解码的聊天机器人 (右) 生成文本的速度远快于启用预测解码的聊天机器人 (左)

开始使用预测性解码 

预测性解码正逐渐成为加速大语言模型推理的核心策略。从草图目标生成与并行验证的基础方法,到EAGLE-3等进阶技术,这类方法有效应对了顺序生成 过程中计算资源闲置的关键挑战。

随着工作负载规模的扩大以及对更低延迟和更高系统效率需求的提升,预测性解码等技术将变得愈发重要。将这些方法与 NVIDIA TensorRT-LLM、SGLANG 和 vLLM 等框架结合使用,有助于开发者在实际应用场景中部署性能更优、实用性更强且更具成本效益的模型。

准备开始了吗?请查看 TensorRT-Model-Optimizer/examples/speculative_decoding GitHub 仓库中的 Jupyter Notebook 教程,尝试将预测性解码应用于您自己的模型。

致谢

感谢 NVIDIA 工程师 Yu Chenhan 和 Guo Hao 为本文的开发与撰写所做出的贡献。

 

标签