Speculative Decoding and EAGLE

Decode throughput for a 70B-class model on H100 is typically 40-80 tokens/second. Each token needs a full forward pass that reads all the model weights from HBM. You can't make the model smaller without changing its output, and you can't push batch size past what memory allows. The only way out is to get more than one token per forward pass.

Autoregressive generation looks inherently serial: x_{t+1} = sample(p(· | x_{1:t})). But there is room for concurrency. Suppose a cheap predictor says "the next 4 tokens are probably [a, b, c, d]". You can verify all 5 positions in a single forward pass of the big model and accept the longest matching prefix.

Leviathan, Kalai, Matias (2023, "Fast Inference from Transformers via Speculative Decoding") made this exact with an accept/reject rule that preserves the target model's sampling distribution. Same output distribution, 2-4× faster.

The Two-Model Setup

• Target model M_p: the big, slow, high-quality model you actually want samples from. Distribution: p(x).
• Draft model M_q: a small, fast, lower-quality model with distribution q(x). It is roughly 10-70× smaller than the target (about 1B for a 70B target, 0.6B for a 32B one, 500M for a 7B one).

Per step:

1. Draft model proposes K tokens autoregressively: x_1, x_2, ..., x_K ~ q.
2. The target model runs one forward pass over all K+1 positions in parallel and produces p(x_k) for each proposed token.
3. Accept/reject each token left-to-right via the modified rejection-sampling rule below. Accept the longest matching prefix.
4. If any token is rejected, sample the replacement from the corrected distribution and stop. Otherwise sample one bonus token from p(· | x_1...x_K).

If the draft matches the target perfectly, you get K+1 tokens per target forward. If the draft is wrong at position 1, you get 1 token.

The Exactness Rule

Speculative decoding produces samples with the same distribution as sampling from p directly, and the rejection rule below is what guarantees it.

For each drafted token x_t:
    r ~ Uniform(0, 1)
    if r < p(x_t) / q(x_t):
        accept x_t
    else:
        sample replacement from residual: (p - q)+ / ||(p - q)+||_1
        stop

where (p - q)+ denotes the positive part of the pointwise difference. When the draft and target agree (p ≈ q) acceptance is close to 1. When they disagree, the residual distribution is built so the overall sample still follows p.

Greedy case. For temperature=0 sampling, check argmax(p) == x_t. If it matches, accept. If not, output argmax(p) and stop.

Expected Speedup

If the draft model's token-level acceptance rate is α, the expected number of tokens per target forward pass is:

E[tokens] = (1 - α^{K+1}) / (1 - α)        # K = draft length, α in [0, 1]

At α = 0.8, K = 4: (1 - 0.8^5)/(1 - 0.8) = 3.36 tokens per forward. One round costs roughly cost_q * K + cost_p: K draft steps plus one target verify. If cost_p >> cost_q * K, the speedup is about 3.36× / 1 = 3.36× on throughput.

The parameter that matters is α, and it comes from how well the draft is aligned with the target. A good draft is what makes the whole thing pay off.

Training the Draft: Distillation

A randomly initialized small model makes a poor draft. The standard recipe is to distill from the target:

1. Pick a small architecture (~1B for a 70B target, ~500M for a 7B target).
2. Run the target model on a large text corpus; store its next-token distributions.
3. Train the draft with KL divergence against the target's distribution (not against ground-truth tokens).

The result is α of 0.6-0.8 on coding and 0.7-0.85 on natural-language chat, which gives 2-3× speedups in production.

EAGLE: Tree Drafting + Feature Reuse

Li, Wei, Zhang, Zhang (2024, "EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty") spotted two inefficiencies in standard speculative decoding:

1. The draft does K serial steps, each through its full stack. It could instead reuse the target's features (hidden states) from the most recent verify, since the target already computed rich representations that the draft re-derives from scratch.
2. The draft outputs a linear chain. If it output a tree of candidates instead (multiple guesses per node), the target's single forward pass could verify several paths in parallel via a tree attention mask and pick the longest accepted branch.

EAGLE-1 changes:

• The draft takes the target's final hidden state at position t as input, not raw tokens.
• The draft is one transformer decoder layer rather than a separate small model.
• The output is a tree of K = 4-8 candidates per depth, with depth 4-6.

EAGLE-2 (2024) adds dynamic tree topology. The tree grows wider where the draft is uncertain and stays narrow where it is confident, which raises α_effective without raising verify cost.

EAGLE-3 (Li et al. 2025, "EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test") removes the fixed top-layer feature dependency and trains the draft with a new "test-time simulation" loss. The draft learns on outputs that match the target's test-time distribution instead of the teacher-forced training distribution. Acceptance rate rises from 0.75 (EAGLE-2) to 0.82 (EAGLE-3), and mean tokens/verify from 3.0 to 4.5.

Tree Attention Verification

When the draft outputs a tree, the target model verifies it in a single forward pass using a tree attention mask, a causal mask that encodes the tree topology rather than a pure line. Each token attends only to its ancestors in the tree. The verify pass is still one forward and one matmul, and the topological mask costs only a few extra KV entries.

        root
       /    \
      a      b
     / \    / \
    c  d   e   f

If a, b are competing first-token candidates and c, d, e, f are second-token candidates, all six positions get verified in one forward pass. The output is the longest prefix along any accepted path.

