EAGLE-3: Soaring Beyond Speed [ The Art of Accelerated Language Models ]

Introduction

Code Snippet

Refer to Experimentation Script

In the rapidly evolving world of Large Language Models (LLMs), a growing trend has emerged: scaling up training data to boost model intelligence without inflating inference costs . This approach, akin to feeding a mind with more knowledge to sharpen its wit, has proven effective for many models. However, when applied to EAGLE a state-of-the-art speculative sampling method,this strategy hits a snag. Why? Because EAGLE’s reliance on feature prediction constraints creates a bottleneck, limiting the gains from additional data.

Enter EAGLE-3, a groundbreaking leap forward. Instead of clinging to feature prediction, EAGLE-3 boldly abandons it in favor of direct token prediction, liberating the draft model from unnecessary constraints. But that’s not all, EAGLE-3 also ushers in a new era of multi-layer feature fusion, replacing the narrow focus on top-layer features with a richer blend of low-, mid-, and high-level semantic information. This is achieved through an innovative technique called training-time test, which simulates multi-step generation during training, ensuring the draft model thrives on scaled-up data.

The result? A transformative enhancement in performance, enabling EAGLE-3 to fully harness the power of expanded training datasets while maintaining lightning-fast inference speeds. With EAGLE-3, the LLM community takes a bold step toward unlocking the true potential of scalable intelligence.

Glance at Approach

Speculative sampling slashes LLM latency by partially parallelizing token generation. In vanilla speculative sampling, a smaller draft model rapidly generates multiple tokens, which are then verified in parallel by the target model. This allows multiple tokens to be produced in a single forward pass, significantly cutting down inference time . However, the draft model operates independently, limiting its ability to fully leverage the target model’s rich context.

EAGLE takes speculative sampling further by reusing the top-layer features of the target model (the features before the LM head). It trains the draft model to autoregressively predict the next feature, then uses the target model’s LM head to derive the draft token. This tight coupling enables EAGLE to achieve superior acceleration compared to vanilla speculative sampling .

Key Differences: Vanilla vs. EAGLE

Aspect	Vanilla Speculative Sampling	EAGLE
Draft Model Input	Operates independently of the target model	Reuses top-layer features from the target model
Token Generation	Predicts tokens directly	Predicts features, then derives tokens via LM head
Acceleration	Moderate speedup	Significantly better acceleration

Comparison with Test-Time Scaling

Test-time scaling methods like DeepSeek-R1 enhance reasoning capabilities by engaging in deliberate multi-step reasoning before responding . While effective, these approaches increase inference time and computational costs. In contrast, EAGLE accelerates inference without compromising output quality, making it a more cost-efficient solution.

Core Concept of EAGLE

EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency) accelerates LLM inference by introducing a draft model that works alongside the base model. The draft model generates multiple tokens in parallel, which are then verified by the base model. This approach reduces the sequential nature of autoregressive token generation, significantly speeding up inference .

The key innovation in EAGLE is its ability to reuse top-layer features from the base model to train the draft model. Later versions like EAGLE-3 further enhance this by replacing feature prediction constraints with direct token prediction and leveraging multi-layer feature fusion.

Prefill vs Decoding Stage

LLM inference consists of two main stages:

Prefill Stage:
- Involves encoding the input prompt and initializing the key-value (KV) cache for subsequent token generation.
- This stage only happens once per input sequence.
- Example: For the input “How can”, the prefill stage computes the KV cache for “How” and “can” .
Decoding Stage:
- Generates tokens one at a time in an autoregressive manner.
- Speculative sampling accelerates this stage by generating multiple tokens in parallel using the draft model and verifying them with the base model.
- Example: After prefilling “How can”, the decoding stage generates “I”, “do”, “it”, etc., token by token .

Speculative Decoding vs Vanilla Inference

In vanilla autoregressive decoding:

Tokens are generated sequentially, one at a time.
Each token depends on the previous token, leading to high latency.

In speculative decoding:

