P-EAGLE：vLLM 中的并行投机解码

TL;DR

P-EAGLE 是一种并行投机解码方法，已集成到 vLLM（从 v0.16.0 起），它消除了自回归草稿的顺序瓶颈。通过在一次前向传播中生成全部 K 个草稿标记，在使用 NVIDIA B200 GPU 的真实工作负载上相较于原生 EAGLE-3 可实现最高 1.69 倍的加速。

解决自回归草稿瓶颈

标准的投机解码方法（如 EAGLE）以自回归方式草稿标记，这意味着生成 K 个草稿标记需要对草稿模型进行 K 次顺序前向传播。这样会导致延迟随投机深度线性增长，限制了系统在不让草稿开销侵蚀整体性能提升的情况下进行激进投机的能力。

P-EAGLE 通过将草稿过程从自回归转为并行，打破了这一上限。它不再采用顺序步骤，而是在一次前向传播中生成完整的草稿标记集合，使草稿标记的数量与所需前向传播次数解耦。

P-EAGLE 架构与机制

P-EAGLE 通过两个主要步骤生成草稿标记：

步骤 1：Prefilling

目标模型处理提示并生成一个新标记。在此过程中，P-EAGLE 捕获内部隐藏状态：每个提示位置的 h_prompt 与新生成标记的 h_context。这些隐藏状态用于指导草稿模型的预测。

步骤 2：Parallel Drafting

草稿模型并行构建每个位置的输入，使用标记嵌入和隐藏状态的组合：

  • 提示位置： 每个提示标记嵌入 emb(p) 与目标模型对应的 h_prompt 配对，并向后移一位，以预测位置 $i$ 的标记。
  • 下一标记预测（NTP）： 第一个位置将新生成的标记嵌入 emb(new) 与 h_context 配对。
  • 多标记预测（MTP）： 对于位置 2 到 K，尚不存在标记嵌入和隐藏状态时，P-EAGLE 使用两个可学习参数：共享的掩码标记嵌入 emb(mask) 和共享隐藏状态 h_shared，作为中性占位符。

所有位置随后通过 N 层 Transformer 和语言模型头部，同时预测草稿标记 $t_1$ 到 $t_K$。

长序列上的训练

并行草稿在训练期间会增加内存需求，因为在长度为 N 的序列上训练 K 组并行会产生 $N \times K$ 个总位置。例如，$N=8,192$ 且 $K=8$ 时，单个训练样本包含 65,536 个位置，导致注意力矩阵规模超过 40 亿元素。

为了解决此问题，P-EAGLE 引入了一种 序列划分算法 用于序列内部拆分。该算法将 $N \times K$ 位置序列划分为连续块，同时保持跨块的正确注意力依赖，并在同一序列的不同块之间累计梯度。

vLLM 实现细节

将并行草稿集成到 vLLM 需要克服与批次元数据和内存管理相关的若干技术挑战：

融合 Triton Kernel

并行草稿由于引入了 MASK 占位符，会打破草稿与验证批次形状的一致性。为避免通过多次 GPU 操作重建批次元数据的开销，vLLM 实现了一个融合的 Triton kernel。该 kernel 在一次传递中完成复制先前标记 ID 与位置、插入奖励标记、用 MASK 标记 ID 填充并行草稿槽位，以及生成必要的元数据（被拒标记掩码、掩码标记掩码以及隐藏状态映射）。

隐藏状态管理

由于隐藏状态远大于标记 ID，vLLM 使用专用的复制 kernel 将学习得到的隐藏状态占位符（parallel_drafting_hidden_state_tensor）广播到 MASK 标记槽位，同时将目标隐藏状态映射到其新位置。

KV 缓存与 CUDA Graphs

  • KV 缓存： 将被拒标记映射为 PADDING_SLOT_ID（-1），以防止产生错误的缓存写入。
  • CUDA Graphs： 捕获范围扩展至 $K \times \text{max_num_seqs}$，以容纳更大的草稿批次规模。

性能基准

在 NVIDIA B200 GPU 上对 GPT-OSS-20B 的评估表明，P-EAGLE 在多个基准（MT-Bench、HumanEval、SPEED-Bench）上显著优于原生 EAGLE-3。

吞吐量提升

在低并发 (c=1) 时，P-EAGLE 提供 55%–69% 更高的吞吐量；在高并发 (c=64) 时，仍保持 5%–25% 的提升。具体相对于 EAGLE-3 的加速比为：

  • SPEED-Bench： 最高 1.69 倍（c=1）
  • HumanEval： 最高 1.55 倍（c=1）
  • MT-Bench： 最高 1.55 倍（c=1）

接受长度 (AL)

P-EAGLE 在每轮的平均接受草稿标记数 (AL) 上高于 EAGLE-3。在投机深度 $K=7$ 时，HumanEval 上的 AL 提高了 30%（3.94 对 3.03），SPEED-Bench 上提高了 31%（3.38 对 2.59）。

投机深度效率

不同于自回归草稿器在 $K=3$ 时达到吞吐峰值，P-EAGLE 在所有并发水平下始终在 $K=7$ 时实现峰值吞吐。这证明 P-EAGLE 能在更深的投机深度下获益，而不会遭受顺序草稿带来的线性延迟惩罚。

部署

在 vLLM 中通过在 SpeculativeConfig 类中设置 "parallel_drafting": true 即可启用并行草稿。以下模型的预训练 P-EAGLE 头已在 HuggingFace 上提供：

  • GPT-OSS 120B
  • GPT-OSS 20B
  • Qwen3-Coder 30B

Sources