Speculative decoding

Open in Colab ▶️ Run this notebook in Colab

Speculative decoding#

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

Prerequisites: 01_inference/01 (KV cache).

Papers:

  • Leviathan et al. 2022, Fast Inference from Transformers via Speculative Decoding (2211.17192).

  • Chen et al. 2023, Accelerating Large Language Model Decoding with Speculative Sampling (2302.01318).

What#

A small “draft” model proposes γ tokens. The expensive “target” model verifies all γ in one forward. Tokens are accepted with probability min(1, p_target(x) / q_draft(x)); on rejection, a residual sample is drawn from the re-normalised difference distribution. The target distribution is preserved exactly.

The acceptance ratio in one line: if the draft underestimated this token (p_target > q_draft), the target likes it more than the draft did — accept it. If the draft overestimated it (p_target < q_draft), accept only with probability p_target / q_draft so the marginal output distribution still equals the target’s. The residual sampler covers the remaining mass exactly. This is what makes speculative sampling distributionally lossless rather than a quality-throughput tradeoff.

Why it works: the target forward dominates the latency. Running it once over γ+1 positions costs about the same as running it once on a single position - you get extra tokens for free whenever the draft agrees.

We simulate with stubbed logit distributions (no real LLM). Two numerical checks that actually matter:

  1. Distributional equivalence. TV distance between the speculative output distribution and the target-only distribution is <0.05 (speculative decoding must not change the output distribution).

  2. Empirical speedup matches the closed-form E[tok/step] = (1 - α^(γ+1)) / (1 - α).

from llm_systems_cookbook.nb import bootstrap

from collections import Counter

import numpy as np

s = bootstrap("01_inference_07_speculative_decoding")

Target and draft distributions#

We work on a 20-word vocabulary. The target distribution is handcrafted per-step (random but seeded). The draft distribution is the target plus some noise - on average it agrees with the target ~70 % of the time (i.e. α ≈ 0.7).

VOCAB = 20
GAMMA = 4


def target_dist(rng: np.random.Generator) -> np.ndarray:
    logits = rng.normal(0, 1.0, size=VOCAB)
    probs = np.exp(logits - logits.max())
    return probs / probs.sum()


def draft_dist(target: np.ndarray, rng: np.random.Generator, noise: float = 0.3) -> np.ndarray:
    perturbation = rng.normal(0, noise, size=VOCAB)
    logits = np.log(target + 1e-12) + perturbation
    p = np.exp(logits - logits.max())
    return p / p.sum()

Speculative sampling loop#

For each step:

  1. Sample γ tokens from the draft.

  2. Compute target probs at all positions (in real life, one forward).

  3. For each drafted token i, accept with min(1, p_t(x)/q_d(x)).

  4. On the first rejection (or if all accepted), sample a residual token from normalize(relu(p_t - q_d)).

accept_trials = 0
accept_hits = 0


def speculative_step(rng: np.random.Generator) -> tuple[list[int], int]:
    '''Return the accepted tokens and how many were generated this step.'''
    global accept_trials, accept_hits
    target = target_dist(rng)
    draft = draft_dist(target, rng)

    accepted: list[int] = []
    drafted = rng.choice(VOCAB, size=GAMMA, p=draft)
    for x in drafted:
        p_t = float(target[x])
        q_d = float(draft[x])
        u = rng.random()
        accept_trials += 1
        if u < p_t / max(q_d, 1e-12):
            accept_hits += 1
            accepted.append(int(x))
        else:
            residual = np.clip(target - draft, 0, None)
            residual = residual / max(residual.sum(), 1e-12)
            accepted.append(int(rng.choice(VOCAB, p=residual)))
            return accepted, len(accepted)
    accepted.append(int(rng.choice(VOCAB, p=target)))
    return accepted, len(accepted)


def target_only_step(rng: np.random.Generator) -> int:
    target = target_dist(rng)
    return int(rng.choice(VOCAB, p=target))


rng_spec = np.random.default_rng(0)
rng_ref = np.random.default_rng(0)
N_STEPS = 5000
spec_tokens: list[int] = []
ref_tokens: list[int] = []
tokens_per_step: list[int] = []
for _ in range(N_STEPS):
    toks, n = speculative_step(rng_spec)
    spec_tokens.extend(toks)
    tokens_per_step.append(n)
for _ in range(len(spec_tokens)):
    ref_tokens.append(target_only_step(rng_ref))

avg_tok_per_step = float(np.mean(tokens_per_step))
print(f"avg accepted tokens per speculative step = {avg_tok_per_step:.3f}")
print(f"max γ+1 = {GAMMA + 1}")