The draft model predicts multiple tokens in parallel.
The base model verifies these tokens, accepting or rejecting them based on probability distributions.
If a token is rejected, the base model samples a new token and discards the remaining draft tokens.

Vanilla vs. EAGLE

Vanilla Speculative Sampling:
Input → [Draft Model] → Draft Tokens → [Target Model] → Verified Output

EAGLE:
Input → [Target Model Features] → [Draft Model] → Feature Prediction → [LM Head] → Draft Tokens → [Target Model] → Verified Output

By reusing top-layer features and integrating tightly with the target model, EAGLE bridges the gap between speed and intelligence, offering a novel approach to efficient LLM inference.

Prefill and Decoding Stages

Input Prompt: "How can"
                |
                v
        +------------------+
        |   Prefill Stage  |
        | - Encode "How"   |
        | - Encode "can"   |
        | - Initialize KV  |
        +------------------+
                |
                v
        +------------------+
        |   Decoding Stage |
        | - Generate "I"   |
        | - Generate "do"  |
        | - Generate "it"  |
        +------------------+

How EAGLE Trains the Draft Model

EAGLE trains the draft model to predict tokens by leveraging features from the base model. Here’s how it works:

Training Data Collection:
- Use sequences generated by the base model as training data.
- Extract top-layer features from the base model to guide the draft model.
Feature Fusion (EAGLE-3):
- Instead of relying solely on top-layer features, EAGLE-3 fuses low-, mid-, and high-level features from the base model.
- This provides richer semantic information, improving the draft model’s accuracy.
Training-Time Test:
- Simulate multi-step generation during training to ensure the draft model adapts to longer sequences.
- Remove the feature prediction constraint, allowing the draft model to focus directly on token prediction.

Training the Draft Model

Base Model Features
        |
        v
+-----------------------------+
| Feature Fusion (Low, Mid, High) |
+-----------------------------+
        |
        v
+-----------------------------+
| Train Draft Model           |
| - Predict Next Token        |
| - Simulate Multi-Step Gen   |
+-----------------------------+

How EAGLE Uses the Draft Model During Inference

During inference, EAGLE alternates between drafting and verification:

Drafting:
- The draft model generates multiple tokens in parallel.
- Example: For the prefix “How can”, the draft model might generate “I do it”.
Verification:
- The base model evaluates the draft tokens.
- Accept or reject each token based on probability distributions.
- Example: The base model accepts “I” and “do” but rejects “it”, replacing it with “help”.

ASCII Diagram: Inference Workflow

Input: "How can"
                |
                v
+---------------------------------+
| Base Model (Prefill)            |
| - Encode "How"                  |
| - Encode "can"                  |
| - Initialize KV Cache           |
+---------------------------------+
                |
                v
+---------------------------------+
| Draft Model (Speculative Decode)|
| - Predict "I do it"             |
+---------------------------------+
                |
                v
+---------------------------------+
| Base Model (Verification)       |
| - Accept "I"                    |
| - Accept "do"                   |
| - Reject "it", Sample "help"    |
+---------------------------------+
                |
                v
Output: "How can I do help"

Key Innovations in EAGLE-3

Direct Token Prediction:
- Removes the feature prediction constraint, allowing the draft model to focus on predicting tokens directly.
- Reduces error accumulation during long sequences.
Multi-Layer Feature Fusion:
- Replaces reliance on top-layer features with fused features from multiple layers.
- Captures richer semantic information, improving accuracy.
Training-Time Test:
- Simulates multi-step generation during training, enabling the draft model to generalize better.

Feature Fusion in EAGLE-3

Feature fusion is a key innovation in EAGLE-3 that replaces reliance on top-layer features from the base model with a richer combination of low-, mid-, and high-level features. Let’s break this down step by step:

1. What Are Features in LLMs?

In transformer-based models like LLaMA, features are representations of the input at different layers:

