Reproduce “P-EAGLE: Parallel-Drafting EAGLE with Scalable Training” on OpenResearch

About this paper

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

Researchers from the University of California, Santa Cruz, and AWS Amazon developed P-EAGLE, a framework that transforms EAGLE-style speculative decoding into a parallel multi-token prediction architecture. This approach enables scalable training on long contexts (up to 20,000 tokens) and achieves end-to-end LLM inference speedups of 1.10x to 1.36x compared to autoregressive EAGLE-3, while maintaining comparable draft quality.

Key insights

• A simple learnable shared hidden state and a trainable mask token embedding are effective for parallel multi-token prediction in EAGLE, outperforming more complex position-aware alternatives.
• Amortized attention mask construction, pre-calculating and then slicing masks, drastically reduces data loading and epoch times (48x and 5x respectively) during training on long sequences.
• Dependency-aware sequence partitioning for within-sequence gradient accumulation allows training on extremely long contexts (up to 20,000 tokens) by managing memory consumption while preserving causal attention dependencies.
• Deeper drafter architectures (e.g., 4 layers) and training with slightly higher speculation depths than inference generally improve parallel draft quality and acceptance length.

Summary by alphaXiv.