5.3 Medusa, EAGLE, and MTP

A standalone Draft is like hiring another assistant: pay is low, but they must re-read the whole text and keep a separate notebook before writing. Another approach lets the assistant reuse what the editor already understood: Target’s hidden states hold its reading of the context, and attaching a few small heads or a light network on those representations is often cheaper than running a full model from tokens again.

Methods in this family are called Self-Draft. This section covers them in order of appearance: Medusa hangs several prediction heads on Target; the EAGLE series trains a small network that reads Target features; MTP is a multi-token prediction module baked into pretraining. All need specially trained weights, and those weights bind to a specific Target version.

⏰ Freshness: paper claims follow each arXiv version; vLLM support is based on v0.32.0 (released 2026-10-18); SGLang parameters follow its official docs.

📑 Contents

1. Self-Draft: what you reuse, what you add
2. Medusa: one hidden state, several future heads
3. Draft Tree and Tree Attention
4. EAGLE: a light draft from features
5. EAGLE-2: spend the tree budget on more confident paths
6. EAGLE-3: change the training objective and input signals
7. MTP: a draft module that ships with pretraining
8. Hands-on: why sibling nodes must not see each other
9. Verification and deploy: tree compute does not auto-guarantee lossless
Summary
Self-check list
References

1. Self-Draft: what you reuse, what you add

📊 Approach	What of Target is reused	What is added	How you get it
Standalone Draft	Only shared token history	A full small model	Pick a compatible same-family model
Medusa	Last-layer hidden state	Several prediction heads	Train those heads on Target
EAGLE / EAGLE-2	High-level features + tokens	A one-ish-layer draft network	Train against Target
EAGLE-3	Low/mid/high multi-layer features + tokens	Draft network	Improved training pipeline
MTP	Shared Embedding and output head	MTP module trained with pretraining	Ships with model weights

All of these improve the “produce a draft” step; Target still verifies. Weights, activations, and cache of the helper module also count against GPU-memory budget.

🔑 Core idea: Reuse Target’s representations to skip re-understanding; keep Target’s verification so the final output is still its call. Whether the distribution is strictly preserved depends on which acceptance rule you use (see §9).

2. Medusa: one hidden state, several future heads

2.1 Several heads guess together

A normal LM Head maps the current hidden state $h_t$ to the next-token distribution. Medusa hangs $K$ more small heads on the same $h_t$ ; head $k$ directly predicts the farther position $t+k+1$ :

q^{(k)}=\operatorname{softmax}\big(g_k(h_t)\big),\qquad k=1,\ldots,K

Each $g_k$ is a single-layer FFN with a residual. All heads read the same $h_t$ , so they can finish in one shot — no serial multi-step run like a standalone Draft.

The cost is clear too: when the 3rd head predicts, it does not know what positions 1 and 2 finally chose. Farther heads are more like “guessing with eyes closed,” and accuracy decays faster.

2.2 Several candidates form a tree

If each head only keeps Top-1, one early miss kills the whole chain. Medusa keeps several candidates per head, organizes possible combinations into a tree, and lets Target verify many branches in one go.

You cannot expand every Top-k combination of every head: $b$ branches per level and depth $d$ means $b^d$ -scale nodes. Medusa’s approach is to measure path hit rates on a calibration set and keep only the most promising nodes, forming a sparse tree.

2.3 Medusa-1 and Medusa-2

The paper gives two training modes:

Medusa-1: freeze Target’s backbone, train only the prediction heads. Target is unchanged; with exact verification you can stay lossless.
Medusa-2: train backbone and heads together so representations better fit predicting far tokens.

Medusa-2 has effectively produced a new Target. Even with exact verification, what it preserves is consistency with the new model; quality evals must compare against the original model before training. Giving the editor an assistant only, versus retraining the editor’s writing habits too, are different degrees of change.

3. Draft Tree and Tree Attention

3.1 Every path on the tree is a different history

After the same prefix you might continue with return or with raise; their follow-ups belong to two mutually exclusive histories. Tree Attention requires: each node may only see the shared prefix and its own ancestors.

