Speculative decoding for faster generation

Question: Speculative decoding for faster generation

Companies: Cursor, Hugging Face

Difficulty: Hard

Problem statement

Implement speculative decoding: a small draft model proposes tokens and the large model verifies them, accepting matching runs to save forward passes. Reported on 1point3acres' Cursor bank; reconstructed in our own words.

Clarify the setup

Pin down four things before writing code. The interviewer usually has answers ready.

• What is a model here? A function from a token sequence to a probability distribution over the next token. The target model is slow and authoritative. The draft model is cheap and often right.
• Greedy or sampled? Greedy decoding (always take the argmax) has a simple accept rule. Sampling needs the rejection-sampling rule so the output distribution still equals the target's. Do greedy first, then sampling.
• What does a "forward pass" cost? The point of the algorithm is that the target scores k proposed tokens in one pass, at about the cost of scoring one, because decoding is memory-bandwidth bound. Model that with a score_block call that counts as one pass.
• Stop conditions. max_new tokens or an end-of-sequence token, whichever comes first.

Approach

Plain decoding runs the target once per token. Speculative decoding runs it once per block:

1. Draft. The draft model proposes k tokens one by one: d1 … dk.
2. Verify. The target runs once over prefix + d1 … dk and returns k + 1 next-token distributions: the one after prefix, after prefix + d1, and so on, up to after prefix + d1 … dk.
3. Accept a prefix of the draft. Walk the drafted tokens in order and stop at the first one the target rejects.
   • Greedy: accept di while it equals the target's argmax at position i. At the first mismatch, emit the target's argmax instead.
   • Sampling: accept di with probability min(1, p(di) / q(di)), where p is the target's distribution and q the draft's. On rejection, sample a replacement from the residual max(0, p − q), normalized.
4. Bonus token. If all k drafts were accepted, the target's last distribution (after dk) is free. Emit one more token from it.

Every block emits at least one token, so the loop never stalls. In the worst case (the draft is always wrong) it costs one target pass per token, the same as plain decoding, plus the wasted draft work.

Why the sampled version is exact. For a drafted token x, the chance it is emitted by acceptance is q(x) · min(1, p(x)/q(x)) = min(q(x), p(x)). The chance of a rejection is 1 − Σ min(p, q), and the residual puts mass max(0, p(x) − q(x)) on x, scaled by 1 / (1 − Σ min(p, q)). The sum is min(p, q) + max(0, p − q) = p(x). So the output is distributed exactly as the target alone would have sampled. This is the property interviewers want you to state: speculation changes speed, never output.

The draft does not have to be a model. In the greedy version any proposer works: a smaller model, an n-gram table, or text copied from the input. For code edits the output mostly repeats the original file, so "find where the current output's last few tokens occur in the original, and propose the next k tokens from there" accepts long runs nearly for free. This is the main trick behind fast "apply edit" features in code editors.

Code

import random
from typing import Callable, List, Optional, Sequence

Dist = List[float]
ModelFn = Callable[[Sequence[int]], Dist]


class CountingModel:
    """Wraps a next-token function and counts forward passes."""

    def __init__(self, fn: ModelFn) -> None:
        self.fn = fn
        self.calls = 0

    def next_dist(self, tokens: Sequence[int]) -> Dist:
        self.calls += 1
        return self.fn(tokens)

    def score_block(self, prefix: Sequence[int], block: Sequence[int]) -> List[Dist]:
        """One forward pass over prefix + block: len(block) + 1 next-token dists."""
        self.calls += 1
        seq = list(prefix)
        dists = [self.fn(seq)]
        for tok in block:
            seq.append(tok)
            dists.append(self.fn(seq))
        return dists


def argmax(dist: Dist) -> int:
    return max(range(len(dist)), key=dist.__getitem__)


def sample(dist: Dist, rng: random.Random) -> int:
    r = rng.random() * sum(dist)
    acc = 0.0
    last = 0
    for i, p in enumerate(dist):
        if p <= 0:
            continue
        acc += p
        last = i
        if r < acc:
            return i
    return last  # float rounding: fall back to the last token with mass


def greedy_decode(target: CountingModel, prompt: Sequence[int], max_new: int,
                  eos: Optional[int] = None) -> List[int]:
    """Baseline: one target pass per token."""
    out = list(prompt)
    for _ in range(max_new):
        tok = argmax(target.next_dist(out))
        out.append(tok)
        if tok == eos:
            break
    return out[len(prompt):]


Proposer = Callable[[Sequence[int], int], List[int]]


