发表时间: 2026-03 · Blog post by Amazon and NVIDIA Team (blog.vllm.ai)
原文: https://blog.vllm.ai/blog/2026-03-13-p-eagle
作者/机构:Amazon and NVIDIA Team (AWS & NVIDIA)
一句话结论 P-EAGLE 提出了一种用于大语言模型推理的并行投机解码方法,通过在单次前向传播中生成所有草稿 token,在 NVIDIA B200 上实现了比原生 EAGLE-3 高达 1.69 倍的推理加速。
要解决什么问题 现有的投机解码(Speculative Decoding)最前沿方法 EAGLE 采用的是自回归的方式生成草稿 token。这种机制存在一个致命的隐性卡点:为了生成 $K$ 个草稿 token,草稿模型必须进行 $K$ 次串行的前向传播。随着现代推理模型输出序列越来越长,为了获得更高的加速比,系统需要预测更多的 token(即增大投机深度 $K$)。但在自回归机制下,草稿模型的延迟会随着投机深度呈线性增长,这种不断累加的串行开销最终会吞噬掉投机解码带来的性能收益,严格限制了系统进行激进投机的能力。此外,如果直接将草稿模型改为并行训练来绕过这个卡点,会引发严重的显存爆炸问题:在长度为 $N$ 的序列上训练 $K$ 个并行组会产生 $N \times K$ 个总位置,导致注意力矩阵的计算量呈平方级增长(例如 65K 个位置的注意力矩阵在 bf16 精度下单次就需要 8GB 显存),而常规的跨样本梯度累加根本无法解决这种单条序列内部的显存溢出。
怎么做的 P-EAGLE 的核心思路是将 EAGLE 的草稿生成过程从自回归改造为完全并行,从而彻底解耦草稿数量与前向传播次数。该方法分为两个阶段。第一阶段是预填充,目标模型像常规推理一样处理提示词并生成新 token,同时 P-EAGLE 会捕获目标模型内部的隐藏状态(提示词位置的 $h_{prompt}$ 和新生成 token 的 $h_{context}$),这些状态编码了目标模型在每个位置的“知识”。第二阶段是 P-EAGLE 草稿模型,它在单次前向传播中并行构造所有位置的输入。每个输入由 token 词嵌入(embedding)和隐藏状态拼接而成。对于提示词位置,输入是当前 token 词嵌入与目标模型对应隐藏状态的拼接,并向后偏移一位以预测下一个 token。对于后续的投机位置,其核心输入构造机制如下:
在第 1 个预测位置(单 token 预测),输入是新生成 token 的词嵌入 $emb(new)$ 与 $h_{context}$ 的拼接;对于第 2 到 $K$ 个预测位置(多 token 预测),由于真实的词嵌入和隐藏状态尚未生成,P-EAGLE 引入了两个在训练期间学习到的固定向量作为中性占位符:掩码 token 词嵌入 $emb(mask)$ 和共享隐藏状态 $h_{shared}$。所有位置的输入被同时送入 $N$ 层 Transformer 和语言模型头,一次性输出 $K$ 个草稿 token。为了解决长序列并行训练时的显存卡点,作者引入了一种序列划分算法,将 $N \times K$ 的位置序列分割成连续的块,在保持跨块注意力依赖正确的前提下,实现同一序列内部的梯度累加。在 vLLM 的工程实现中,由于并行生成引入了掩码占位符,导致草稿批次形状与验证批次形状不再一致,作者编写了融合的 Triton 算子在 GPU 上直接复制和扩展目标批次、插入掩码并重建元数据,同时使用专用的拷贝算子将 $h_{shared}$ 广播到掩码槽位中,从而用极低的算子下发开销抵消了元数据重建的代价。
效果如何 实验在单张 NVIDIA B200 GPU 上基于 vLLM 框架进行,测试模型为 GPT-OSS-20B。对比的基线方法是公开的原生 EAGLE-3 检查点,它代表了传统的自回归投机解码路线。测试涵盖了三个基准:MT-Bench(多轮指令遵循)、SPEED-Bench(长上下文代码生成)和 HumanEval(函数级代码合成)。在并发度为 1 时,P-EAGLE 相比 EAGLE-3 实现了 55% 到 69% 的吞吐量(TPS)提升(其中 SPEED-Bench 达到最高 1.69 倍加速);在并发度高达 64 时,仍能保持 5% 到 25% 的性能优势。量化结果揭示了并行路线的核心优势:原生 EAGLE-3 在投机深度 $K=3$ 时吞吐量就达到了极限,而 P-EAGLE 在所有并发级别下都能在 $K=7$ 时达到吞吐量峰值。在 $K=7$ 时,P-EAGLE 的平均接受长度(AL,即每轮投机中被验证器接受的草稿 token 数量)显著更高,例如在 HumanEval 上达到 3.94,比基线的 3.03 高出 30%。该方法的代价是需要专门训练支持并行的草稿模型头,且在系统底层打破了草稿与验证阶段的 token 布局一致性,需要通过手写的融合算子来处理复杂的输入准备、注意力元数据管理和 KV 缓存槽位映射(例如将拒绝的 token 映射为 -1 以防止无效写入)。
EAGLE是大型语言模型(LLM)推理中推测解码(speculative decoding)的最先进方法。然而,其自回归的草稿生成方式产生了一个隐性瓶颈:推测的token越多,草稿生成器(drafter)需要的连续前向传播次数就越多。最终,这些开销会吞噬推测解码带来的性能收益。
为了解决这一问题,本文提出了P-EAGLE (Parallel-EAGLE),它通过在单次前向传播中生成所有 $K$ 个草稿token来打破这一上限。在真实的业务负载下,P-EAGLE在NVIDIA B200 GPU上相比原生的EAGLE-3实现了高达1.69倍的加速。用户只需下载或训练一个支持并行的草稿生成头,并在vLLM的服务流水线中添加 "parallel_drafting": true 配置,即可解锁这一性能提升。目前,预训练的P-EAGLE头已在HuggingFace上开源,支持的模型包括GPT-OSS 120B、GPT-OSS 20B和Qwen3-Coder 30B。该方法已从v0.16.0版本开始集成入vLLM中。
EAGLE的草稿生成瓶颈:EAGLE相比标准的自回归解码可以实现2-3倍的加速,并且已经被广泛部署在包括vLLM、SGLang和TensorRT-LLM在内的生产推理框架中。然而,EAGLE是以自回归的方式生成token草稿的。为了生成 $K$ 个草稿token,它需要通过草稿模型进行 $K$ 次前向传播。随着草稿模型在生成长序列输出方面变得越来越强大,这种草稿生成的开销变得非常显著——草稿生成器的延迟与推测深度呈线性比例扩展,这严重限制了系统进行激进推测的能力。
P-EAGLE架构概述:P-EAGLE的核心思想是将EAGLE从自回归的草稿生成转化为并行的草稿生成。P-EAGLE能够在单次前向传播中生成全部 $K$ 个草稿token。其架构主要包含以下两个步骤:
步骤1:预填充(Prefilling)
目标模型(target model)处理提示词(prompt)并像在正常推理过程中一样生成一个新的token。在此过程中,P-EAGLE会捕获模型内部的隐藏状态(hidden states):为每个提示词位置捕获 $h_{prompt}$,为新生成的token捕获 $h_{context}$。这些隐藏状态编码了目标模型在每个位置所“知道”的信息,并将在后续指导草稿模型的预测。这一步骤与自回归的EAGLE完全相同。
步骤2:P-EAGLE草稿生成器(Drafter)
草稿生成器并行地为每个位置构建输入。每个输入由一个token嵌入(embedding)与一个隐藏状态拼接而成。
在长序列上训练P-EAGLE的挑战与方案:
长序列匹配需求:现代推理模型通常会产生极长的输出。在UltraChat数据集上,GPT-OSS 120B生成的序列(包括提示词)中位长度为3,891个token,P90长度达到了10,800个token。草稿模型必须在匹配的上下文长度上进行训练,才能在推理时保持有效性。
内存放大问题:并行草稿生成在训练期间会极大地放大内存需求。在长度为 $N$ 的序列上训练 $K$ 个并行组会产生 $N \times K$ 个总位置。当 $N = 8192$ 且 $K = 8$ 时,单个训练样本将包含65,536个位置。注意力机制要求每个位置关注每一个有效位置—— $65K \times 65K$ 意味着需要计算超过40亿个元素,这在bf16精度下将消耗8GB的内存。
vLLM中的工程实现细节:
[token, MASK, MASK, …])。这些额外的位置仅为草稿生成而存在,导致草稿批次的形状不再与验证批次的形状匹配。由于无法重用验证元数据,系统必须重建批次元数据:扩展输入token ID、隐藏状态和位置以插入掩码token/嵌入的槽位,增加每个请求的位置,然后根据更新后的位置重新计算槽位映射(slot mapping)和每个请求的起始索引。# 将目标隐藏状态复制到它们的新位置
self.hidden_states[out_hidden_state_mapping] = target_hidden_states
# 用学习到的并行草稿隐藏状态填充掩码位置
mask = self.is_masked_token_mask[:total_num_output_tokens]
torch.where(
mask.unsqueeze(1),
self.parallel_drafting_hidden_state_tensor,
self.hidden_states[:total_num_output_tokens],
out=self.hidden_states[:total_num_output_tokens],
)
其中,\1 是从模型的 \1 缓冲区加载的,这是一个学习到的表示,用于告知模型这些位置应该预测未来的token。
PADDING_SLOT_ID (-1),以防止发生虚假的缓存写入。对于CUDA图(CUDA graphs),系统将捕获范围扩展了 $K \times max\_num\_seqs$,以容纳由并行草稿生成引入的更大的草稿批次。实验环境配置:
--max-num-seqs 1024, --max-model-len 100000, --max-cudagraph-capture-size 4096, 关闭前缀缓存和分块预填充(--no-enable-prefix-caching, --no-enable-chunked-prefill),设定 --kv-cache-dtype fp8 并开启异步调度。运行GPT-OSS-20B与EAGLE草稿生成器需应用vLLM补丁(PR#36684)。实验结果与分析:
实验对推测深度 $K \in {3,5,7}$ 和并发级别 $C \in {1,2,4,8,16,32,64}$ 进行了全面扫描,以测量每秒生成的token数(TPS),并在给定服务条件下选择最大化TPS的 $K$ 值。
吞吐量(TPS)全面提升:
P-EAGLE在低并发($c=1$)下提供了55–69%的更高吞吐量,在高并发($c=64$)下依然维持了5–25%的性能增益。
推测深度的扩展性优势:
数据显示出一个一致的模式:P-EAGLE在所有并发级别上均在 $K=7$ 时达到峰值TPS。相反,原生的EAGLE-3在 $K=3$ 时达到最高TPS(尽管在更高并发时其最佳深度偶尔会向更高值偏移)。这一行为反映了并行草稿生成的根本优势:P-EAGLE在单次前向传播中生成所有 $K$ 个草稿token,使其能够从更深的推测中获益,而不会产生额外的顺序开销。而自回归草稿生成器必须逐步生成推测token,这限制了它们有效扩展到更大 $K$ 值的极力。
更高的接受长度(Acceptance Length, AL):
除了减少草稿生成开销外,P-EAGLE的吞吐量增益还得益于更好的接受长度(即每轮推测中验证器平均接受的草稿token数)。更高的AL意味着更多的草稿工作转化为实际输出。
在相同推测深度 $K$ 下,P-EAGLE始终实现比EAGLE-3更高的AL。以 $K=7$ 为例:
P-EAGLE成功消除了推测解码中的顺序生成瓶颈,在真实工作负载上相比原生EAGLE-3实现了高达1.69倍的加速。通过将草稿生成的数量与前向传播的次数解耦,现在的系统可以探索更大规模的草稿架构,这甚至能够实现比单层基线模型更高的接受率。
在工程实现层面,P-EAGLE通过手写的融合内核(fused kernels)精心处理了输入准备、注意力元数据管理以及KV缓存槽位映射等底层复杂性。尽管该方法需要部署专门训练的模型,但其带来的显著性能优势使其成为vLLM推测解码功能的宝贵补充。随着未来有更多经过并行训练的模型开源可用,预计这种架构高效、基础设施稳健的方法将成为生产环境中LLM部署的首选路径,为追求最大化推理性能和降低延迟的系统提供明确的解决方案。