DL S069 · Speculative Decoding — 2–3× Free Speedup

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 69 / 80 · Module M11 · ~2 hours

The story

Here's the puzzle. Decode is memory-bound (S066): you read 14 GB of weights to produce one 2-byte token. Ridiculous. If only you could produce more tokens per memory read.

You can. The trick, from Leviathan et al. (Google, 2022) and simultaneously Chen et al. (DeepMind, 2023): use a small draft model (say Llama-3.2-1B) to generate the next $k$ tokens greedily/sampled. Then run the big target model (Llama-3-70B) once on all $k$ tokens in parallel — a single forward pass — and accept as many drafted tokens as agree with what the target would have produced. On acceptance, you got $k$ tokens for one big-model forward pass. On rejection at position $j$, you keep the first $j-1$ tokens and take the target's token at $j$ (always at least one token per iteration, so it's never worse than plain decode).

The kicker: this is provably lossless. Under proper rejection sampling, the output distribution is exactly the target's. You are not trading quality for speed. You are trading "some wasted small-model compute" for "much less wasted big-model memory bandwidth". And on real workloads it's 2–3× faster.

Every serious inference stack — vLLM, TensorRT-LLM, TGI, MLC, llama.cpp — ships speculative decoding now. Understanding it changes how you think about all of decode.

The 2022 idea in one image

Imagine a professor grading exam papers. Each paper takes 60 seconds. There are 100 papers. That's 100 minutes. Now imagine a TA who guesses each paper's grade in 3 seconds — mostly right, occasionally wrong. Professor's new workflow: TA scribbles guesses on all 100 papers (5 minutes). Professor reads them, agrees on 80%, corrects the other 20 in one sitting. Total: 5 minutes + one glance per paper (much less than 60s each because the answer is already there to verify) ≈ 20 minutes. 5× speedup, same grades.

That's speculative decoding. The TA is the draft model. The professor is the target model. The "verify all at once" trick works because reading 100 papers in a row is not 100× slower than reading one — the professor's eyes are already open. Same for a GPU: one forward pass on a 5-token sequence is not 5× slower than one forward pass on a 1-token sequence. It's memory-bound; you pay the weight-read cost once either way. The extra 4 tokens are essentially free compute.

Leviathan's original paper reported 2–3× speedups on T5. In 2024–2025, with tree-based drafts (Medusa, EAGLE), the same idea hits 3–5×. It might be the single highest-ROI inference trick invented in the last five years.

Prerequisites

• Session 066 — why decode is memory-bound (the whole reason speculative works).
• Session 053 — softmax, temperature, top-p.
• Session 068 — distillation is one way to get a good draft model.

1 · The loop, in words

Target model $M_t$ (big, slow). Draft model $M_d$ (small, fast). Speculation length $k$ (typically 4–8).

Repeat until EOS:

1. Draft: starting from the current context, autoregressively sample $k$ tokens from $M_d$: $\tilde{x}_1, \ldots, \tilde{x}_k$, recording draft probabilities $q_i(\cdot)$ at each step.
2. Target forward pass (parallel): feed the context + $\tilde{x}_1, \ldots, \tilde{x}_k$ into $M_t$ as one sequence. This is one forward pass, same cost as producing a single token in vanilla decode (memory-bound: you read weights once). Extract target probabilities $p_i(\cdot)$ at positions $1..k+1$.
3. Rejection-sample sequentially for $i = 1 \ldots k$:
   • Accept $\tilde{x}_i$ with probability $\min\left(1, \frac{p_i(\tilde{x}_i)}{q_i(\tilde{x}_i)}\right)$.
   • On acceptance: continue to $i+1$.
   • On rejection at $i$: sample a replacement from the corrected residual $p'_i(x) \propto \max(0, p_i(x) - q_i(x))$, break.
4. Free token: if all $k$ tokens were accepted, sample one more from $p_{k+1}(\cdot)$ — that's already there in the target's parallel pass output.

At the end of one iteration: you've produced between 1 and $k+1$ new tokens for one big-model forward pass. Amortized, expected tokens per big pass is $1 + \sum_i \alpha_i$ where $\alpha_i$ is per-position acceptance probability.

2 · Why it's lossless

The claim: tokens accepted by this scheme are distributed exactly as if you'd sampled from $p_i$ directly.

Sketch: consider a single position. Probability of emitting token $x$ under the scheme is:

