P-EAGLE: Parallel-Drafting EAGLE with Scalable Training

Mude Hui * 1 2 Xin Huang * 3 Jaime Campos Salas 3 Yue Sun 3 Nathan Pemberton 3 Xiang Song 3 Ashish Khetan 3 George Karypis 3

A1 主要贡献

本文主要解决的问题是大型语言模型（LLM）推理中的效率瓶颈，特别是针对推理能力模型（Reasoning LLMs）产生长输出序列的场景。现有的EAGLE方法虽然通过投机解码加速了推理，但其Drafting过程仍是自回归的（串行生成）。现有的并行Drafting方法（如ParallelSpec和PARD）在处理长上下文训练时，由于注意力机制的二次方复杂度，面临显存溢出（OOM）或计算不可行的问题。

主要贡献如下：

1. 长上下文的可扩展训练框架：提出了一套包含“摊销掩码构建（Amortized Mask Construction）”和“序列分块（Sequence Partitioning）”的框架。这解决了并行预测训练中注意力显存随序列长度和预测深度乘积呈二次方增长的问题，使得模型能够在长序列（如20K tokens）上进行有效训练。
2. 基于EAGLE的并行Drafting架构（P-EAGLE）：引入了一个可学习的共享隐藏状态（Shared Hidden State），将EAGLE从自回归生成转变为并行多token预测。理论分析证明，利用RoPE的位置编码已足够包含所需的位置信息，无需额外的位置特定隐藏状态。
3. 优化的训练配方：通过系统性的消融实验，确立了P-EAGLE的最佳实践，包括架构深度（推荐4层）、解冻Embedding层以及训练与推理时的预测深度对齐策略。
4. 生产级部署与验证：在vLLM中实现了P-EAGLE，并在GPT-OSS 120B/20B和Qwen3-Coder 30B上进行了评估。结果显示，相比于高度优化的自回归EAGLE-3，P-EAGLE实现了1.10×–1.36×的端到端加速。

Fig 1展示了推理类模型在UltraChat数据集上的长序列特征（中位数接近4K，P90超过10K tokens），强调了长上下文训练的必要性。

Table 1对比显示，ParallelSpec和PARD在长上下文（8K+）训练时面临OOM或不可行问题，而P-EAGLE可扩展至20K tokens。

A2 方法细节

2. P-EAGLE 架构

架构设计
P-EAGLE旨在消除自回归Drafting的开销。传统的EAGLE生成$K$个draft tokens需要$K$次串行前向传播，因为每一步都依赖上一步的token和隐藏状态。P-EAGLE通过并行化设计，在一次前向传播中生成所有$K$个tokens。
* 输入构建：
* 对于Next-Token Prediction (NTP) 位置（即Pos 1）：利用目标模型（Target Model）产生的真实隐藏状态（来自层索引2, L/2, L-1的拼接）和当前Token Embedding作为输入。这与标准EAGLE一致。
* 对于Multi-Token Prediction (MTP) 位置（即Pos 2-K）：由于缺乏前序步骤的真实输出，P-EAGLE引入了两个可学习参数来替代：
1. 共享隐藏状态（Shared Hidden State, $h_{shared}$）：替代缺失的前序隐藏向量。
2. Mask Token Embedding：替代未知的上一个Token。

• 处理流程：所有位置的输入（Embedding + 投影后的隐藏状态）被并行送入$N$层Transformer层（基于LLaMA 3架构，含RoPE），最后通过LM Head输出预测结果。
• 设计选择：解冻从目标模型继承的Token Embeddings，以便学习Mask Token的有意义表示；采用更深的网络结构（4层），以补偿并行预测带来的信息缺失。

Fig 2展示了P-EAGLE架构。上方为目标模型提供特征，下方Draft模型中，Pos 1使用真实特征，Pos 2-4使用可学习的共享隐藏状态$h_{shared}$进行并行预测。

3. 面向长上下文的可扩展训练框架

挑战分析
训练并行预测模型需要将长度为$n$的序列扩展以适应$K$个并行预测深度。未优化的情况下，总位置数为$n \times K$，注意力复杂度为$O((nK)^2)$，导致长序列OOM。虽然PARD方法引入了条件Drop-token (COD) 采样（按几何衰减率$r$保留后续深度位置）来减少总位置数，但其随机采样导致每个样本需要构建不同的注意力掩码（Mask），构建成本高达$O((nK)^2)$，在长序列下极其缓慢。

3.1 摊销掩码构建 (Amortized Mask Construction)
* 观察：跨预测深度的因果结构具有位置不变性（Position-invariance）。无论序列总长度如何，位置0到$n$的注意力模式是相同的。
* 解决方案：
* 在训练初始化阶段，一次性预计算最大序列长度的注意力掩码。
* 在训练过程中，针对每个样本，通过张量切片（Tensor Slicing） 操作直接获取所需的掩码子矩阵（即预计算掩码的左上角子矩阵）。
* 该操作是常数时间复杂度的视图操作，无需额外内存分配。

