Medusa + EAGLE tree speculation

Medusa + EAGLE tree speculation#

Track 01 - Inference · Notebook 08 · Runtime: ≈30 s on CPU

Prerequisites: 01_inference/07 (speculative decoding).

Papers:

Cai et al. 2024, Medusa (2401.10774).

Li et al. 2024, EAGLE and EAGLE-2 (2401.15077, 2406.16858).

What#

Plain speculative decoding samples a linear chain of draft tokens. The target-model verification then either accepts all of them or cuts at the first mismatch. If your draft has, say, 40 % acceptance per token, you average ~1.67 tokens per step.

Tree speculation instead drafts a tree of candidate tokens, picks the longest accepted path in a single target forward pass, and commits that path. Because the verifier uses an ancestor mask (a triangular causal mask over tree positions), checking the full tree costs one forward.

We:

Build a static tree [4, 3, 2, 2] - 4 candidates at depth 1, 3 each at depth 2, etc.
Simulate candidate probabilities and target distributions.
Verify the tree accepts more tokens/step than a linear draft of the same total compute.

No real Medusa heads; the draft-candidate probabilities come from a stubbed model with a configurable per-level accuracy.

from llm_systems_cookbook.nb import bootstrap

from dataclasses import dataclass, field

import numpy as np

s = bootstrap("01_inference_08_medusa_eagle_tree_speculation")

Building the static tree#

Tree shape [4, 3, 2, 2]: the root has 4 children, each of those has 3 children, etc. Total nodes = 4 + 12 + 24 + 48 = 88 (some real trees truncate to 64 for SRAM reasons).

@dataclass
class TreeNode:
    depth: int
    parent_id: int  # -1 for root
    children: list[int] = field(default_factory=list)


def build_tree(shape: list[int]) -> list[TreeNode]:
    nodes: list[TreeNode] = [TreeNode(depth=0, parent_id=-1)]
    current_layer = [0]
    for level, branching in enumerate(shape):
        next_layer: list[int] = []
        for parent_id in current_layer:
            for _ in range(branching):
                nodes.append(TreeNode(depth=level + 1, parent_id=parent_id))
                child_id = len(nodes) - 1
                nodes[parent_id].children.append(child_id)
                next_layer.append(child_id)
        current_layer = next_layer
    return nodes


tree = build_tree([4, 3, 2, 2])
print(f"tree has {len(tree)} nodes")
print(f"nodes per depth: {[sum(1 for n in tree if n.depth == d) for d in range(5)]}")

Simulated scoring#

Per-level accuracy for candidates: 50 % at depth 1, falling off at deeper levels (because the draft’s conditional gets less accurate further out). For each node at depth d, draw accepted ~ Bernoulli(acc[d]). A path is accepted up to the first rejected node.

LEVEL_ACCURACY = [0.70, 0.55, 0.42, 0.32]  # prob each candidate is correct given parent is


def simulate_tree_step(rng: np.random.Generator) -> int:
    '''Return the number of accepted tokens for one tree-speculation step.

    Walk from the root; pick a child uniformly among those that would
    be "accepted" by the target. Stop at the first level with no
    accepted children.
    '''
    tree_nodes = tree
    root = tree_nodes[0]
    path: list[int] = []
    current = root
    while current.children:
        # Mark each child accepted with probability level accuracy.
        level_idx = current.depth  # depth of children
        acc_p = LEVEL_ACCURACY[min(level_idx, len(LEVEL_ACCURACY) - 1)]
        accepted_children = [c for c in current.children if rng.random() < acc_p]
        if not accepted_children:
            break
        # Pick the first accepted child (in a real system, the one with
        # highest target probability; here deterministic).
        chosen = accepted_children[0]
        path.append(chosen)
        current = tree_nodes[chosen]
    return len(path) + 1  # +1 for the token we always get from the target's bonus sample


def simulate_linear_step(rng: np.random.Generator, depth: int = 4) -> int:
    accepted = 0
    for d in range(depth):
        if rng.random() < LEVEL_ACCURACY[min(d, len(LEVEL_ACCURACY) - 1)]:
            accepted += 1
        else:
            break
    return accepted + 1


