Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

P-EAGLE: Parallel-Drafting EAGLE with Scalable Training

Mude Hui, Xin Huang, Jaime Campos Salas, Yue Sun, Nathan Pemberton, Xiang Song, Ashish Khetan, George Karypis

TL;DR

P-EAGLE addresses the bottleneck of autoregressive decoding in reasoning-enabled LLMs by converting EAGLE-style drafting into parallel multi-token prediction using a learnable shared hidden state and mask-token strategy. It introduces a scalable training framework with amortized mask construction and sequence partitioning to manage the quadratic attention memory growth that arises with long contexts. The approach is implemented in vLLM and demonstrates consistent end-to-end speedups of 1.10×–1.36× across GPT-OSS 120B, 20B, and Qwen3-Coder 30B while maintaining comparable acceptance lengths to autoregressive baselines. These results establish parallel drafting as a viable production technique for accelerating long-context reasoning in large language models.

Abstract

Reasoning LLMs produce longer outputs, requiring speculative decoding drafters trained on extended sequences. Parallel drafting - predicting multiple tokens per forward pass - offers latency benefits over sequential generation, but training complexity scales quadratically with the product of sequence length and parallel positions, rendering long-context training impractical. We present P(arallel)-EAGLE, which transforms EAGLE from autoregressive to parallel multi-token prediction via a learnable shared hidden state. To scale training to long contexts, we develop a framework featuring attention mask pre-computation and sequence partitioning techniques, enabling gradient accumulation within individual sequences for parallel-prediction training. We implement P-EAGLE in vLLM and demonstrate speedups of 1.10-1.36x over autoregressive EAGLE-3 across GPT-OSS 120B, 20B, and Qwen3-Coder 30B.

P-EAGLE: Parallel-Drafting EAGLE with Scalable Training

TL;DR

P-EAGLE addresses the bottleneck of autoregressive decoding in reasoning-enabled LLMs by converting EAGLE-style drafting into parallel multi-token prediction using a learnable shared hidden state and mask-token strategy. It introduces a scalable training framework with amortized mask construction and sequence partitioning to manage the quadratic attention memory growth that arises with long contexts. The approach is implemented in vLLM and demonstrates consistent end-to-end speedups of 1.10×–1.36× across GPT-OSS 120B, 20B, and Qwen3-Coder 30B while maintaining comparable acceptance lengths to autoregressive baselines. These results establish parallel drafting as a viable production technique for accelerating long-context reasoning in large language models.

Abstract

Reasoning LLMs produce longer outputs, requiring speculative decoding drafters trained on extended sequences. Parallel drafting - predicting multiple tokens per forward pass - offers latency benefits over sequential generation, but training complexity scales quadratically with the product of sequence length and parallel positions, rendering long-context training impractical. We present P(arallel)-EAGLE, which transforms EAGLE from autoregressive to parallel multi-token prediction via a learnable shared hidden state. To scale training to long contexts, we develop a framework featuring attention mask pre-computation and sequence partitioning techniques, enabling gradient accumulation within individual sequences for parallel-prediction training. We implement P-EAGLE in vLLM and demonstrate speedups of 1.10-1.36x over autoregressive EAGLE-3 across GPT-OSS 120B, 20B, and Qwen3-Coder 30B.
Paper Structure (23 sections, 4 theorems, 5 figures, 11 tables, 1 algorithm)

This paper contains 23 sections, 4 theorems, 5 figures, 11 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Architecture
  3. Scalable Training Framework for Long Contexts
  4. Amortized Mask Construction
  5. Sequence Partitioning
  6. Training Recipe
  7. Hidden State Design
  8. Increasing Model Capacity
  9. Unfreezing the Embedding Layer
  10. Training vs. Inference Speculation Depth
  11. Extended Training Duration
  12. Longer Training Sequences
  13. Experiments
  14. Experimental Setup
  15. Acceptance Length Comparison
  16. ...and 8 more sections

Key Result

Theorem 2.1

The mapping $\delta \mapsto R_\delta$ from relative position to RoPE rotation matrix is injective, i.e., $R_{\delta_1} = R_{\delta_2} \implies \delta_1 = \delta_2$.

Figures (5)

  • Figure 1: Sequence length (prompt + generation) distribution on UltraChat dataset with GPT-OSS 120B. Reasoning level: Medium. Median: 3,891 tokens; P90: 10,800 tokens; P99: 20,000 tokens.
  • Figure 2: P-EAGLE architecture. The target model (top) processes prompt tokens and produces hidden states from layer indexes 2, $L/2$, and $L-1$ (concatenated to $3d$ dimensions), where $L$ is the number of decoder layers. The P-EAGLE drafter (bottom) takes these hidden states for the Next-Token Prediction (NTP) position (Pos 1), which operates like standard autoregressive prediction with actual context. Multi-Token Prediction (MTP) positions (Pos 2-4) use a learnable shared hidden state $h_{\text{shared}}$ since they lack preceding hidden states. Token embeddings are combined with projected hidden states and processed through $N$ transformer layers.
  • Figure 3: Position-invariance of causal attention across prediction depths. G0, G1, G2 in the figure denote prediction depths 0, 1, 2, where depth $d$ predicts the token $d+1$ positions ahead. The mask for a shorter sequence (right) is exactly the top-left submatrix of a longer sequence's mask (left), enabling constant-time retrieval.
  • Figure 4: Sequence partitioning for within-sequence gradient accumulation. Example: $n=16$ tokens, $K=4$ prediction depths (G0--G3 denote groups in PARD terminology, where group $g$ predicts the token $g+1$ positions ahead). Depth 0 contains all positions; depths 1--3 contain progressively fewer due to COD sampling. Partitioning by depth-0 indices causes dependency violations: position 8 at depth 2 depends on position 7 at depth 1, but they land in different segments. Our algorithm tracks assignments iteratively across depths to preserve dependencies.
  • Figure 5: Learnable $\alpha$ trajectory and comparison with baseline. Left:$\alpha$ decreases 71% from initialization, converging to $\sim$0.03. Center: MTP accuracy comparison---baseline (no context injection) achieves 57.9% versus the regularized variant's 54.6%. Right: NTP-MTP gap comparison---baseline achieves a lower gap (24.6%) than the regularized variant (27.4%). The model actively learns to minimize context injection because it hurts performance.

Theorems & Definitions (7)

  • Theorem 2.1: liu2025rethinking
  • Lemma 2.2
  • proof
  • Theorem 2.3: Attention Score-Level Injectivity
  • proof
  • Corollary 2.4: Absolute Position Recovery
  • proof