• 效果：消除了每样本构建掩码的开销。如Table 2所示，在2048 tokens下，将PARD的数据加载速度提升了48倍。

Fig 3展示了如何通过切片预计算的大掩码来获得短序列的掩码，利用了注意力模式的位置不变性。

3.2 序列分块 (Sequence Partitioning)
* 挑战：即使掩码构建加速了，显存占用仍随序列增长。对于8192长度、$K=8$、保留率$r=0.8$的序列，总位置数约3.4万，注意力显存呈$O(L^2)$增长。单张卡无法容纳整个序列进行前向传播。
* 问题：传统的梯度累积将Batch拆分为Micro-batch，但在单条序列超显存时无效。必须将单条序列拆分。然而，COD的随机采样结构破坏了简单的索引切分——位置$p$（在深度$d$）必须关注位置$p-1$（在深度$d-1$），随机采样可能导致依赖对分布在不同切片中，破坏因果注意力。
* 解决方案算法：提出了一种迭代式的序列分块技术。
1. 初始化：根据显存限制设定分块数量$S$和边界。
2. 深度0-1分配：对于深度0和1的位置，直接根据其索引分配到对应的分块段中。
3. 依赖传播（深度$\ge$2）：对于深度$d \ge 2$的每个位置$p$，将其分配到与其依赖项（深度$d-1$的位置$p-1$）相同的分块段中。这种迭代传播保证了所有跨深度的注意力依赖都在同一分块内。
4. 因果补全：为了满足因果注意力，每个分块段额外包含直到其边界的所有深度0（NTP）位置。

• 结果：通过将序列分为$S$段，峰值注意力显存从$O(L^2)$降低到$O(L^2/S^2)$，实现了“序列内梯度累积”，同时保留了所有跨深度的注意力依赖。

Fig 4展示了直接按索引分块会破坏依赖关系（红色箭头跨越虚线），而所提算法能确保依赖项在同一分块中。

A3 背景知识/关键Observation/设计原则

在方法细节之前，论文阐述了关于隐藏状态设计的关键理论观察。
* 背景：自回归解码效率低，投机解码（如EAGLE）利用目标模型特征层作为Draft模型的输入，比独立训练的小模型更高效。
* 关键Observation：在并行预测中，是否需要为不同的预测深度（Predicting 1-step ahead, 2-steps ahead...）提供特定的隐藏状态编码？
* 理论分析表明，不需要。RoPE（旋转位置编码）已经编码了绝对位置信息。由于预测深度$g$是序列位置$p$的确定性函数（在COD结构中，每个位置映射到一个深度），且绝对位置$p$可以从RoPE注意力分数中唯一恢复（见附录B证明），因此引入显式的深度编码或注入NTP上下文是信息冗余的。
* 基于此原则，P-EAGLE采用了最简单的共享隐藏状态（Shared Hidden State） 设计，即所有MTP位置共享同一个可学习向量，而非区分深度的向量。这不仅简化了模型，而且实验证明比复杂的上下文注入方案性能高出7-15%。

A4 实验环境

• 数据集：

  • 训练集：UltraChat (对话), GSM-8K (数学推理), OpenCodeInstruct (代码)。
  • 评估集：HumanEval (代码生成), MT-Bench (多轮对话), GSM-8K Test (数学)。

• 模型配置：

  • 目标模型：GPT-OSS 120B, GPT-OSS 20B [OpenAI, 2025], Qwen3-Coder 30B [Team, 2025]。
  • P-EAGLE Drafter：4层Decoder Layer（基于LLaMA 3.1结构），Embedding层解冻。

• 训练参数：

  • 最大序列长度：8192 tokens。
  • 并行预测组数（深度）：$K_{train} = 8$。
  • COD下采样率：0.8。
  • Batch Size：8（配合序列分块和梯度累积）。
  • 硬件：8× H200 GPUs。

• 基线：自回归（AR）EAGLE-3（单层，使用HCA损失优化，是一个强基线）。由于ParallelSpec和PARD无法在8K+长上下文下训练，故不作为主要对比对象。

A4 实验结果

1. 接受长度 (Acceptance Length) 对比

实验旨在验证并行Drafting是否会降低生成质量。结果显示P-EAGLE在质量上不输于甚至略优于自回归EAGLE。
* 总体表现：在所有9个模型-数据集组合中，P-EAGLE的平均每步接受长度均匹配或超过AR EAGLE-3。
* GPT-OSS 120B：平均提升 +4.5%。
* GPT-OSS 20B：平均提升 +2.5%。
* Qwen3-Coder 30B：平均提升 +2.0%。即使在Qwen3-Coder这种AR基线已经非常强（HumanEval上接受长度达4.36/5）的情况下，P-EAGLE仍实现了提升。