Low-level features: These capture basic information such as word embeddings or syntactic patterns (e.g., “The cat sat”).
Mid-level features: These encode more abstract patterns, such as relationships between words or phrases (e.g., “The cat sat on the mat” implies spatial context).
High-level features: These represent complex semantic information, such as intent or reasoning (e.g., “Why did the cat sit on the mat?”).

Traditionally, methods like EAGLE reused only the top-layer features (high-level features) from the base model to train the draft model. However, top-layer features are inherently limited because they focus solely on predicting the next token. This constraint makes it challenging for the draft model to predict multiple tokens accurately, especially for longer sequences.

2. How Does Feature Fusion Work?

EAGLE-3 introduces multi-layer feature fusion, where it combines features from low-, mid-, and high-level layers of the base model. Here’s how it works:

Extract Features: During the forward pass of the base model, EAGLE-3 extracts features from multiple layers:
- Low-level: Captures syntax and structure.
- Mid-level: Captures relationships between words/phrases.
- High-level: Captures semantics and reasoning.
Concatenate Features: These features are concatenated into a single vector:
```
Combined Feature = [Low-Level Feature, Mid-Level Feature, High-Level Feature]
```
For example, if each feature has a dimensionality of 1024, the combined feature would have a dimensionality of 3072.
Reduce Dimensions: A fully connected (FC) layer reduces the combined feature vector back to the original dimensionality (e.g., 1024). This ensures compatibility with the draft model.
Train the Draft Model: The fused feature vector serves as input to the draft model, enabling it to generate tokens based on richer semantic information.

Example of Feature Fusion

Imagine the input sentence:
“The quick brown fox jumps over the lazy dog.”

Low-level features: Represent individual words like “fox,” “jumps,” and “dog.”
Mid-level features: Capture relationships such as “fox jumps over dog.”
High-level features: Encode the overall meaning, such as “a fox performing an action.”

By fusing these features, EAGLE-3 provides the draft model with a holistic understanding of the input, improving its ability to predict future tokens accurately.

Why Is Feature Fusion Important?

Feature fusion allows the draft model to:

Predict multiple tokens without relying solely on the next-token prediction capability of top-layer features.
Mitigate error accumulation, as the draft model has access to richer contextual information .

Training-Time Test in EAGLE-3

The training-time test is another critical innovation in EAGLE-3. It simulates multi-step generation during training, enabling the draft model to adapt to longer sequences and removing the feature prediction constraint.

1. What Is the Feature Prediction Constraint?

In earlier versions of EAGLE, the draft model was trained to predict the next feature (a representation of the next token) rather than the token itself. This created two challenges:

Error Accumulation: If the predicted feature deviated from the ground truth, subsequent predictions became less accurate.
Limited Expressiveness: The draft model was constrained by the need to align its output with the base model’s top-layer features.

EAGLE-3 removes this constraint by directly predicting tokens instead of features.

2. How Does Training-Time Test Work?

During training, EAGLE-3 simulates the inference process to ensure the draft model learns to handle multi-step generation. Here’s how it works:

Step 1: Generate a Token:
- The draft model generates the first token using fused features from the base model.
- Example: Given the prefix “The quick brown fox”, the draft model predicts “jumps”.
Step 2: Simulate Multi-Step Generation:
- Instead of stopping after one token, EAGLE-3 feeds the generated token back into the draft model as input.
- Example: After generating “jumps”, the draft model uses it to predict the next token, e.g., “over”.
Repeat the Process:
- This continues for multiple steps, simulating the entire generation process during training.
- Example: The draft model generates “The quick brown fox jumps over the lazy dog.” step by step.
Loss Calculation:
- The loss function evaluates the accuracy of all generated tokens, ensuring the draft model adapts to longer sequences.

Example of Training-Time Test

Consider the input prefix:
“How can I improve my writing skills?”

Step 1: The draft model predicts the first token, e.g., “Read”.
Step 2: The draft model uses “Read” as input to predict the next token, e.g., “more”.
Step 3: The draft model uses “Read more” to predict the next token, e.g., “books”.
Final Output: The draft model generates “Read more books to enhance your vocabulary and style.”

