Speculative Decoding — Module 12.3

How This Lesson Fits the Module & Volume

KV caches and Flash Attention make each forward pass cheaper, but decode still advances roughly one token per large-model step. Speculative decoding uses a cheap draft model to propose several tokens, then the target model verifies them in one parallel pass—accepting a prefix that matches the target distribution.

Serving engines (vLLM, TensorRT-LLM, HF assistants) expose draft+verify pipelines to cut wall-clock latency without changing the final sampling distribution (when acceptance is done correctly).

Learning Objectives

By the end of this lesson, students should be able to:

Describe the draft-then-verify loop for speculative decoding.
Explain why verification can accept multiple tokens per target forward.
State the lossless / distribution-preserving goal of rejection sampling variants.
Identify when speculation helps (aligned draft) vs hurts (poor draft).
Sketch a minimal draft/verify API in code.
Relate speculation to KV cache updates on accept/reject.

Definition

Speculative decoding accelerates autoregressive generation by letting a smaller (or otherwise cheaper) draft model propose K tokens, then running the large target model once over those proposals to accept the longest valid prefix under a verification rule. Rejected positions resample from a corrected distribution so outputs match ordinary target sampling.

Draft + Verify Pipeline

1. Draft

Small model proposes K tokens

2. Prefill verify

Target scores all proposals

3. Accept

Keep matching prefix

4. Repair

Resample on first reject

Role	Cost	Job
Draft model	Cheap / fast	Propose speculative tokens
Target model	Expensive	Verify in parallel; define truth
Acceptance rule	CPU-light	Decide keep vs resample
KV cache	Must rewind on reject	Commit only accepted tokens

Why Multiple Tokens per Target Step?

During verification the target attends over the draft tokens as if they were a short prefill. One expensive forward yields logits at every draft position. If the draft often agrees with the target, you accept several tokens—amortizing the target’s weight-load across more than one output token.

Good draft

Same tokenizer / aligned family.
High acceptance rate.
Clear latency win.

Bad draft

Divergent predictions.
Frequent rejects.
Extra draft cost wasted.

Serving notes

Medusa / EAGLE: draft heads on target.
Tree speculation: branch candidates.
Must keep sampling unbiased.

Minimal Conceptual Loop

import torch

def speculative_step(draft, target, input_ids, K=4):
    """Illustrative: greedy accept if draft token == target argmax."""
    draft_ids = input_ids
    proposals = []
    for _ in range(K):
        logits = draft(draft_ids).logits[:, -1, :]
        tok = logits.argmax(dim=-1, keepdim=True)
        proposals.append(tok)
        draft_ids = torch.cat([draft_ids, tok], dim=-1)

    # One target forward over prompt + proposals
    cand = torch.cat([input_ids] + proposals, dim=-1)
    t_logits = target(cand).logits  # positions aligned with cand

    accepted = []
    pos = input_ids.size(1) - 1
    for i, tok in enumerate(proposals):
        pred = t_logits[:, pos + i, :].argmax(dim=-1, keepdim=True)
        if torch.equal(pred, tok):
            accepted.append(tok)
        else:
            accepted.append(pred)  # repair: take target token
            break
    return torch.cat([input_ids] + accepted, dim=-1)

# Production systems use proper rejection sampling so temperature/top-p stay exact.

Strengths and Tradeoffs

Strengths

Lower latency without changing target quality (when done right).
Exploits spare FLOPs / parallel verify.
Composes with KV cache and continuous batching.

Tradeoffs

Needs a good draft (extra model or heads).
Cache rewind logic is subtle.
Low acceptance can increase latency.

Common Misconception

“Any faster small model can replace the large one.” Speculation does not substitute the target—it only proposes. If you skip verification, you get the draft’s quality, not the target’s.

Knowledge Check

Short Answer: What are the two roles in speculative decoding? Answer: Draft (propose) and target (verify).
True/False: Verification can accept multiple tokens per target forward. Answer: True.
Multiple Choice: Correct rejection sampling aims to: (a) bias toward the draft, (b) match the target distribution, (c) drop the tokenizer. Answer: (b).
Short Answer: What happens to the KV cache on a reject? Answer: Rewind/discard speculative KV beyond the accepted prefix.
True/False: A poor draft always still speeds up decoding. Answer: False—low acceptance can hurt.
Multiple Choice: Medusa/EAGLE-style methods often: (a) train draft heads on the target, (b) delete attention, (c) freeze the tokenizer only. Answer: (a).
Short Answer: Why is one target forward enough to score K drafts? Answer: Parallel prefill-style pass over the proposed tokens.
True/False: Speculative decoding requires Flash Attention to work. Answer: False—orthogonal optimization.
Multiple Choice: The draft model should ideally share: (a) CSS theme, (b) tokenizer / similar distribution, (c) disk filesystem. Answer: (b).
Short Answer: What caching idea comes next for shared prompts? Answer: Prefix cache.

Key Takeaways

Draft proposes; target verifies—multiple accepts per expensive step.
Proper acceptance preserves the target sampling distribution.
Acceptance rate determines real speedup.
KV caches must commit only accepted tokens.
Next: Prefix Cache.

Trainer’s Guide

Hands-on idea: Measure acceptance rate of a 1B draft vs 7B target on the same prompts; estimate tokens per target forward.

Discussion prompt: When would you prefer a same-family draft vs multi-head self-speculation?

Recap: Speculative decoding trades cheap draft proposals for fewer large-model steps while verifying quality. Continue with Prefix Cache.