• 图表引用：详见 Table 9。

2. 端到端吞吐量 (vLLM)

在vLLM框架中实测输出tokens每秒 (OTPS)。
* 加速效果：P-EAGLE利用单次前向传播生成更多tokens（$K=5 \sim 7$），相比AR EAGLE（通常$K=3$最优）减少了迭代次数。
* 并发度 C=2：加速比最为显著。GPT-OSS 20B达到 1.27×–1.36×，120B为 1.04×–1.10×，Qwen 30B为 1.04×–1.17×。
* 并发度 C=4：20B模型仍保持高加速（1.24×–1.27×），大模型加速比因验证阶段成为瓶颈而略有下降。

• 图表引用：详见 Table 10。

3. 消融实验结论 (基于LLaMA 3.1 8B)

• 隐藏状态设计：Table 3 显示，简单的共享隐藏状态优于所有复杂的上下文注入/深度编码变体（性能高7-15%）。
• 模型容量：Table 4 显示，从1层增加到4层，HumanEval上的接受长度提升了45.7%，证明并行预测需要更大的模型容量来补偿信息缺失。
• Embedding解冻：Table 5 显示，解冻Embedding层带来约5%的提升，使Mask Token能学到有意义的先验。
• 训练深度：Table 6 显示，训练时使用更长的预测深度（$K=8$）比训练推理一致（$K=5$）能带来+4%的提升。
• 训练时长与序列长度：延长训练epoch和增加训练序列长度均能进一步提升性能（Table 7, Table 8）。

A5 结论

本文提出了P-EAGLE，成功将EAGLE投机解码框架从串行自回归转变为并行多token预测。针对并行训练中长上下文导致的显存瓶颈，提出了一套包含摊销掩码构建和序列分块的可扩展训练框架。通过在vLLM上的实现，P-EAGLE在多个大规模模型（最高120B）上展现了相对于强基线AR EAGLE-3的显著加速（1.10×–1.36×），且不牺牲生成质量。这证明了并行Drafting在生产级LLM加速中的可行性和优势。未来工作可结合更长上下文的优化技术。

A6 附录

附录 A: ParallelSpec 和 PARD 的训练配置

• 这两个对比方法在8张H200 GPU上训练，全局Batch Size 64，8步梯度累积。目标模型为GPT-OSS 120B。因长上下文OOM问题，主要用于Table 1的可扩展性对比。

附录 B: 隐藏状态增强冗余性的理论证明

• 核心论点：证明在MTP位置增强共享隐藏状态（如添加深度Embedding或注入NTP上下文）是冗余的。
• 证明逻辑：
  1. 引用Liu & Zhou (2025) 证明从相对位置$\delta$到RoPE旋转矩阵$R_{\delta}$的映射是单射。
  2. 定理 B.3：证明注意力分数函数$f_{q,k}(\delta) = q^T R_{\delta} k$在几乎所有情况下（勒贝格测度为零的集合除外）关于相对位置$\delta$是单射。
  3. 推论 B.4：绝对位置$p$可以从RoPE注意力分数中唯一恢复。
  4. 结论：由于预测深度$g$是位置$p$的函数（$g(p)$），既然$p$可被恢复，模型即可推断出深度$g$。因此，显式的深度Embedding $e(g)_{depth}$ 不提供额外信息，反而可能引入冗余干扰优化。

附录 C: 2层 vs 4层 P-EAGLE 对比

• 权衡：2层模型延迟更低，4层模型接受率更高。
• 数据：Table 11 显示，2层P-EAGLE在GPT-OSS 20B上性能下降明显（-12.4%），而4层模型实现提升（+2.5%）。但在Qwen和120B上，2层模型仍具有竞争力（仅下降3-6%），可作为低延迟场景的备选。

A7 补充细节

4.1 隐藏状态消融细节

• 为了验证共享隐藏状态的优越性，作者设计了多种变体：

  • Baseline：仅使用单一$h_{shared}$。
  • + depth-specific encoding：添加深度特定的Embedding。
  • + NTP hidden：注入投影后的NTP隐藏状态。
  • + regularized NTP hidden：添加带Dropout和可学习缩放系数$\alpha$的NTP上下文：$h_{MTP} = h_{shared} + \alpha \cdot \text{dropout}(\text{proj}(h_{NTP}))$。

• 实验发现：所有变体均不如Baseline（见Table 3）。

• 深入分析：对于regularized variant，作者追踪了可学习参数$\alpha$的轨迹。如图5所示，$\alpha$从初始值0.1在训练过程中指数衰减至0.029（下降71%），且此时MTP准确率仍低于Baseline。这直接证明了模型在主动“关闭”上下文注入，因为这种冗余信息损害了性能。

Table 3数据表明，任何形式的显式上下文注入或深度编码都会导致性能下降。

Fig 5显示模型在训练中主动减小注入系数$\alpha$，证实了上下文注入的负面影响。