By simulating this process during training, EAGLE-3 ensures the draft model can handle real-world scenarios where multiple tokens are generated sequentially.

Why Is Training-Time Test Important?

Training-time test enables EAGLE-3 to:

Remove the feature prediction constraint, allowing the draft model to focus directly on token prediction .
Adapt to longer sequences, reducing error accumulation and improving acceptance rates.

Feature Fusion and Training-Time Test

Input: "The quick brown fox"
                |
                v
+-----------------------------------+
| Extract Features from Base Model |
| - Low-Level: Syntax               |
| - Mid-Level: Relationships        |
| - High-Level: Semantics           |
+-----------------------------------+
                |
                v
+-----------------------------------+
| Concatenate and Fuse Features     |
| Combined Feature = [Low, Mid, High]|
+-----------------------------------+
                |
                v
+-----------------------------------+
| Reduce Dimensionality via FC Layer|
+-----------------------------------+
                |
                v
+-----------------------------------+
| Train Draft Model                 |
| - Direct Token Prediction         |
| - Simulate Multi-Step Generation  |
+-----------------------------------+
                |
                v
Output: "jumps over the lazy dog"

Draft Model Training (Eagle’s Style)

Training the draft model in EAGLE-3 is a masterclass in innovation, blending feature fusion, self-attention adjustments, and LM head integration to simulate multi-step generation. Let’s break it down step by step, with examples and visuals to keep things crisp and engaging.

1. Feature Fusion: The Heart of EAGLE-3

Instead of relying solely on top-layer features like its predecessors, EAGLE-3 fuses low-, mid-, and high-level features from the target model. This fusion provides the draft model with richer semantic information, enabling it to predict tokens more accurately.

How It Works

During the target model’s forward pass, features from different layers are extracted:
- Low-level: Captures syntax and basic structure.
- Mid-level: Encodes relationships between words/phrases.
- High-level: Represents complex semantics and reasoning.
These features are concatenated into a single vector and passed through a fully connected (FC) layer to reduce dimensionality while preserving context.

Example

For the input “The quick brown fox”, the fused feature vector might encode:

Low-level: “fox” as a noun.
Mid-level: Relationship between “quick” and “fox”.
High-level: Semantic meaning of “a fast animal.”

This rich context empowers the draft model to generate multiple tokens confidently during speculative sampling.

2. Self-Attention Mechanism: Simulating Multi-Step Generation

EAGLE-3 modifies the self-attention mechanism to simulate multi-step token generation during training—a process called training-time test. This ensures the draft model learns to handle sequential dependencies effectively.

Key Adjustment

Traditional self-attention operates sequentially, attending only to previous tokens. In EAGLE-3:
- The attention mask is adjusted dynamically to simulate future token predictions.
- For example, when generating “jumps over the lazy dog”, the model predicts “jumps” based on “The quick brown fox” and then simulates the next steps (“over”, “the”, etc.) during training.

Visualizing the Process

Here’s how the attention causal masks evolve:

Step 1: [How] → [can]
Step 2: [How, can] → [I]
Step 3: [How, can, I] → [do]

Each step feeds back into the draft model, creating a recursive loop that mimics real-world multi-step generation.

3. LM Head: Converting Features to Tokens

Finally, the LM head converts the fused features into token predictions. Unlike vanilla speculative sampling, which relies on feature prediction constraints, EAGLE-3 directly outputs tokens—removing unnecessary bottlenecks.

Example

Using the fused feature vector for “fox”, the LM head predicts the next token, say “jumps”. This direct approach eliminates error accumulation caused by intermediate feature mismatches.

Transformer Components Involved

To visualize the process, here’s an ASCII diagram highlighting the key transformer components:

Input: "The quick brown fox"
        |
        v
+-----------------------------+
| Extract Features            |
| - Low-Level: Syntax         |
| - Mid-Level: Relationships  |
| - High-Level: Semantics     |
+-----------------------------+
        |
        v