When you flatten tree nodes into an array for Attention, you cannot use a normal lower-triangular mask — otherwise a later B branch would see the A branch, treating two exclusive continuations as one continuous text.

The figure below draws the tree mask on the same tree used in §8 (parents = [-1, 0, 0, 1, 2]). Yellow is the return branch, blue the raise branch; solid blue cells in the matrix are visible; red hatches are positions a lower-triangular mask would wrongly see; the rightmost column is Position ID numbered by depth:

5-node Draft Tree and tree mask: each node sees only itself and ancestors; node 4 cannot see other-branch nodes 1 and 3; nodes 1 and 2 both have Position ID 1

3.2 Position IDs are numbered by depth

Sibling nodes at the same level sit at the same sequence position; Position IDs should follow tree depth, not increment by flattened array index. The mask decides “who you can see”; Position ID decides “which position you are.” Both must be right.

All nodes share the prefix KV; after verification only one accepted path is committed, and temporary KV of other branches is discarded.

📌 Key point: Tree Attention is a batching trick for computing several candidate histories at once. It keeps each path’s compute correct; which path is chosen is the verification rule’s job.

4. EAGLE: a light draft from features

4.1 Why draft on features

A standalone small model re-runs the network from tokens, redoing understanding Target already did. EAGLE’s observation: the top-level feature just before Target’s LM Head (the paper’s second-to-top-layer feature) is more regular than tokens. A small network of “one fully connected layer + one Decoder layer” can autoregressively predict the next feature from the current one, then map it through Target’s LM Head into token candidates.

4.2 Features alone are not enough — tokens matter too

A feature sequence does not fully determine the future. Whether sampling picked return or raise, the next feature goes in completely different directions. So EAGLE’s draft network takes two inputs at once: the current feature, and the next token already sampled (the token sequence one step ahead). Token information removes the “which branch” uncertainty.

Draft-network outputs still go to Target for verification. Because it reads Target features, structure, hidden size, which feature layer, and the Target used in training must all match.

⚠️ Note: EAGLE weights bind to a specific Target version. Two 8B models with the same size and similar names cannot share one EAGLE weight file.

5. EAGLE-2: spend the tree budget on more confident paths

5.1 Waste in a fixed tree

EAGLE and Medusa both use static trees: depth and branching stay fixed every round. In practice some contexts have an almost determined continuation (good for going deep); others start with several plausible writings (good for going wide).

EAGLE-2 found draft-model confidence correlates strongly with actual acceptance, so each round it expands and prunes the candidate tree by confidence, putting a limited node budget where acceptance is most likely.

5.2 Look at the whole path, not just the current node

A node may have confidence 0.95 itself, but if its parent only has a 0.1 chance of being accepted, its chance of being committed is small too. Tree paths must pass continuously from the root, so EAGLE-2 estimates node value by the product of confidences along the path:

v(u)\approx\prod_{j\in\operatorname{path}(u)}c_j

where $c_j$ is the draft model’s confidence on path node $j$ . Expansion keeps growing the highest- $v$ nodes; the rearrange stage then picks the highest- $v$ batch among all nodes for verification.

5.3 Confidence only allocates budget

“Confidence correlates with acceptance” is an empirical relation measured on the paper’s models and tasks; change model, temperature, or task and re-check. Also, $v(u)$ only decides which nodes are sent to verification; verification itself is still the standard rule from §5.1 (core speculative-decoding principles). Confidence does not decide accept/reject.

6. EAGLE-3: change the training objective and input signals

EAGLE-3’s authors found that adding training data to EAGLE helped little, because the constraint “must predict Target’s next feature” limited the draft network. EAGLE-3 makes three changes:

Drop the feature-prediction loss; train only on token prediction so the draft network can learn representations useful for prediction.
Fuse multi-layer features: concatenate Target’s low, mid, and high layer features as input, instead of depending only on the last layer.
Training-time Test (TTT): at inference, from the second step on the draft network reads its own previous output, not Target’s true features. Training simulates that process so the network gets used to continuing from its own intermediates.

