P-EAGLE: Faster LLM Inference with Parallel Speculative Decoding in vLLM

发表时间: 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。对于后续的投机位置，其核心输入构造机制如下：

$$\begin{aligned} Input_i = \begin{cases} emb(new) \oplus h_{context}, & i = 1 \\ emb(mask) \oplus h_{shared}, & 2 \le i \le K \end{cases} \end{aligned}$$


在第 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中。

图1：在单张NVIDIA B200显卡上，并发度为1时，P-EAGLE与SPEED-BENCH上其他方法的对比。
图1：在单张NVIDIA B200显卡上，并发度为1时，P-EAGLE与SPEED-BENCH上其他方法的对比。

背景知识与动机

EAGLE的草稿生成瓶颈：EAGLE相比标准的自回归解码可以实现2-3倍的加速，并且已经被广泛部署在包括vLLM、SGLang和TensorRT-LLM在内的生产推理框架中。然而，EAGLE是以自回归的方式生成token草稿的。为了生成 $K$ 个草稿token，它需要通过草稿模型进行 $K$ 次前向传播。随着草稿模型在生成长序列输出方面变得越来越强大，这种草稿生成的开销变得非常显著——草稿生成器的延迟与推测深度呈线性比例扩展，这严重限制了系统进行激进推测的能力。

P-EAGLE 核心方法细节

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）与一个隐藏状态拼接而成。

    • 提示词位置处理：输入将每个提示词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 $t_1, t_2, t_3$ 和 $t_4$。
图2：P-EAGLE架构概述。
图2：P-EAGLE架构概述。

在长序列上训练P-EAGLE的挑战与方案：

  • 长序列匹配需求：现代推理模型通常会产生极长的输出。在UltraChat数据集上，GPT-OSS 120B生成的序列（包括提示词）中位长度为3,891个token，P90长度达到了10,800个token。草稿模型必须在匹配的上下文长度上进行训练，才能在推理时保持有效性。
    图3：在UltraChat数据集上使用GPT-OSS 120B的序列长度（提示词 + 生成）分布。推理级别：中等。

  • 内存放大问题：并行草稿生成在训练期间会极大地放大内存需求。在长度为 $N$ 的序列上训练 $K$ 个并行组会产生 $N \times K$ 个总位置。当 $N = 8192$ 且 $K = 8$ 时，单个训练样本将包含65,536个位置。注意力机制要求每个位置关注每一个有效位置—— $65K \times 65K$ 意味着需要计算超过40亿个元素，这在bf16精度下将消耗8GB的内存。

  • 序列划分算法（Sequence Partition Algorithm）：位置采样技术【An等人，2025】通过随机跳过位置来减少内存，但跳过过于激进会降低草稿质量。梯度累积是内存受限训练的标准解决方案，但它是在不同的训练样本之间进行拆分的；当单个序列本身就超出内存时，则无法拆分。为此，P-EAGLE引入了一种用于序列内拆分的序列划分算法。该算法将 $N \times K$ 的位置序列划分为多个连续的块（chunks），在块边界之间保持正确的注意力依赖关系，并在同一序列的各个块之间累积梯度。

vLLM中的工程实现细节：

  • 并行草稿生成带来的批处理重构挑战：在许多推测解码设置中，草稿生成和验证阶段共享相同的每个请求（per-request）token布局。对于EAGLE而言，草稿生成器消耗的窗口已经与验证器将要检查的窗口相匹配（即 $K$ 个草稿token和1个额外的采样token）。然而，并行草稿生成打破了这种一致性。为了在草稿生成器的一次前向传播中预测 $K$ 个token，系统必须追加MASK占位符（例如，[token, MASK, MASK, …]）。这些额外的位置仅为草稿生成而存在，导致草稿批次的形状不再与验证批次的形状匹配。由于无法重用验证元数据，系统必须重建批次元数据：扩展输入token ID、隐藏状态和位置以插入掩码token/嵌入的槽位，增加每个请求的位置，然后根据更新后的位置重新计算槽位映射（slot mapping）和每个请求的起始索引。
  • 融合的Triton内核优化：为了抵消重建批次元数据带来的开销，P-EAGLE实现了一个融合的Triton内核。该内核通过在GPU上复制和扩展目标模型批次，直接在设备端填充草稿生成器的输入批次。在一次执行中，该内核将先前的token ID和位置从目标批次复制到新的目标槽位中，并插入目标模型采样的每个请求的额外token。接着，它用特殊的MASK token ID填充额外的并行草稿生成槽位。最后，它生成轻量级的元数据：一个拒绝token掩码、一个用于并行草稿生成槽位的掩码token掩码、用于采样草稿token的新token索引，以及一个隐藏状态映射。如果不进行融合，这些逻辑将需要多次GPU操作（复制/散布 + 插入 + 填充 + 掩码 + 重映射）。将其融合为一个内核显著减少了启动开销和额外的内存流量，从而保持了草稿生成设置的低成本。
  • 隐藏状态管理（Hidden State Management）：对于将隐藏状态传递给草稿模型的基于EAGLE的方法，并行草稿生成需要单独处理这些字段的填充。由于隐藏状态的体积明显大于输入批次的其余部分，P-EAGLE对工作进行了拆分：由Triton内核输出一个映射关系，然后由一个专用的复制内核将学习到的隐藏状态占位符广播（broadcast）到掩码token槽位中。具体实现代码如下：