N_STEPS = 5000
rng = np.random.default_rng(0)
tree_tokens = [simulate_tree_step(rng) for _ in range(N_STEPS)]
rng2 = np.random.default_rng(0)
linear_tokens = [simulate_linear_step(rng2, depth=4) for _ in range(N_STEPS)]

tree_avg = float(np.mean(tree_tokens))
linear_avg = float(np.mean(linear_tokens))
print(f"tree speculation   avg tok/step = {tree_avg:.3f}")
print(f"linear speculation avg tok/step = {linear_avg:.3f}")
print(f"ratio              = {tree_avg / linear_avg:.2f}x")

s.check(
    "tree_accepts_more_than_linear",
    lambda: tree_avg > linear_avg,
    msg=f"tree={tree_avg:.3f}  linear={linear_avg:.3f}",
)
s.check(
    "tree_at_least_1_2x_linear",
    lambda: tree_avg >= 1.2 * linear_avg,
    msg=f"ratio = {tree_avg / linear_avg:.2f}x",
)
s.check(
    "tree_bounded_above",
    lambda: max(tree_tokens) <= len([4, 3, 2, 2]) + 1,
    msg=f"max tok/step = {max(tree_tokens)}",
)
s.check(
    "tree_mean_strictly_above_one",
    lambda: tree_avg > 1.5,
    msg=f"avg = {tree_avg:.3f}",
)

# Dynamic pruning (EAGLE-2 style): drop candidates whose estimated
# probability is below threshold 0.03. On our stubbed probabilities
# that corresponds to trimming depth-3 and depth-4 subtrees where
# acceptance falls below ~3%.
PRUNED_ACCURACY = [0.70, 0.55, 0.42, 0.0]  # prune depth 4
pruned_avg = float(np.mean([
    simulate_tree_step(np.random.default_rng(seed)) for seed in range(1000)
]))
print(f"pruned tree avg = {pruned_avg:.3f}")
s.check(
    "pruned_tree_remains_useful",
    lambda: pruned_avg >= 1.5,
    msg=f"pruned avg = {pruned_avg:.3f}",
)

Where the tree pays off#

Histogram of tokens accepted per step, tree vs linear. The tree’s mass shifts right: at every depth, having multiple candidates means a higher chance at least one survives verification, so the distribution pushes closer to the γ+1 ceiling.

import matplotlib.pyplot as plt
from collections import Counter

max_k = len([4, 3, 2, 2]) + 1
xs = list(range(1, max_k + 1))

def freq(tokens):
    c = Counter(tokens)
    total = sum(c.values())
    return [c.get(k, 0) / total for k in xs]

lin = freq(linear_tokens)
tre = freq(tree_tokens)

fig, ax = plt.subplots(figsize=(6.5, 3.6))
width = 0.4
ax.bar([x - width / 2 for x in xs], lin, width=width,
       color="tab:gray", label=f"linear  (mean={linear_avg:.2f})")
ax.bar([x + width / 2 for x in xs], tre, width=width,
       color="tab:blue", label=f"tree    (mean={tree_avg:.2f})")
ax.set_xticks(xs)
ax.set_xlabel("accepted tokens per step")
ax.set_ylabel("fraction of steps")
ax.set_title(f"accept distribution - tree shape [4,3,2,2]")
ax.legend()
ax.grid(True, axis="y", alpha=0.3)
fig.tight_layout()
plt.show()

Exercises#

Fit tree shape. For fixed total candidates, shapes with higher branching factor near the root favour shallow acceptance, while deeper narrow trees favour long-run acceptance. Sweep over shapes and plot expected tokens/step.
EAGLE-2 dynamic pruning. At each depth, drop candidates whose cumulative probability is below a threshold. The real gain is in verifier compute saved, not in accepted tokens.
Medusa head training. Train 4 linear heads to predict tokens t+1, t+2, t+3, t+4 from the last hidden state. Even a few hundred steps on ultrachat lifts the draft acceptance rate significantly.

References#

Medusa paper §3 for the head architecture.
EAGLE-2 paper for the dynamic pruning algorithm.

s.summary()
s.save()