Here “Test” means simulating multi-step drafting of the inference stage during training — unrelated to eval sets or post-deploy retraining.

Analogy: EAGLE asks the assistant to first copy the sketch in the editor’s head, then write; EAGLE-3 only asks them to write what the editor would write, and during training already practice “continue from what you just wrote.”

7. MTP: a draft module that ships with pretraining

7.1 From training objective to inference speedup

The methods above all say “after the model is trained, attach a draft module.” Multi-Token Prediction (MTP) flips the order: during pretraining the model already predicts several future tokens; that auxiliary objective can learn better representations; at inference the trained MTP module can serve as the draft.

DeepSeek-V3 is the flagship. Its MTP module differs from Medusa’s “parallel small heads” in one key way: predict in order, keep the causal chain. Module $k$ reads the previous depth’s representation $h^{k-1}$ and the Embedding of token $i+k$ , RMSNorms each, concatenates, projects back with $M_k$ , runs a Transformer Block, then predicts token $i+k+1$ with the output head shared with the main model:

h'^{k}_i=M_k\big[\operatorname{RMSNorm}(h^{k-1}_i);\ \operatorname{RMSNorm}(\operatorname{Emb}(t_{i+k}))\big]

Embedding and output head are shared with the main model, so the MTP module is light relative to the main model. The structure almost matches EAGLE’s “feature + token” idea; the difference is it is pretrained with the main model, not trained afterward alone.

7.2 Practical results

DeepSeek-V3 trains with only one extra token predicted (depth $D=1$ ). The paper reports that using the MTP module for speculation at inference, the second token’s acceptance rate is 85%–90%, and decode speed (TPS) reaches 1.8× the baseline. Those numbers are for that model and deploy setup; other models or loads need re-measurement.

📌 Key point: MTP’s upside is no separate draft training. Weights train and ship with the model; with official weights you need not worry about version mismatch. If you fine-tune or quantize the main model yourself, the MTP module does not auto-adapt — remeasure acceptance. Premise: the model must ship with an MTP module; without one you cannot use it.

7.3 Shape in vLLM

vLLM v0.32.0 enables this family uniformly with "method": "mtp"; older model-specific names like deepseek_mtp and qwen3_5_mtp are deprecated and auto-rewritten to mtp. Doc examples include XiaomiMiMo/MiMo-7B-Base, and Gemma 4 with the official assistant checkpoint (assistant layers share KV Cache with Target). Official advice is to start tuning from num_speculative_tokens: 1.

7.4 Farther exploration: a whole draft block in one forward

EAGLE and MTP still generate drafts token by token; longer drafts get slower. DFlash replaces the draft model with a small Block Diffusion model conditioned on Target features, generating a whole candidate block in one forward. Authors report 6×+ lossless speedups across several models and tasks, up to 2.5× faster than EAGLE-3. vLLM v0.32.0 already offers "method": "dflash", but this is 2026 work — conclusions still need more independent reproduction.

8. Hands-on: why sibling nodes must not see each other

The code below only builds the tree mask and Position IDs — no Attention compute. Node 0 is the root; 1 and 2 are siblings; 3 hangs under 1; 4 under 2. The shared prefix is outside this small matrix.

parents = [-1, 0, 0, 1, 2]
n = len(parents)
mask = [[0] * n for _ in range(n)]
depth = [0] * n

for node, parent in enumerate(parents):
    assert parent < node, "require parent ordered before child"
    depth[node] = 0 if parent == -1 else depth[parent] + 1
    cursor = node
    while cursor != -1:          # walk up the ancestor chain; mark visible
        mask[node][cursor] = 1
        cursor = parents[cursor]

causal = [[1 if j <= i else 0 for j in range(n)] for i in range(n)]
print("relative Position ID:", depth)
for i, (tree_row, causal_row) in enumerate(zip(mask, causal)):
    leaked = [j for j in range(n) if causal_row[j] and not tree_row[j]]
    print(f"node {i}: tree mask {tree_row}  lower-tri would also see {leaked}")

assert mask[4] == [1, 0, 1, 0, 1]
assert depth[1] == depth[2] == 1