# 将目标隐藏状态复制到它们的新位置
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。

  • KV Cache与CUDA图管理：在KV缓存槽位映射方面，有效的token接收正常的槽位分配，而拒绝的token则被映射到 PADDING_SLOT_ID (-1)，以防止发生虚假的缓存写入。对于CUDA图（CUDA graphs），系统将捕获范围扩展了 $K \times max\_num\_seqs$，以容纳由并行草稿生成引入的更大的草稿批次。

实验环境与结果

实验环境配置：

  • 评估基准与数据集：采用了三个基准测试：MT-Bench（用于多轮指令遵循评估）、SPEED-Bench Code（用于长期代码生成评估）以及HumanEval（用于函数级代码合成评估）。
  • 模型配置：目标模型为GPT-OSS-20B。P-EAGLE草稿生成器是一个轻量级的4层模型，经训练可并行预测最多10个token。基线方法为公开可用的原生EAGLE-3 checkpoint。两者均采用线性草稿生成。
  • 硬件配置：所有实验均在单张 NVIDIA B200 (Blackwell) GPU上进行。
  • 软件与系统配置：使用vLLM框架启动。核心配置参数包括 --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$ 值。

  1. 吞吐量（TPS）全面提升：
    P-EAGLE在低并发（$c=1$）下提供了55–69%的更高吞吐量，在高并发（$c=64$）下依然维持了5–25%的性能增益。

    • MT-Bench结果：P-EAGLE与EAGLE-3的加速比在不同并发下分别为：1.55x ($c=1$), 1.29x ($c=2$), 1.35x ($c=4$), 1.28x ($c=8$), 1.27x ($c=16$), 1.09x ($c=32$), 以及 1.05x ($c=64$)。详见图4。
    • HumanEval结果：加速比分别为：1.55x ($c=1$), 1.53x ($c=2$), 1.45x ($c=4$), 1.35x ($c=8$), 1.31x ($c=16$), 1.37x ($c=32$), 以及 1.23x ($c=64$)。详见图5。
    • Speed-bench结果：加速比分别为：1.69x ($c=1$), 1.61x ($c=2$), 1.54x ($c=4$), 1.45x ($c=8$), 1.40x ($c=16$), 1.22x ($c=32$), 以及 1.25x ($c=64$)。详见图6。

图4：在不同并发级别下，GPT-OSS-20B上P-EAGLE与EAGLE-3的MT-Bench吞吐量（TPS）。P/E加速比分别为：1.55x (c=1), 1.29x (c=2), 1.35x (c=4), 1.28x (c=8), 1.27x (c=16), 1.09x (c=32), 和 1.05x (c=64)。
图5：在不同并发级别下，GPT-OSS-20B上P-EAGLE与EAGLE-3的HumanEval吞吐量（TPS）。P/E加速比分别为：1.55x (c=1), 1.53x (c=2), 1.45x (c=4), 1.35x (c=8), 1.31x (c=16), 1.37x (c=32), 和 1.23x (c=64)。
图6：在不同并发级别下，GPT-OSS-20B上P-EAGLE与EAGLE-3的Speed-bench吞吐量（TPS）。P/E加速比分别为：1.69x (c=1), 1.61x (c=2), 1.54x (c=4), 1.45x (c=8), 1.40x (c=16), 1.22x (c=32), 和 1.25x (c=64)。

  1. 推测深度的扩展性优势：
    数据显示出一个一致的模式：P-EAGLE在所有并发级别上均在 $K=7$ 时达到峰值TPS。相反，原生的EAGLE-3在 $K=3$ 时达到最高TPS（尽管在更高并发时其最佳深度偶尔会向更高值偏移）。这一行为反映了并行草稿生成的根本优势：P-EAGLE在单次前向传播中生成所有 $K$ 个草稿token，使其能够从更深的推测中获益，而不会产生额外的顺序开销。而自回归草稿生成器必须逐步生成推测token，这限制了它们有效扩展到更大 $K$ 值的极力。

  2. 更高的接受长度（Acceptance Length, AL）：
    除了减少草稿生成开销外，P-EAGLE的吞吐量增益还得益于更好的接受长度（即每轮推测中验证器平均接受的草稿token数）。更高的AL意味着更多的草稿工作转化为实际输出。
    在相同推测深度 $K$ 下，P-EAGLE始终实现比EAGLE-3更高的AL。以 $K=7$ 为例：

    • 在HumanEval上，P-EAGLE比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在HumanEval上的AL增加了0.92（从3.02升至3.94），而EAGLE-3仅增加了0.38（从2.65升至3.03）。这种在更高 $K$ 值下不断扩大的差距，进一步验证了P-EAGLE单次传递并行草稿生成不会因更深推测而产生额外成本的理论。

结论与展望

P-EAGLE成功消除了推测解码中的顺序生成瓶颈，在真实工作负载上相比原生EAGLE-3实现了高达1.69倍的加速。通过将草稿生成的数量与前向传播的次数解耦，现在的系统可以探索更大规模的草稿架构，这甚至能够实现比单层基线模型更高的接受率。

在工程实现层面，P-EAGLE通过手写的融合内核（fused kernels）精心处理了输入准备、注意力元数据管理以及KV缓存槽位映射等底层复杂性。尽管该方法需要部署专门训练的模型，但其带来的显著性能优势使其成为vLLM推测解码功能的宝贵补充。随着未来有更多经过并行训练的模型开源可用，预计这种架构高效、基础设施稳健的方法将成为生产环境中LLM部署的首选路径，为追求最大化推理性能和降低延迟的系统提供明确的解决方案。