• (draft produces $x$) × (accept it) + (draft produces something else, gets rejected) × (residual sample lands on $x$)
• $= q(x) \cdot \min(1, p(x)/q(x)) + \sum_{y \ne x} q(y) \cdot (1 - \min(1, p(y)/q(y))) \cdot \frac{\max(0, p(x) - q(x))}{\sum_z \max(0, p(z) - q(z))}$

Grind through — comes out to exactly $p(x)$. The residual construction is the entire trick: it exactly cancels the bias introduced by the draft. Same principle as rejection sampling in classical MCMC.

3 · Expected speedup

Let $\alpha$ be the per-token acceptance probability (assume constant for intuition). Expected tokens per iteration: $E[\text{tokens}] = \sum_{n=0}^{k} \alpha^n = \frac{1 - \alpha^{k+1}}{1 - \alpha}$

For $\alpha = 0.7, k = 5$: $E \approx 3.03$ tokens per big-model pass. So ~3× throughput on decode, minus overhead of the draft passes and the extra tokens in the target's forward.

Overhead calibration: target parallel forward of $k+1$ tokens vs 1 token. On memory-bound decode, this is nearly free — you're re-reading the same weights either way. On prefill (compute-bound), the extra tokens cost real FLOPs; that's why speculative decoding is a decode optimization, not a prefill one.

4 · Code: a working speculative decoder

The chunk everyone gets wrong is the sequential rejection sampling. Let's write it explicitly.

import torch
import torch.nn.functional as F
 
@torch.no_grad()
def draft(model, input_ids, k, temperature=1.0):
    """Autoregressive sample k tokens from `model`. Returns tokens (1, k) and probs list."""
    tokens, probs = [], []
    ctx = input_ids
    for _ in range(k):
        logits = model(ctx).logits[:, -1, :]  # (1, V)
        p = F.softmax(logits / temperature, dim=-1)
        tok = torch.multinomial(p, num_samples=1)  # (1, 1)
        tokens.append(tok)
        probs.append(p)
        ctx = torch.cat([ctx, tok], dim=1)
    return torch.cat(tokens, dim=1), probs  # (1, k), list of (1, V)

Now the target verification + accept/reject:

@torch.no_grad()
def spec_step(target, draft_model, input_ids, k=5, temperature=1.0):
    # 1) Draft k tokens
    drafted, q_list = draft(draft_model, input_ids, k, temperature)
    # 2) One target forward on context + drafted
    full = torch.cat([input_ids, drafted], dim=1)         # (1, T + k)
    logits = target(full).logits                          # (1, T + k, V)
    # Positions of interest: predictions for slots T..T+k inclusive
    T = input_ids.shape[1]
    p_all = F.softmax(logits[:, T - 1 : T + k, :] / temperature, dim=-1)  # (1, k+1, V)
 
    accepted = []
    for i in range(k):
        p_i = p_all[0, i]              # target prob at position i
        q_i = q_list[i][0]             # draft prob at position i
        x = drafted[0, i].item()
        r = torch.rand(1, device=input_ids.device).item()
        ratio = (p_i[x] / q_i[x]).clamp(max=1.0).item()
        if r < ratio:
            accepted.append(x)
        else:
            # rejection: sample from corrected residual
            residual = torch.clamp(p_i - q_i, min=0.0)
            residual = residual / residual.sum()
            new_x = torch.multinomial(residual, 1).item()
            accepted.append(new_x)
            return torch.tensor([accepted], device=input_ids.device)  # (1, i+1)
    # All k accepted: bonus token from target's k-th slot prediction
    bonus = torch.multinomial(p_all[0, k], 1).item()
    accepted.append(bonus)
    return torch.tensor([accepted], device=input_ids.device)  # (1, k+1)

Driver:

@torch.no_grad()
def spec_generate(target, draft_model, prompt_ids, max_new=200, k=5):
    ids = prompt_ids
    while ids.shape[1] < prompt_ids.shape[1] + max_new:
        new = spec_step(target, draft_model, ids, k=k)
        ids = torch.cat([ids, new], dim=1)
        if (new == tokenizer.eos_token_id).any():
            break
    return ids

Real implementations swap target(full).logits for a KV-cached version — the target only needs to process the newly appended tokens because its cache already covers input_ids. Same for the draft. In this simple version we recompute for clarity; in production this is critical for the speedup to materialize.

5 · Draft model choice

Options, in order of setup complexity:

