P-EAGLE：通过 vLLM 中的并行投机采样实现更快的 LLM 推理

EAGLE 是目前大语言模型 (LLM) 推理中用于推测解码的最先进方法，但其自回归的草稿生成方式制造了一个隐蔽的瓶颈：你推测的 token 越多，草稿模型所需的顺序前向传递就越多。最终，这些开销会抵消掉你的性能提升。P-EAGLE 通过在单次前向传递中生成所有 K 个草稿 token 移除了这一上限，在 NVIDIA B200 上的实际工作负载中，相比原生 EAGLE-3 可带来高达 1.69 倍的加速。

你可以通过下载（或训练）一个支持并行的草稿头，并在你的 vLLM 服务流水线中添加 "parallel_drafting": true 来开启这一性能提升。预训练的 P-EAGLE 草稿头已在 HuggingFace 上提供，适用于 GPT-OSS 120B、GPT-OSS 20B 以及 Qwen3-Coder 30B，你可以立即开始使用！

在这篇文章中，我们将解释 P-EAGLE 的工作原理、我们如何从 v0.16.0 (PR#32887) 开始将其集成到 vLLM 中，以及如何使用我们的预训练检查点进行服务。以下是所使用的工件列表：

• ArXiv 论文
• HuggingFace 模型 (GPT-OSS 120B, GPT-OSS 20B, Qwen3-Coder-30B-A3B-Instruct)
• vLLM 集成 统一并行草稿
• vLLM-Speculators (RFC, PR)

P-EAGLE 快速入门

你可以通过在 SpeculativeConfig 类中进行单一配置更改来启用并行草稿。

# vllm/config/speculative.py
   parallel_drafting: bool = True

这是在 vLLM 中使用 P-EAGLE 作为草稿模型启用并行草稿的示例命令：

  vllm serve openai/gpt-oss-20b \
   --speculative-config '{"method": "eagle3", "model": "amazon/gpt-oss-20b-p-eagle", "num_speculative_tokens": 5, "parallel_drafting": true}'

EAGLE 的草稿瓶颈

EAGLE 相比标准自回归解码可实现 2-3 倍的加速，并已广泛部署在 vLLM、SGLang 和 TensorRT-LLM 等生产推理框架中。EAGLE 以自回归方式起草 token。为了生成 K 个草稿 token，它需要对草稿模型进行 K 次前向传递。随着草稿模型在撰写长输出方面变得更好，这种草稿开销变得显著——草稿模型的延迟随推测深度线性增加，限制了我们进行推测的激进程度。

我们的方法：Parallel-EAGLE (P-EAGLE)

我们提出了 P-EAGLE，它将 EAGLE 从自回归草稿生成转变为并行草稿生成。在 B200 GPU 上，P-EAGLE 在 MT-Bench、HumanEval 和 SpeedBench 上的 GPT-OSS 20B 模型上实现了相比原生 EAGLE-3 1.05 倍至 1.69 倍的加速。它现已集成到 vLLM 中，以解锁并行推测解码，并准备好加速实际部署。

P-EAGLE 在单次前向传递中生成 K 个草稿 token。图 2 展示了其架构，包括两个步骤。

第 1 步：预填充 (Prefilling)。 目标模型处理提示词并生成一个新 token，正如在正常推理过程中那样。在此过程中，P-EAGLE 捕获模型的内部隐藏状态：每个提示位置的 h_prompt，以及新生成 token 的 h_context。这些隐藏状态编码了目标模型在每个位置“知道”的内容，并将引导草稿模型的预测。此步骤与自回归 EAGLE 完全相同。

第 2 步：P-EAGLE 草稿器。 草稿器并行构建每个位置的输入。每个输入由 token 嵌入与隐藏状态拼接而成。

对于提示位置，输入将每个提示 token 嵌入 emb(p) 与其来自目标模型的相应 h_prompt 配对。遵循与自回归 EAGLE 相同的约定，位置偏移一位。位置 i 接收来自位置 i-1 的 token 和隐藏状态，使其能够预测位置 i 的 token。

对于位置 1，即下一 Token 预测 (NTP)，输入将新生成的 token 嵌入 emb(new) 与 h_context 配对。该位置的操作与标准的自回归 EAGLE 完全相同。对于位置 2 到 K，即多 Token 预测 (MTP)，所需的输入（token 嵌入和隐藏状态）尚不存在。P-EAGLE 用两个可学习参数填充这些位置：共享掩码 token 嵌入 emb(mask) 和共享隐藏状态 h_shared。这些是训练期间学习到的固定向量，充当中性占位符。

所有位置共同通过 N 个 Transformer 层，然后通过语言模型头，在单次前向传递中预测出草稿 token t1、t2、t3 和 t4。

在长序列上训练 P-EAGLE

现代推理模型会产生长输出。如图 3 所示，GPT-OSS 120B 在 UltraChat 数据集上生成的序列（包括提示词）中位数长度为 3,891 个 token，P90 为 10,800 个 token。为了在推理时有效，草稿模型必须在匹配的上下文长度上进行训练。

一个关键挑战是并行草稿会放大训练期间的内存需求。在长度为 N 的序列上训练 K 个并行组会产生 N × K 个总位置。当 N = 8,192 且 K = 8 时，单个训练示例包含 65,536 个位置。注意力机制要求每个位置关注所有有效位置——65K × 65K 意味着超过 40 亿个元素，以 bf16 格式消耗 8GB 内存。

位置采样 [An et al., 2025] 通过随机跳过位置来减少内存，但跳过过于激进会降低草稿质量。梯度累积是内存受限训练的标准解决方案，但它是在不同的训练示例之间进行分割的。当单个序列超出内存限制时，便无法再进行分割。

P-EAGLE 引入了一种用于序列内分割的序列划分算法。该算法将 N × K 个位置序列划分为连续的块，保持跨块边界的正确注意力依赖关系，并跨同一序列的块累积梯度。详细信息请参阅 P-EAGLE 论文。

vLLM 中的实现

并行草稿的挑战

在许多推测解码设置中，起草和验证共享相同的每个请求 token 布局。对于 EAGLE 来说基本如此：草稿器消耗的窗口已经与验证器将要检查的内容相匹配；K 个草稿 token 和一个额外的采样 token。

并行草稿破坏了这种一致性。为了在一次草稿器前向传递中预测 K 个 token，我们附加了 MASK 占位符（例如 [token, MASK, MASK, …]）。这些额外位置仅用于起草，因此草稿批次形状不再与验证批次形状匹配。因为我们无法重用验证元数据，所以必须重建批次元数据。我们扩展了输入 token ID、隐藏状态和位置以插入掩码 token/嵌入槽，增加每个请求的位置，然后从更新的位置重新计算槽映射和每个请求的起始索引。

Triton 内核

为了抵消重建批次元数据的开销，我们实现了一个融合的 Triton 内核，通过复制和扩展目标模型批次，在 GPU 上填充草稿器的输入批次。在单次传递中，内核将目标批次中的先前 token ID 和位置复制到新的目标槽中，并插入由目标模型采样的每个请求的额外 token。然后，它用特殊的 MASK token ID 填充额外的并行草稿槽。最后，它生成轻量级元数据：拒绝 token 掩码、用于并行草稿槽的掩码 token 掩码、用于采样草稿 token 的新 token 索引以及隐藏状态映射。

否则，此逻辑将需要多次 GPU 操作（复制/分散 + 插入 + 填充 + 掩码 + 重映射）。将其融合到一个内核中可减少启动开销和额外的内存流量，从而保持较低的起草设置成本。

隐藏状态管理

对于将隐藏状态传递给草稿模型的 EAGLE 类方法，并行草稿会分别处理这些字段的填充。由于隐藏状态明显大于输入批次的其余部分，我们将工作分开：Triton 内核输出一个映射，而专用的复制内核将学习到的隐藏状态占位符广播到掩码 token 槽中。

# Copy target hidden states to their new positions
 self.hidden_states[out_hidden_state_mapping] = target_hidden_states
 
# Fill masked positions with the learned Parallel Drafting hidden state
 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],
 )