def model_proposer(draft: CountingModel) -> Proposer:
    def propose(seq: Sequence[int], k: int) -> List[int]:
        seq = list(seq)
        block = []
        for _ in range(k):
            tok = argmax(draft.next_dist(seq))
            block.append(tok)
            seq.append(tok)
        return block
    return propose


def speculative_greedy(target: CountingModel, propose: Proposer, prompt: Sequence[int],
                       max_new: int, k: int = 4, eos: Optional[int] = None) -> List[int]:
    """Same output as greedy_decode(target, ...), in fewer target passes."""
    out = list(prompt)
    new: List[int] = []
    while len(new) < max_new:
        block = propose(out, min(k, max_new - len(new)))
        dists = target.score_block(out, block)
        emitted = []
        for i, tok in enumerate(block):
            best = argmax(dists[i])
            emitted.append(best)          # equals tok when accepted
            if best != tok:
                break                     # first mismatch: keep target's token, drop the rest
        else:
            emitted.append(argmax(dists[len(block)]))  # bonus token
        for tok in emitted[: max_new - len(new)]:
            out.append(tok)
            new.append(tok)
            if tok == eos:
                return new
    return new


def speculative_sample(target: CountingModel, draft: CountingModel, prompt: Sequence[int],
                       max_new: int, k: int, rng: random.Random,
                       eos: Optional[int] = None) -> List[int]:
    """Output distributed exactly as sampling from the target alone."""
    out = list(prompt)
    new: List[int] = []
    while len(new) < max_new:
        kk = min(k, max_new - len(new))
        seq, block, qs = list(out), [], []
        for _ in range(kk):
            q = draft.next_dist(seq)
            tok = sample(q, rng)
            block.append(tok)
            qs.append(q)
            seq.append(tok)
        ps = target.score_block(out, block)
        emitted = []
        for i, tok in enumerate(block):
            p, q = ps[i], qs[i]
            if rng.random() < min(1.0, p[tok] / q[tok]):
                emitted.append(tok)
                continue
            residual = [max(0.0, pi - qi) for pi, qi in zip(p, q)]
            emitted.append(sample(residual, rng) if sum(residual) > 0 else sample(p, rng))
            break
        else:
            emitted.append(sample(ps[kk], rng))
        for tok in emitted[: max_new - len(new)]:
            out.append(tok)
            new.append(tok)
            if tok == eos:
                return new
    return new


def lookup_proposer(reference: Sequence[int], ngram: int = 3) -> Proposer:
    """Draft by copying: find the latest place the output's last `ngram`
    tokens occur in `reference` and propose what followed them there."""
    ref = list(reference)

    def propose(seq: Sequence[int], k: int) -> List[int]:
        tail = list(seq[-ngram:])
        if len(tail) < ngram:
            return ref[:k]
        for start in range(len(ref) - ngram, -1, -1):
            if ref[start:start + ngram] == tail:
                return ref[start + ngram:start + ngram + k]
        return []
    return propose

With an empty proposal, speculative_greedy still emits the target's next token (the bonus), so a proposer that has nothing to say degrades to plain decoding.

Worked numbers

Let α be the chance the target accepts a drafted token (assume independent), k the block size and c the draft's cost per token relative to one target pass.

• Expected tokens per target pass: 1 + α + α² + … + α^k = (1 − α^(k+1)) / (1 − α).
• With α = 0.8, k = 4: (1 − 0.328) / 0.2 = 3.36 tokens per pass.
• Each block costs 1 + k·c target-pass equivalents. With a draft 20× smaller, c ≈ 0.05, so a block costs 1.2 and the speedup is 3.36 / 1.2 ≈ 2.8×.
• Larger k is not free: at α = 0.6, k = 4 gives 2.31 tokens per pass and k = 8 gives 2.47, but the draft cost climbs from 1.2 to 1.4, so speedup drops from 1.92× to 1.76×. Pick k by measuring α on real traffic.
• Copy drafting for edits often sees α > 0.95 on unchanged regions. With k = 16: (1 − 0.95^17)/0.05 ≈ 11.6 tokens per pass at almost zero draft cost.

Test cases

• Greedy speculation returns exactly the target's greedy output, with a good draft, a useless draft and an empty proposer.
• A perfect draft needs about max_new / (k + 1) target passes; an always-wrong draft needs max_new.
• eos stops generation in the middle of an accepted block.
• max_new is never exceeded even when the bonus token would overshoot.
• The sampled version matches the target's distribution (checked empirically on one step with a fixed seed).
• Copy drafting over a reference that the target reproduces accepts whole blocks.

def _target_fn(seq):
    last = seq[-1] if seq else 0
    nxt = (last * 3 + 1) % 5
    dist = [0.05] * 5
    dist[nxt] = 0.8
    return dist


def _good_draft_fn(seq):
    return _target_fn(seq)


