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 以自回归方式生成草稿。要生成 K 个草稿 token，它需要对草稿模型进行 K 次前向传播。随着草稿模型在长文本生成方面表现越来越好，这种草稿生成开销变得显著——草稿模型的延迟随投机深度线性增长，限制了我们投机策略的激进程度。

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

我们提出了 P-EAGLE，它将 EAGLE 从自回归式转变为并行草稿生成。在 B200 GPU 上，P-EAGLE 在 GPT-OSS 20B 模型的 MT-Bench、HumanEval 和 SpeedBench 测试中，相比原生 EAGLE-3 实现了 1.05 倍至 1.69 倍的加速。它现已集成到 vLLM 中，支持并行投机解码，随时可用于加速实际部署。

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

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

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

对于 prompt 位置，输入将每个 prompt token 的嵌入 emb(p) 与来自目标模型的对应 h_prompt 配对。遵循与自回归 EAGLE 相同的惯例，位置向后移动一位。位置 i 接收来自 i-1 位置的 token 和隐藏状态，从而能够预测位置 i 的 token。

对于第一个位置（Next-Token-Prediction, NTP），输入将新生成的 token 嵌入 emb(new) 与 h_context 配对。该位置的运作方式与标准的自回归 EAGLE 完全一致。对于第 2 到第 K 个位置（Multi-Token-Prediction, MTP），所需的输入（token 嵌入和隐藏状态）尚不存在。P-EAGLE 使用两个可学习参数填充这些位置：共享掩码 token 嵌入 emb(mask) 和共享隐藏状态 h_shared。这些是训练过程中学到的固定向量，作为中性占位符。

这些位置一起通过 N 层 Transformer，然后通过语言模型头，从而在单次前向传播中预测出草稿 token t1、t2、t3 和 t4。

在长序列上训练 P-EAGLE

现代推理模型倾向于产生长输出。如图 3 所示，GPT-OSS 120B 在 UltraChat 数据集上生成的序列（包括 prompt）中位数长度为 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 代码测试，以及用于函数级代码合成的 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 倍的加速。通过将草稿 token 数与前向传播次数解耦，我们现在可以探索更大的草稿架构，这甚至能实现比单层基线更高的接受率。这种实现通过手写的融合内核，仔细处理了输入准备、注意力元数据管理和 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 博客。