parallel_drafting_hidden_state_tensor 是从模型的 mask_hidden 缓冲区加载的，这是一种学习到的表示，告诉模型这些位置应该预测未来的 token。

对于 KV 缓存槽映射，有效 token 接收正常的槽分配，而被拒绝的 token 被映射到 PADDING_SLOT_ID (-1) 以防止虚假的缓存写入。对于 CUDA 图，我们将捕获范围扩展了 K × max_num_seqs，以适应并行草稿引入的更大草稿批次。

P-EAGLE 的 vLLM 基准测试

我们在 GPT-OSS-20B 上训练了 P-EAGLE，并在三个基准测试中进行了评估：用于多轮指令遵循的 MT-Bench、用于长期代码生成的 SPEED-Bench Code，以及用于函数级代码合成的 HumanEval。与公开可用的 原生 EAGLE-3 检查点 相比，P-EAGLE 在低并发（c=1）时实现了 55-69% 的吞吐量提升，在高并发（c=64）时保持 5-25% 的增益。结果如图 4-6 所示。

P-EAGLE 草稿器是一个轻量级的 4 层模型，训练用于并行预测最多 10 个 token。为了评估性能，我们在并发级别 C ∈ {1,2,4,8,16,32,64} 下扫描推测深度 K ∈ {3,5,7}。我们的目标是为 P-EAGLE 和原生 EAGLE-3 确定正确的部署配置。线性起草同时用于 P-EAGLE 和原生 EAGLE-3。在这种情况下，“最佳 P-EAGLE”和“最佳 EAGLE-3”指的是在给定服务条件下实现峰值吞吐量的配置。它们以每秒 token 数 (TPS) 来衡量，对于给定的推测深度 K。对于每种方法，我们选择在给定服务条件下使 TPS 最大化的 K。

