P-EAGLE: Parallel Speculative Decoding in vLLM

TL;DR

P-EAGLE is a parallel speculative decoding method integrated into vLLM (starting from v0.16.0) that eliminates the sequential bottleneck of autoregressive drafting. By generating all K draft tokens in a single forward pass, it achieves up to 1.69x speedup over vanilla EAGLE-3 on real-world workloads using NVIDIA B200 GPUs.

Solving the Autoregressive Drafting Bottleneck

Standard speculative decoding methods like EAGLE draft tokens autoregressively, meaning that producing K draft tokens requires K sequential forward passes through the draft model. This creates a linear scaling of latency relative to speculation depth, which limits how aggressively a system can speculate without the drafting overhead eating into the overall performance gains.

P-EAGLE removes this ceiling by transforming the drafting process from autoregressive to parallel. Instead of sequential steps, P-EAGLE generates the entire set of draft tokens in one single forward pass, decoupling the number of draft tokens from the number of forward passes required.

P-EAGLE Architecture and Mechanism

P-EAGLE operates in two primary steps to generate draft tokens:

Step 1: Prefilling

The target model processes the prompt and generates a new token. During this process, P-EAGLE captures the internal hidden states: h_prompt for each prompt position and h_context for the newly generated token. These hidden states guide the drafter's predictions.

Step 2: Parallel Drafting

The drafter constructs inputs for each position in parallel using a combination of token embeddings and hidden states:

  • Prompt Positions: Each prompt token embedding emb(p) is paired with its corresponding h_prompt from the target model, shifted by one position to enable prediction of the token at position $i$.
  • Next-Token-Prediction (NTP): The first position pairs the newly generated token embedding emb(new) with h_context.
  • Multi-Token-Prediction (MTP): For positions 2 through K, where token embeddings and hidden states do not yet exist, P-EAGLE uses two learnable parameters: a shared mask token embedding emb(mask) and a shared hidden state h_shared. These serve as neutral placeholders.

All positions then pass through N transformer layers and a language model head to predict draft tokens $t_1$ through $t_K$ simultaneously.

Training on Long Sequences

Parallel drafting increases memory requirements during training because training K parallel groups on a sequence of length N creates $N \times K$ total positions. For example, with $N=8,192$ and $K=8$, a single training example contains 65,536 positions, leading to massive attention matrices (over 4 billion elements).

To address this, P-EAGLE introduces a sequence partition algorithm for intra-sequence splitting. This algorithm divides the $N \times K$ position sequence into contiguous chunks while maintaining correct attention dependencies across boundaries and accumulating gradients across chunks of the same sequence.

vLLM Implementation Details

Integrating parallel drafting into vLLM required overcoming several technical challenges related to batch metadata and memory management:

Fused Triton Kernel

Parallel drafting breaks the consistency between drafting and verification batch shapes because of the added MASK placeholders. To avoid the overhead of rebuilding batch metadata through multiple GPU operations, vLLM implements a fused Triton kernel. This kernel handles copying previous token IDs and positions, inserting the bonus token, filling parallel-drafting slots with MASK token IDs, and generating necessary metadata (rejected-token masks, masked-token masks, and hidden-state mappings) in a single pass.

Hidden State Management

Because hidden states are significantly larger than token IDs, vLLM uses a dedicated copy kernel to broadcast the learned hidden state placeholder (parallel_drafting_hidden_state_tensor) into the mask token slots, while target hidden states are mapped to their new positions.

KV Cache and CUDA Graphs

  • KV Cache: Rejected tokens are mapped to PADDING_SLOT_ID (-1) to prevent spurious cache writes.
  • CUDA Graphs: The capture range is extended by $K \times \text{max_num_seqs}$ to accommodate the larger draft batch size.

Performance Benchmarks

Evaluations on GPT-OSS-20B using an NVIDIA B200 GPU show that P-EAGLE significantly outperforms vanilla EAGLE-3 across multiple benchmarks (MT-Bench, HumanEval, and SPEED-Bench).

Throughput Gains

At low concurrency (c=1), P-EAGLE delivers 55–69% higher throughput. At high concurrency (c=64), gains of 5–25% are sustained. Specifically, the speedup ratios over EAGLE-3 are:

  • SPEED-Bench: Up to 1.69x (at c=1)
  • HumanEval: Up to 1.55x (at c=1)
  • MT-Bench: Up to 1.55x (at c=1)

Acceptance Length (AL)

P-EAGLE achieves a higher average number of accepted draft tokens per round (AL) than EAGLE-3. At a speculation depth of $K=7$, P-EAGLE's AL is 30% higher on HumanEval (3.94 vs 3.03) and 31% higher on SPEED-Bench (3.38 vs 2.59).

Speculation Depth Efficiency

Unlike autoregressive drafters, which peak in throughput at $K=3$, P-EAGLE consistently achieves peak throughput at $K=7$ across all concurrency levels. This confirms that P-EAGLE can benefit from deeper speculation without the linear latency penalty associated with sequential drafting.

Deployment

Parallel drafting can be enabled in vLLM by setting "parallel_drafting": true in the SpeculativeConfig class. Pre-trained P-EAGLE heads are available on HuggingFace for the following models:

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

Sources