Distributional equivalence#

TV distance between empirical speculative output distribution and target-only output distribution. Should be small.

def empirical(tokens: list[int]) -> np.ndarray:
    c = Counter(tokens)
    total = len(tokens)
    return np.array([c.get(i, 0) / total for i in range(VOCAB)])


tv = 0.5 * np.abs(empirical(spec_tokens) - empirical(ref_tokens)).sum()
print(f"TV distance (spec vs target-only) = {tv:.3f}")

s.check(
    "tv_distance_small",
    lambda: tv < 0.10,
    msg=f"TV = {tv:.3f}",
)

Empirical vs closed-form speedup#

For α ≈ 0.7 and γ = 4:

E[tok/step] = (1 - α^(γ+1)) / (1 - α)
            ≈ (1 - 0.7^5) / 0.3
            ≈ 2.76

# True alpha = fraction of per-token acceptance decisions that accepted.
alpha = accept_hits / max(accept_trials, 1)
closed_form = (1 - alpha ** (GAMMA + 1)) / (1 - alpha)
print(f"empirical alpha (direct) = {alpha:.3f}")
print(f"empirical tok/step       = {avg_tok_per_step:.3f}")
print(f"closed-form E[tok/step]  = {closed_form:.3f}")
print(f"error                    = {abs(avg_tok_per_step - closed_form) / closed_form:.1%}")

s.check(
    "alpha_estimate_reasonable",
    lambda: 0.4 <= alpha <= 0.9,
    msg=f"alpha = {alpha:.3f}",
)
s.check(
    "tok_per_step_above_one",
    lambda: avg_tok_per_step > 1.5,
    msg=f"{avg_tok_per_step:.3f}",
)
s.check(
    "closed_form_matches_empirical",
    lambda: abs(avg_tok_per_step - closed_form) / closed_form < 0.15,
    msg=f"empirical={avg_tok_per_step:.3f}  closed_form={closed_form:.3f}",
)
s.check(
    "spec_never_exceeds_max_tokens",
    lambda: max(tokens_per_step) <= GAMMA + 1,
    msg=f"max tokens/step = {max(tokens_per_step)}",
)

Draft length vs expected speedup#

Plot E[tok/step] as a function of γ for a few acceptance rates α - this is the closed-form (1 - α^(γ+1)) / (1 - α). Overlay the empirical measurement for our (α, γ=4) point. The right panel shows the wall-clock win if each target forward costs 10 ms and each draft token 1 ms.

import matplotlib.pyplot as plt

gammas = list(range(1, 9))
alphas = [0.3, 0.5, 0.7, 0.9]

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9, 3.4))
for a in alphas:
    ax1.plot(gammas, [(1 - a ** (g + 1)) / (1 - a) for g in gammas],
             "o-", label=f"α = {a}")
ax1.plot([GAMMA], [avg_tok_per_step], marker="*", markersize=14,
         color="tab:red", linestyle="", label=f"empirical (α≈{alpha:.2f})")
ax1.set_xlabel("draft length γ")
ax1.set_ylabel("E[accepted tokens / step]")
ax1.set_title("closed-form accept curve + empirical")
ax1.legend(fontsize=8)
ax1.grid(True, alpha=0.3)

# Toy cost model: target forward = 10 ms, draft token = 1 ms.
t_target, t_draft = 0.010, 0.001
target_only = 1 / t_target
spec_rate   = avg_tok_per_step / (t_target + GAMMA * t_draft)
ax2.bar(["target-only", f"speculative γ={GAMMA}"],
        [target_only, spec_rate],
        color=["tab:gray", "tab:blue"])
ax2.set_ylabel("tokens / sec (toy model)")
ax2.set_title(f"wall-clock speedup ≈ {spec_rate / target_only:.2f}x")
ax2.grid(True, axis="y", alpha=0.3)
fig.tight_layout()
plt.show()

Exercises#

  1. Vary γ. Plot E[tok/step] vs γ for α ∈ {0.3, 0.5, 0.7, 0.9}. There’s a sweet spot: too-large γ wastes draft compute on likely rejections.

  2. Quality of the draft. Reduce draft noise so α rises to 0.9; measure the speedup. The ceiling is γ+1 (you’d need a perfect draft, which defeats the purpose).

  3. Real draft + target. Pair SmolLM2-135M as draft with Qwen2.5-1.5B as target via transformers.generate + assistant_model. Measure real wall-clock speedup.

References#

  • Leviathan 2022 §3 for the derivation of the rejection rule.

  • Hugging Face transformers assistant_model docs for the production two-model loop.

s.summary()
s.save()