一个一致的模式出现了。P-EAGLE 在所有并发级别下均在 K=7 时达到峰值 TPS。相比之下，原生 EAGLE-3 在 K=3 时达到最高 TPS，其改进深度偶尔会根据并发情况向更高值移动。这种行为反映了并行起草的一个根本优势。P-EAGLE 在单次前向传递中生成所有 K 个草稿 token，使其能够受益于更深层的推测而无需承担额外的顺序开销。相比之下，自回归草稿器必须一步步生成推测 token，这限制了它们有效扩展到更大 K 的能力。

所有实验均在单张 NVIDIA B200 (Blackwell) GPU 上使用 vLLM 进行，采用以下服务配置。

VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8=1 \
vllm serve openai/gpt-oss-20b \
    --speculative-config '{
      "method": "eagle3",
      "model": "amazon/GPT-OSS-20B-P-EAGLE",
      "num_speculative_tokens": 7,
      "parallel_drafting": true}' \
    --port 8000 \
    --max-num-seqs 1024 \
    --max-model-len 100000 \
    --max-num-batched-tokens 100000 \
    --max-cudagraph-capture-size 4096 \
    --no-enable-prefix-caching \
    --no-enable-chunked-prefill \
    --kv-cache-dtype fp8 \
    --async-scheduling \
    --stream-interval 20

注意。 使用 EAGLE 草稿器为 GPT-OSS-20B 提供服务目前需要一个 vLLM 的单行补丁 (PR#36684)。请在启动前应用它。此修复预计将在即将发布的 vLLM 版本中合并。

除了减少草稿开销外，P-EAGLE 的吞吐量增益还得益于更好的接受长度 (AL)，即验证器在每轮推测中平均接受的草稿 token 数量。更高的 AL 意味着更多的草稿工作转化为实际输出，这直接提升了有效的 OTPS/TPS。

下表比较了 P-EAGLE 和原生 EAGLE-3 在 GPT-OSS-20B 上三个基准测试的 AL。

P-EAGLE (AL)

EAGLE-3 (AL)

在相同的推测深度 K 下，P-EAGLE 始终实现比 EAGLE-3 更高的 AL。在 K=7 时，P-EAGLE 在 HumanEval 上比 EAGLE-3 高出 30%（3.94 vs 3.03），在 SPEED-Bench 上高出 31%（3.38 vs 2.59），在 MT-Bench 上高出 13%（3.70 vs 3.27）。值得注意的是，P-EAGLE 从更深层的推测中获益更多。从 K=3 增加到 K=7，P-EAGLE 的 AL 在 HumanEval 上增加了 0.92（3.02 到 3.94），而 EAGLE-3 仅增加了 0.38（2.65 到 3.03）。在较高的 K 值下，这种差距进一步拉大，这与 P-EAGLE 的单次并行起草方式相一致，因为更深层的推测不会导致额外的成本。

结果复现

启动服务器后，使用 vllm bench serve 运行基准测试。

#MT-Bench
export MODEL="openai/gpt-oss-20b"
export BASE_URL="https://:8000"
vllm bench serve \
    --dataset-name hf \
    --dataset-path philschmid/mt-bench \
    --num-prompts 80 \
    --max-concurrency 1 \
    --model $MODEL \
    --base-url $BASE_URL \
    --temperature 0.0 \
    --hf-output-len 2048
 
#HumanEval command:
#Download HumanEval dataset openai/openai_humaneval
vllm bench serve \
    --dataset-name custom \
    --dataset-path <dataset path> \
    --num-prompts 164 \
    --max-concurrency 1 \
    --model $MODEL \
    --base-url $BASE_URL \
    --temperature 0.0 \
    --custom-output-len 2048

总结

P-EAGLE 移除了推测解码中的顺序瓶颈，在实际工作负载中相比原生 EAGLE-3 实现了高达 1.69 倍的加速。通过将草稿数量与前向传递次数解耦，我们现在可以探索更大的草稿架构，与单层基线相比，这甚至可以实现更高的接受率。该实现通过手工编写的融合内核仔细处理了输入准备、注意力元数据管理和 KV 缓存槽映射的复杂性。虽然它需要专门训练的模型，但其带来的性能优势使其成为 vLLM 推测解码能力的宝贵补充。

随着更多并行训练模型的可用，我们预计这种方法将成为生产 LLM 部署的首选。P-EAGLE 的架构效率与 vLLM 的稳健基础设施相结合，为追求最大推理性能和降低延迟的用户提供了一条清晰的路径。

立即尝试：从 HuggingFace 下载预训练的 P-EAGLE 头，在你的 vLLM 配置中为任何支持的模型设置 "parallel_drafting": true，亲眼见证速度提升吧。

致谢

AWS: Xin Huang, Florian Saupe, Jaime Campos Salas, Ashish Khetan, George Karypis

NVIDIA: Benjamin Chislett, Max Xu, Zeyuan (Faradawn) Yang, Kaihang Jiang, Xin Li, Omri Almog

我们也特别感谢维护人员和 vLLM 社区提供的审核、指导以及在此功能构建之上的卓越基础设施。

也发布于 AWS 博客。