+-----------------------------+
| Concatenate & FC Layer      |
| Combined Feature Vector     |
+-----------------------------+
        |
        v
+-----------------------------+
| Draft Model                 |
| - Adjusted Self-Attention   |
| - Simulates Multi-Step Gen  |
+-----------------------------+
        |
        v
+-----------------------------+
| LM Head                     |
| Outputs Predicted Token     |
+-----------------------------+
Output: "jumps"

Why This Matters

By combining feature fusion, adjusted self-attention, and direct token prediction, EAGLE-3 trains the draft model to excel at speculative sampling. This approach not only enhances accuracy but also scales gracefully with larger datasets .

In essence, EAGLE-3 transforms the draft model into a nimble yet powerful assistant, capable of accelerating LLM inference without compromising quality.

Why Use a Small Draft Model?

A smaller draft model is crucial for speculative sampling because:

Speed: Smaller models generate tokens faster, reducing latency.
Memory Efficiency: They consume less GPU memory, enabling larger batch sizes and better utilization of consumer-grade hardware like NVIDIA A40 GPUs.
Scalability: By decoupling the draft model from the target model, speculative sampling can scale to larger models without retraining.

For example, in our experiments, we used lmsys/sglang-EAGLE-llama2-chat-7B as the draft model and meta-llama/Llama-2-13b-chat-hf as the base model. This combination achieved significant speedups while maintaining accuracy.

Preparing the Draft Model

Training a draft model for EAGLE involves several steps:

Data Collection: Use sequences generated by the target model to create training data.
Feature Fusion: Train the draft model to predict tokens directly using fused features from multiple layers of the target model.
Training-Time Test: Simulate multi-step generation during training to ensure the draft model generalizes well to longer sequences.

While fine-tuning a draft model from scratch can be challenging due to incomplete documentation, pre-trained models like lmsys/sglang-EAGLE-llama2-chat-7B provide a reliable starting point.

Our Experiments

Refer to Experimentation Script - Read more about challenges collecting necessary metrics from Sglang’s metrics publihsed to Prometheus related to Speculative Decoding.

We conducted four experiments to evaluate the performance of EAGLE-3 under different configurations:

Vanilla Decoding
Baseline configuration with no speculative sampling. Metrics collected included throughput, latency, and acceptance rate.
EAGLE Decoding
Enabled speculative sampling with the draft model. Parameters included:
- --speculative-num-steps 5
- --speculative-eagle-topk 8
- --speculative-num-draft-tokens 64
EAGLE + Torch Compile
Added Torch Compile optimizations to further accelerate inference. Key parameters:
- --enable-torch-compile
- --cuda-graph-max-bs 2
EAGLE-3 Simulation
Simulated EAGLE-3 behavior by leveraging the same draft model and optimizing for multi-layer feature fusion.

Despite challenges in collecting meaningful metrics (e.g., metrics not appearing in logs), all experiments ran successfully. We observed qualitative improvements in speed and responsiveness, particularly with EAGLE-3.

Conclusion

Speculative Sampling

EAGLE redefines the boundaries of Large Language Model (LLM) inference acceleration by marrying speculative sampling with innovative architectural advancements. At its core, EAGLE abandons restrictive feature prediction constraints, embracing direct token prediction and multi-layer feature fusion to unlock unprecedented scalability and performance . By simulating multi-step generation during training through the training-time test, EAGLE ensures that its draft model fully benefits from expanded training data, achieving a remarkable 6.5x speedup over vanilla autoregressive decoding while maintaining lossless output quality .

This leap in efficiency not only surpasses predecessors like EAGLE-2 but also sets a new benchmark for speculative sampling methods. With its ability to dynamically adapt to complex contexts and scale gracefully, EAGLE exemplifies how rethinking traditional paradigms can propel LLMs into a faster, smarter future .

For a deeper dive into the technical intricacies, explore the EAGLE paper or visit the GitHub repository.

Eagle Speculative Sampling Related Posts