• Same-family smaller model. Llama-3-70B target ↔ Llama-3.2-1B draft. Same tokenizer (essential), similar distribution → $\alpha \approx 0.7$.
• Distilled draft (S068). Train a 1B student on the target's outputs. $\alpha \approx 0.8$.
• N-gram / lookup draft (Prompt Lookup Decoding, Saxena 2023). Just grep the prompt for repeated substrings. Zero extra params, $\alpha \approx 0.3$ but no compute cost. Surprisingly good for code, RAG, structured output.
• Self-speculation (Medusa, Eagle). Add "prediction heads" to the target itself that predict multiple future tokens in one pass. No separate model, higher $\alpha$, need to train the heads.

Medusa

Cai et al. 2024: attach $K$ extra decoder heads to the target ($h_1, \ldots, h_K$), each predicting the token $k$ steps ahead. Fine-tune only these heads (base frozen) on a small dataset. At inference, sample from each head to get candidates, verify with a tree-attention pass. ~2.5× speedup, no separate draft model.

EAGLE / EAGLE-2 / EAGLE-3

Li et al. 2024. Sharper draft model that predicts the feature vector (pre-LM-head hidden state) at the next position, then applies the frozen LM head. More faithful to the target than a naive small model. State of the art in 2024: 3× on Llama-3-70B.

EAGLE-2 (2024) added a dynamic draft tree: instead of a fixed-depth candidate tree, expand deeper along high-confidence branches. Reported 4× speedup on Vicuna-33B.

EAGLE-3 (Li et al., March 2025) added "training-time test" — during draft-head training, they simulate the drift that happens when the draft's own output feeds back as input, and train against that distribution. Result: acceptance rate stays high as tree depth grows. Independent benchmarks (agentnative.dev 500-prompt suite, 2026) put EAGLE-3 at 2.89× median speedup on Llama-3-70B vs Medusa-2 at 2.21× vs a standard 1B draft at 1.95× vs prompt-lookup at 1.38×.

2025 default: EAGLE-3

As of mid-2025, vLLM v1 ships EAGLE-3 as the default speculator for English + code workloads. Medusa remains competitive when you can't afford the EAGLE draft training run (EAGLE-3 draft heads take ~8 H100-hours per target model). The industry consensus is: use EAGLE-3 if you can, prompt-lookup if you can't train, and turn spec off entirely if acceptance is below ~0.5.

Further reading — the speculative frontier:

• Cai et al., Medusa (ICML 2024) + Medusa-2 (2024).
• Li et al., EAGLE (ICML 2024), EAGLE-2, EAGLE-3 (2025).
• Fu et al., Lookahead Decoding (2024) — tree-based Jacobi iteration alternative.
• Miao et al., SpecInfer (ASPLOS 2024) — tree verification with token-tree attention masks.
• HuggingFace, Assisted Generation blog (updated 2024).

6 · When speculative decoding hurts

7 · Mermaid: one iteration

8 · How to actually run this today

Easiest path — HuggingFace assisted_generation:

from transformers import AutoModelForCausalLM, AutoTokenizer
 
tok = AutoTokenizer.from_pretrained("meta-llama/Llama-3-70B-Instruct")
target = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3-70B-Instruct", torch_dtype="auto", device_map="auto")
draft = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B-Instruct", torch_dtype="auto", device_map="auto")
 
inputs = tok("Explain speculative decoding briefly.", return_tensors="pt").to(target.device)
out = target.generate(**inputs, assistant_model=draft, max_new_tokens=200)
print(tok.decode(out[0], skip_special_tokens=True))

Prompt-lookup (zero-cost draft):

out = target.generate(**inputs, prompt_lookup_num_tokens=10, max_new_tokens=200)

vLLM: --speculative_model meta-llama/Llama-3.2-1B-Instruct --num_speculative_tokens 5.

TensorRT-LLM has both draft-target and Medusa/EAGLE built in.

Recall

Stretch: implement prompt-lookup decoding from scratch: given a prompt, find its longest suffix that appears earlier in the prompt, and use the following few tokens as the draft. Measure acceptance rate on a code-generation task vs a chat task.

In your own words: speculative decoding turns __________________________ into __________________________.

Spaced review: S053 sampling strategies, S066 KV cache (memory-bound decode is the whole reason this works).

Next session (S070): continuous batching — how servers pack multiple requests of different lengths into one GPU without wasting cycles waiting for the slowest.

Bring back tomorrow:

• The 4-step spec-decoding loop.
• Why it's lossless (rejection sampling identity).
• When it stops helping (batch > 8, tokenizer mismatch, domain drift).

Previous: ← DL S068 · Next: DL S070 →