When It Wins, When It Doesn't

Wins:

• Chat and completion with predictable text (code, common English, structured output), where α is high.
• Setups with idle GPU compute during decode (the memory-bound phase). Tree drafting puts the spare FLOPs to work.

Loses / no win:

• Highly stochastic outputs (creative writing at high temperature), where α drops toward 1/|vocab|.
• Batch serving at very high concurrency. Batching already fills the FLOPs, so there is little room for tree verification.
• Very small target models, where the draft isn't much smaller.

Production deployments typically report 2-3× wall-clock speedup on chat, 3-5× on code generation, and close to none on creative writing.

main.py runs on synthetic toy distributions so the math stays visible. It needs only numpy. There is no neural network in it and no GPT-2 distillation; the target is a fixed distribution over a small vocabulary, and the draft is that target blended with a uniform distribution.

• speculative_step(target, draft, K, rng) implements the exact rejection rule. The harness compares the empirical total-variation distance from the target for plain sampling and for speculative sampling, using the first token of each round.
• measure_alpha measures the acceptance rate of the draft. The alpha flag is only a hint for how make_draft blends, so the measured value comes out different.
• measure_throughput counts tokens per round at draft length K and prints it next to expected_tokens(alpha, K), the closed form from The Concept.
• build_tree and tree_attention_mask construct a depth-3 tree with branching (3, 2, 2), and validate_tree_mask checks that every node attends to its ancestors and nothing else. The tree here is topology only. There is no draft model filling it in.

The snippet below is illustrative. It compresses the same idea and adds a temperature argument that main.py does not have.

def speculative_step(p_target, q_draft, K, temperature=1.0):
    """One round of speculative decoding. Returns list of accepted tokens."""
    # 1. Draft K tokens
    draft_tokens = []
    q_probs = []
    state = draft_state_init()
    for _ in range(K):
        probs = softmax(q_draft(state) / temperature)
        t = np.random.choice(len(probs), p=probs)
        draft_tokens.append(t)
        q_probs.append(probs[t])
        state = draft_step(state, t)
 
    # 2. Target computes p at every drafted position + 1 extra
    p_probs_all = target_forward_batched(p_target, draft_tokens, temperature)
 
    # 3. Accept/reject left-to-right
    accepted = []
    for k, tok in enumerate(draft_tokens):
        r = np.random.uniform()
        if r < p_probs_all[k][tok] / q_probs[k]:
            accepted.append(tok)
        else:
            residual = np.maximum(p_probs_all[k] - q_probs[k], 0)
            residual /= residual.sum()
            accepted.append(np.random.choice(len(residual), p=residual))
            return accepted
    # 4. All K accepted → sample bonus token from target
    accepted.append(np.random.choice(len(p_probs_all[-1]), p=p_probs_all[-1]))
    return accepted

• vLLM and SGLang ship first-class speculative decoding. Flags: --speculative_model, --num_speculative_tokens. EAGLE-2/3 support comes via the --spec_decoding_algorithm eagle flag.
• NVIDIA TensorRT-LLM supports Medusa and EAGLE trees natively.
• Reference draft models: Qwen/Qwen3-0.6B-spec (drafts for Qwen3-32B) and meta-llama/Llama-3.2-1B-Instruct-spec (drafts for 70B).
• Medusa heads (Cai et al. 2024, "Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads"): instead of a draft model, add K parallel prediction heads to the target itself. Simpler to deploy, with slightly lower acceptance than EAGLE.

This lesson produces outputs/skill-speculative-tuning.md, a skill that profiles a target model's workload and chooses the draft model, K (draft length), tree width, and temperature, and decides when to fall back to plain decode.

1. Implement the exact rejection rule and check it empirically. Run 10K samples via speculative_step and via plain target sampling; compute TV distance between the two output distributions. The distance should be below 0.01.

2. Compute the speedup formula. Given fixed α and K, plot expected tokens per target-forward. Find the optimal K for α ∈ {0.5, 0.7, 0.9}.

3. Train a tiny draft. Take a 124M GPT-2 target and distill a 30M GPT-2 draft on 100M tokens with KL loss. Measure α on held-out text. Expect 0.6-0.7.

4. Implement EAGLE-style tree drafting. Instead of a chain, have the draft output top-3 branches at each depth. Build the tree attention mask. Verify the target accepts the longest correct branch.

5. Measure failure modes. Run speculative decode at temperature=1.5 (high stochasticity). Show α collapses and the algorithm is slower than plain decode due to draft overhead.

• Leviathan, Kalai, Matias, 2023, "Fast Inference from Transformers via Speculative Decoding": the exact rejection rule and the theoretical speedup analysis
• Chen, Borgeaud, Irving et al., 2023, "Accelerating Large Language Model Decoding with Speculative Sampling": concurrent speculative-sampling paper at DeepMind
• Cai, Li, Geng, Wang, Wang, Zhu, Dao, 2024, "Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads": parallel-heads alternative to a draft model
• Li, Wei, Zhang, Zhang, 2024, "EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty": feature reuse and tree drafting
• Li et al., 2024, "EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees": dynamic tree topology
• Li et al., 2025, "EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test": train-time test-time matching
• Fu, Haotian, Peng et al., 2024, "Break the Sequential Dependency of LLM Inference Using Lookahead Decoding": Jacobi/lookahead decoding, a speculator-free alternative