def _bad_draft_fn(seq):
    last = seq[-1] if seq else 0
    nxt = (last * 3 + 2) % 5
    dist = [0.05] * 5
    dist[nxt] = 0.8
    return dist


def test_greedy_matches_target_with_any_draft():
    expected = greedy_decode(CountingModel(_target_fn), [1], 20)
    for fn in (_good_draft_fn, _bad_draft_fn):
        t = CountingModel(_target_fn)
        got = speculative_greedy(t, model_proposer(CountingModel(fn)), [1], 20, k=4)
        assert got == expected
    t = CountingModel(_target_fn)
    assert speculative_greedy(t, lambda seq, k: [], [1], 20) == expected


def test_pass_counts():
    t = CountingModel(_target_fn)
    speculative_greedy(t, model_proposer(CountingModel(_good_draft_fn)), [1], 20, k=4)
    assert t.calls == 4  # 5 tokens per pass
    t = CountingModel(_target_fn)
    speculative_greedy(t, model_proposer(CountingModel(_bad_draft_fn)), [1], 20, k=4)
    assert t.calls == 20


def test_eos_and_max_new():
    expected = greedy_decode(CountingModel(_target_fn), [1], 20)
    eos = expected[6]
    t = CountingModel(_target_fn)
    got = speculative_greedy(t, model_proposer(CountingModel(_good_draft_fn)), [1], 20, k=4, eos=eos)
    assert got == expected[: expected.index(eos) + 1]
    t = CountingModel(_target_fn)
    assert len(speculative_greedy(t, model_proposer(CountingModel(_good_draft_fn)), [1], 7, k=4)) == 7

import random as _random


def test_sampling_matches_target_distribution():
    p = [0.5, 0.3, 0.2]
    q = [0.2, 0.2, 0.6]
    target = CountingModel(lambda seq: p)
    draft = CountingModel(lambda seq: q)
    rng = _random.Random(7)
    counts = [0, 0, 0]
    n = 40000
    for _ in range(n):
        counts[speculative_sample(target, draft, [0], 1, k=1, rng=rng)[0]] += 1
    for got, want in zip(counts, p):
        assert abs(got / n - want) < 0.01


def test_copy_drafting_accepts_whole_blocks():
    reference = [7, 1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3, 0]

    def copy_target(seq):
        # The target reproduces the reference after the prompt [7].
        i = len(seq)
        dist = [0.0] * 10
        dist[reference[i] if i < len(reference) else 0] = 1.0
        return dist

    t = CountingModel(copy_target)
    got = speculative_greedy(t, lookup_proposer(reference, ngram=1), [7], 12, k=6)
    assert got == reference[1:13]
    assert t.calls <= 4

Follow-ups

1. KV cache bookkeeping. Both models append the drafted tokens to their KV caches. After verification, truncate each cache to the accepted length. The target's cache already holds the accepted tokens because it scored them; only the bonus or correction token still needs a pass, and it rides along in the next block.
2. Batching. Different requests in a batch accept different lengths. Servers either pad to the longest block or keep per-request lengths in a ragged batch. Speculation helps most at small batch sizes, where the GPU is idle on memory reads; at large batch sizes it competes for compute and the win shrinks.
3. Tree speculation. Propose several candidate continuations as a tree and verify all branches in one pass with a tree attention mask. Raises the expected accepted length for the same number of target passes.
4. Draft heads instead of a draft model. Extra lightweight heads on the target predict tokens 2, 3, 4 ahead from the same hidden state. No second model to deploy and keep in sync.
5. Adaptive k. Track acceptance per request. Grow k while the draft keeps hitting, shrink it after rejections. Copy drafting works well in unchanged regions of a file and badly inside the edited lines, so a per-region policy pays off.

Interviewer expectations

Mid-level (L4)

• Implements greedy speculation correctly: draft k, verify in one pass, accept the matching prefix, take the target's token at the first mismatch.
• Remembers the bonus token and handles max_new and end-of-sequence inside a block.
• Tests that the output equals plain greedy decoding.

Senior (L5)

• Implements the sampled version with the min(1, p/q) accept rule and the residual distribution, and explains why the output distribution is unchanged.
• Works out expected tokens per pass and the speedup as a function of α, k and draft cost, and says why bigger k is not always better.
• Mentions KV cache truncation after rejections.

Staff (L6+)

• Connects it to the product: copy drafting from the original file for code edits, acceptance that varies by region, and adaptive k.
• Discusses serving: ragged acceptance in a batch, when speculation stops paying at high batch sizes, tree speculation and draft heads.
• Proposes how to measure it in production: acceptance rate, tokens per target pass and time to last token per request class.