Node 4 sits last in the array; a lower-triangular mask would see nodes 1 and 3, which belong to the other branch. The printout lists each node’s “extra seen” nodes.

A real implementation also handles the shared prefix, paged-KV block tables, per-request tree boundaries inside a Batch, and an Attention Kernel that supports this sparse mask. This example only shows the dataflow constraint.

9. Verification and deploy: tree compute does not auto-guarantee lossless

9.1 Medusa’s Typical Acceptance

Besides exact rejection sampling, Medusa also proposes Typical Acceptance: if a candidate’s probability under Target’s distribution exceeds an entropy-related threshold, treat it as “reasonable enough” and accept. That can accept longer stretches and, in the paper’s evals, mostly keeps generation quality; but it does not satisfy the distribution-preserving proof in §5.1 — it is an approximate scheme.

Eval reports should separate these cases clearly: frozen backbone or joint training; greedy match, exact random sampling, or Typical Acceptance. Writing only “used Medusa” does not answer what quality guarantee the output has.

9.2 Three things people mix up

Parallel forward: one Target call computes several positions or tree nodes — a compute-level optimization.
Token-level verification: candidates are checked left to right by the rule; stop at the first reject — the standard rule from §5.1 (core speculative-decoding principles).
Block Verification: decide how much of a whole candidate block to accept and how to correct — another sampling algorithm.

Block Verification proves that, while preserving the distribution, joint verification’s expected commit length is at least that of per-token verification. It needs matching probability rules; it is not the same as “putting several tokens in one Tensor.” vLLM v0.32.0 exposes this via "rejection_sample_method": "block", currently only in the V2 Model Runner sampler.

9.3 Paper structure and framework implementation must be checked separately

A framework method name is a load-and-run convention; it may not reproduce every structure in the paper. Trees, for example:

vLLM v0.32.0: eagle, eagle3, and mtp all draft as a linear chain — no tree verification.
SGLang: the EAGLE family builds a draft tree when --speculative-eagle-topk > 1, with --speculative-num-steps (draft steps) and --speculative-num-draft-tokens (nodes finally sent to verify) controlling the budget. Its default-on V2 Overlap Scheduler only supports topk=1; trees need an explicit change.

At deploy time check item by item: Target revision, helper-weight source, feature-layer convention, vocab, dtype, quantization, and the backend that actually runs. More tree nodes mean more verify compute and temporary KV; set the node budget from measured performance curves.

📝 Summary

Self-Draft reuses Target representations for cheap drafts, but helper modules need special training and bind to Target version.
Medusa uses several parallel heads to predict multiple future positions at once; farther is less accurate. After Medusa-2 joint training, the quality baseline must be the original model.
Tree Attention lets a node see only the shared prefix and ancestors; Position IDs follow depth; only one path is committed.
EAGLE drafts with “feature + one-step-ahead token”; EAGLE-2 allocates tree budget dynamically by path-confidence product.
EAGLE-3 drops the feature loss, fuses multi-layer features, and uses Training-time Test so training matches multi-step inference.
MTP ships an ordered prediction module from pretraining and uses it as draft at inference; DeepSeek-V3 reports 85%–90% second-token acceptance.
Typical Acceptance is approximate; Block Verification preserves the distribution with expected accept length ≥ per-token verification; vLLM and SGLang differ on tree support.

🎯 Self-check list

Explain the difference and cost of reusing Target representations vs a standalone small model
Explain why Medusa heads can run in parallel, and why farther heads are less accurate
Distinguish Medusa-1 vs Medusa-2 training and which quality baseline to use
Write a correct Tree Attention Mask and Position IDs for a small tree
Explain why EAGLE needs token information, and why EAGLE weights cannot be reused across Targets
Explain why EAGLE-2 uses path-confidence product rather than a node’s own confidence
List EAGLE-3’s three main changes relative to EAGLE
Contrast DeepSeek-V3’s MTP module with Medusa prediction heads
Distinguish parallel forward, token-level verification, and Block Verification, and note that Typical Acceptance does not preserve the distribution

Search