Speculative decoding.

Speculative decoding makes a big model write faster by letting a small, cheap model guess the next few words, then having the big model check all the guesses in one go instead of writing each word itself. The clever part is the rule for which guesses to keep: it leaves the big model’s answers exactly as they would have been. Same words, same odds, just sooner.

On your Spark the big model spends nearly all its time waiting for its weights to arrive from memory, so checking five guesses costs about the same as writing one word. A decent guesser can make your server up to two or three times faster without changing a single answer, which is a rare kind of free lunch on the scoreboard against Fireworks. It is also a favourite interview question, because “how can it be faster and exact?” separates people who have heard of the trick from people who understand it.

You already have the two facts it rests on: decoding is limited by reading memory, not by math (Why ~17 tokens per second?), and the model keeps notes on everything it has read so far (The KV cache, by hand). Quantization attacks the same speed limit by shrinking the model; this attacks it by getting more words out of each trip to memory, and the two stack. Afterwards, Train your own draft model builds a better guesser, and Inside vLLM and SGLang shows how production engines run this alongside everything else.

Today

• Say exactly what a model’s “output distribution” is: at every step, a list of odds, one for each possible next word-piece. Say what always taking the top choice (greedy decoding) and rolling the dice by those odds (sampling) each do with it.
• Explain why checking 5 guessed word-pieces costs your Spark about the same as writing 1, using the memory-speed math you already know.
• State the rule for keeping a guess (keep it with probability min(1, p/q), where p and q are how likely the big and the small model each thought it was; otherwise draw a replacement from the leftover max(0, p − q), the part the small model under-guessed) and prove, in four lines, that the output follows the big model’s odds exactly. Not approximately. Exactly.
• Predict the speedup from three numbers (how often guesses are kept, how many guesses per round, and how expensive the guesser is next to the big model), and pick the best number of guesses.
• Build it yourself with Llama 3.2 1B guessing for Llama 3.1 8B, check with a statistics test that the output odds really match, tune it, try a version that needs no second model at all, and add a row to SCOREBOARD.md.

You need Why ~17 tokens per second? (each written word-piece reads every weight once) and the server, harness and scoreboard from Baseline: my Spark vs Fireworks. The KV cache, by hand helps but isn’t required. Plan on an evening for the idea and the proof, two or three evenings to build and test, and one benchmark run.

The Problem

Your scoreboard says the Spark decodes Llama 3.1 8B at something like 13 tokens per second, under a ceiling of about 17. You know why: every token drags all 16 GB of weights through a 273 GB/s pipe, and the GPU’s math units spend more than 99% of each step waiting. You also know the cruel part. Prefill, where the model reads a whole prompt in one pass, uses the same weight read for hundreds of tokens at once. Decode can’t, because token 5 depends on token 4, which doesn’t exist yet.

So here is a thought experiment. Suppose an oracle whispered the next four tokens to you. You could feed all four into the 8B at once, like a tiny prefill, and in one pass get the model’s prediction at every position. If the oracle was right, you just made four tokens for the price of one weight read.

There is no oracle. But there is Llama 3.2 1B, which reads about 2.5 GB per token instead of 16 and guesses the next word of ordinary text surprisingly often. It is wrong a lot, too. So the idea needs a second half: a way for the big model to check the small model’s guesses, keep the good ones, and fix the first bad one, all in that one pass.

In the diagnostic you said: “I know it follows the same distribution… a big model will stop it and correct it if it’s out of distribution.” Close, but it hides the interesting part. The big model doesn’t throw out guesses for being “out of distribution”. It throws them out at random, with a probability chosen so that the final text is statistically indistinguishable from what the big model would have written alone. Not nearly as good: the same distribution, provably. By the end you’ll have written the proof and checked it with a chi-square test on your own hardware.

The Idea

First: what an output distribution is

Before speculation makes sense, one word needs to become concrete. You have written the last line of a transformer: the final hidden vector times the output matrix gives one number per vocabulary entry. For Llama 3, that is 128,256 numbers. Those raw scores are logitslogitsThe raw scores a model outputs for each possible next token, before they are turned into probabilities. Any real number; bigger means more likely.. Push them through a softmaxsoftmaxThe function that turns a list of scores into probabilities: exponentiate each one, then divide by the total. All results are positive and add up to 1. and you get 128,256 positive numbers that add up to exactly 1.

That list is the output distribution. Not a vibe, not a style: a vector. At every step, given everything so far, the model hands you one probability per possible next token. “The capital of France is” might give ” Paris” 0.82, ” a” 0.05, ” the” 0.04, and a long tail of tiny numbers for the other 128,253 tokens.

What you do with the vector is the decoding strategydecoding strategyThe rule for turning the probability vector into one chosen token. Greedy and sampling are the two basic ones; top-k, top-p and temperature are tweaks to sampling.:

• Greedy decoding picks the biggest entry, the argmaxargmax“The argument of the maximum”: which entry in a list is the biggest, as opposed to how big it is. argmax of [0.1, 0.7, 0.2] is position 1.. Deterministic: same prompt, same output, every time.
• Sampling rolls a weighted die. Lay the probabilities end to end as slices of the interval from 0 to 1, draw a uniform random number, and take whichever slice it lands in. ” Paris” comes out 82% of the time, ” a” 5% of the time. temperaturetemperatureDivide the logits by a number T before softmax. T below 1 sharpens the distribution toward the top token; T above 1 flattens it. T near 0 approaches greedy. reshapes the vector before the roll.

Do this once per step and you get a sequence, whose probability is the product of each token’s probability given the ones before it. Two decoding procedures “follow the same distribution” if every possible output sequence has the same probability under both. That is the bar speculative decoding clears.

From here on, two letters. p is the target model’s (the 8B’s) probability vector at a step. q is the draft model’s (the 1B’s) vector at the same step, given the same prefix. Both are lists of 128,256 numbers that sum to 1. p(x) means “the entry for token x”.

Why checking is cheap

Here is the fact that makes the whole trick worth doing, and you already own it.

In Why ~17 tokens per second?, a decode step for one token reads all 16 GB of weights, does about 1 FLOP per byte, and leaves the tensor cores idle. Feeding the model 5 tokens instead of 1 in the same forward pass changes only the math side: each weight, once loaded, is used 5 times instead of once. The bytes don’t change. On a machine whose ridge pointridge pointPeak FLOP/s divided by memory bandwidth. Below it, work is limited by memory; above it, by math. On the Spark it’s in the hundreds of FLOPs per byte. sits in the hundreds of FLOPs per byte, going from 1 to 5 FLOPs per byte is still deep in memory-bound territory. The step takes about the same time.

So the 8B can score 5 positions in one pass: its probability vector after the current text, after the current text plus guess 1, after guesses 1 and 2, and so on. This is what verificationverificationRunning the target model once over the draft’s guessed tokens, to get its own probability vector at every guessed position. Same mechanics as prefill on a short block. means. It is a mini-prefill. The KV cache makes it cheap in the other dimension too: the past is already cached, so only the 5 new positions are computed.

You will measure this in the first task: time the 8B on 1, 2, 4, 8, 16, 32 and 64 new tokens on top of a cached prompt. Expect the curve to stay nearly flat for small counts, then bend upward once there’s enough math per byte to matter.

Draft, then verify

The loop has two players. The draft modeldraft modelA small, fast model that guesses the next few tokens. Its guesses never appear in the output unless the target model approves them. is cheap and often right. The target modeltarget modelThe big model whose output distribution you actually want. Here, Llama 3.1 8B. Also called the verifier. is the authority. One round:

1. Draft. The 1B generates K tokens the normal way, one at a time: d1, d2, …, dK. It records its probability vector q at each step. K small forward passes, each reading 2.5 GB instead of 16.
2. Verify. The 8B runs once over the current text plus d1..dK and returns K + 1 probability vectors: p1 (what it thinks should come after the current text), p2 (after d1), …, p(K+1) (after dK).
3. Accept, left to right. Check d1 against p1 and q1 with the rule below. If it’s kept, check d2 against p2 and q2. Stop at the first rejection.
4. Fix or bonus. At the rejected position, sample a replacement token from a corrected distribution (next section). If all K were kept, you get one more token for free: sample it from p(K+1), which the verify pass already computed.

Every round emits at least one token and at most K + 1. Even the worst round, d1 rejected immediately, still emits one correctly distributed token for one 8B pass plus the wasted drafting.

The acceptance rule

Now the part you were fascinated by. The draft proposed token x. You have two numbers: q(x), how likely the draft thought x was, and p(x), how likely the target thinks it is. What should you do?

The greedy case is easy and worth getting out of the way. With greedy decoding, the target’s “distribution” puts all its weight on its argmax. Keep x if it equals the target’s argmax; otherwise emit the target’s argmax instead and stop. That’s it. The output is token for token the same as greedy decoding with the 8B alone, because every emitted token is the 8B’s argmax given the text before it.

Sampling is where it gets clever. There is no single right answer to match. ” Paris” is right 82% of the time and ” a” is right 5% of the time. So you can’t ask “is x correct?” You have to ask “how often should x come out here, and how often does the draft propose it?”

That’s the rule, stated by both papers:

draft proposes x, sampled from q

keep x  with probability  min(1, p(x) / q(x))

if rejected, sample a replacement from the residual distribution:
    r(y)  =  max(0, p(y) − q(y))  /  Σ_z max(0, p(z) − q(z))

In words: if the target likes x at least as much as the draft does, always keep it. If the target likes it less, keep it only in the fraction p/q of cases. And when you reject, don’t draw from p again. Draw from the residual distributionresidual distributionThe part of p the draft under-covers: for each token, how much more probability the target gives it than the draft does (zero if the draft already over-covers it), rescaled to sum to 1., which only contains the tokens the draft was too stingy with.

The residual is the subtle part, and the reason your diagnostic answer missed the shape. Rejection isn’t “the guess was out of distribution, ask the big model again”. It is “the guess landed in a region the draft over-covers, so redirect this probability to the region the draft under-covers”. Every unit of probability the draft put in the wrong place gets moved, by rejection, to exactly where it was missing.

The figure below is the rule with a toy six-word vocabulary. The top left is the draft’s q, the top right the target’s p. The solid blue in both is the same thing, min(p, q): the probability that a draft is proposed and kept. The wash in p is max(0, p − q): what’s left for rejections to fill, drawn rescaled at bottom left. Press “Sample 1” a few times and read the trace. Then press “Sample 10,000”.

Things to try:

• The defaults. α is 0.70: 70% of drafts are kept. The draft over-proposes A and SAT; rejections refill THE, CAT and DOG, the tokens it under-proposes. After 10,000 samples every bar sits on its tick, and χ² is usually under 11.07, the 5% threshold for 5 degrees of freedom (about 1 run in 20 will exceed it by chance; that’s what 5% means).
• The bug. Turn on “Bug: on reject, sample p” and sample again. This is the most natural wrong version: “if the guess is bad, let the big model just pick”. It sounds right, and it is biased. After 10,000 samples χ² is in the hundreds, and it keeps growing with every batch you add.
• One guess. Switch the draft to “One guess (lookup)”. Now q puts all its weight on THE. The rule still works: keep THE with probability p(THE) = 0.40, otherwise sample from p with THE removed. This is how draft-free methods like prompt lookup stay exact.
• Closer to p. Slide toward p. α climbs to 1 and the residual empties. A perfect draft is never rejected.

Why the output is exactly p

Here is the proof. It is four lines, and it is worth being able to do on a whiteboard.

Fix one step. What is the probability that the token that comes out is some particular y? There are two disjoint ways for y to come out:

1. The draft proposed y and it was kept. The draft proposes y with probability q(y); it’s kept with probability min(1, p(y)/q(y)). Multiply:

q(y) · min(1, p(y)/q(y))  =  min(q(y), p(y))

2. Something was rejected, then the residual produced y. The chance of a rejection, over all possible proposals, is one minus the chance of a keep. Summing case 1 over every token, the chance of a keep is Σ min(p, q), which is exactly α. So:

P(reject)  =  1 − Σ_z min(p(z), q(z))
           =  Σ_z [ p(z) − min(p(z), q(z)) ]       (because Σ_z p(z) = 1)
           =  Σ_z max(0, p(z) − q(z))

That last sum is precisely the normalizer of the residual. So the probability of rejecting and then drawing y is

P(reject) · r(y)  =  Σ_z max(0, p(z) − q(z))  ·  max(0, p(y) − q(y)) / Σ_z max(0, p(z) − q(z))
                  =  max(0, p(y) − q(y))

3. Add the two cases:

P(output = y)  =  min(q(y), p(y))  +  max(0, p(y) − q(y))  =  p(y)

The last equality is just arithmetic: if p(y) ≥ q(y) it’s q(y) + p(y) − q(y); if p(y) is smaller it’s p(y) + 0. Either way, p(y). The figure was this proof drawn as bars: solid blue plus wash equals p, for every token.

Two things to notice. First, nothing assumed the draft was any good. q can be terrible; the output is still exactly p. A bad draft only costs speed. Second, the step that makes it work is the normalizer: the chance of rejecting is exactly the total size of the residual. That’s not a coincidence of these numbers, it’s forced by both p and q summing to 1.

From one token to a whole sequence. One round emits several tokens, so you need one more step, and it’s induction. Consider the first draft d1: by the argument above, the first emitted token has distribution p1, the target’s distribution after the current text. If d1 was kept, then d1 is the text so far, and d2 was drafted from q(· | text + d1) and is checked against p(· | text + d1), so by the same argument the second emitted token follows the target’s distribution given the first. Continue. When a rejection happens, the round ends and the next round starts from the true emitted text. The bonus token is sampled from p(K+1) directly. So every emitted token is distributed according to the target given all previously emitted tokens, and the probability of any full sequence is the same product of target probabilities that plain sampling gives. Same distribution, exactly.

The proof also tells you the conditions it needs, which are exactly the bugs you can write:

• q must be the distribution you actually drew from. If you draft greedily but plug the softmax of the draft’s logits in as q, the rule is wrong. A greedy draft’s q is one-hot.
• p must be computed on the true prefix. That means the target’s KV cache must not contain rejected tokens. More on that below.
• Temperature, top-k and top-p have to be applied to p before the rule (and to q before drafting). Speculation reproduces whatever p you feed it, warped or not.

How many tokens per pass, and the best K

Each round, how many tokens come out? Call α the acceptance rateacceptance rateThe probability that one draft token is kept. For sampling it equals Σ min(p, q), the overlap of the two distributions; for greedy, how often the draft’s top pick matches the target’s.. To get a clean formula, pretend each draft is kept independently with the same probability α. (Real acceptance varies a lot from token to token; the formula is a good average, not an exact law.) The first draft is kept with probability α, the first two with probability α², and so on. The round always emits one extra token, the fix or the bonus. So the expected number of tokens per target pass is

E[tokens per round]  =  1 + α + α² + ... + α^K  =  (1 − α^(K+1)) / (1 − α)

Tokens per pass isn’t speedup, because drafting isn’t free. Let c be the cost of one draft step divided by the cost of one target step. A round costs K draft steps plus one verify pass, so 1 + Kc target-passes of time (assuming verifying K + 1 positions costs about one pass, which your first task checks). Leviathan et al. give the resulting speedup:

speedup  =  (1 − α^(K+1))  /  ( (1 − α) · (1 + K·c) )

Bandwidth alone predicts c ≈ 2.5 GB / 16 GB ≈ 0.16 for the 1B drafting for the 8B. In practice it’s often worse, because the small model’s step is dominated by fixed costs (kernel launches, Python, sampling) that don’t shrink with model size. That’s the number to watch in your measurements.

Play with it until three facts feel obvious:

• α sets the ceiling. At α = 0.5 you top out around 1.6x even with nearly free drafts. At 0.9 you can reach 4x or more.
• c sets the best K, and punishes greed. At α = 0.7: c = 0.05 wants K = 6 (2.35x), c = 0.1 wants K = 4 (about 2x), c = 0.25 wants K = 2 (about 1.46x). Past the best K, extra drafts are almost never kept but always paid for.
• It can lose. Low α with high c drops below 1x. Speculation with a bad, slow draft is slower than not speculating.

Where drafts come from

Anything that produces guesses with a known q can be a draft. The proof doesn’t care. The options are a trade-off between how good the guesses are (α) and how much they cost (c).

• A smaller model from the same family. What you’ll build. Llama 3.2 1B for Llama 3.1 8B, or Qwen2.5 0.5B for Qwen2.5 7B. The draft must use the same tokenizer, so that token 1234 means the same string to both. Distilling a small model specifically to imitate your target raises α further; that’s the draft-model lesson, where you train your own.
• N-gram and prompt lookup. No model at all. prompt lookup decodingprompt lookup decodingFind the last few generated tokens somewhere earlier in the prompt or output, and propose whatever followed them there. Free to compute; works when the output copies from the input. (Apoorv Saxena, 2023) looks for the most recent 2-3 tokens earlier in the context and proposes what came next. Its q is one-hot, so acceptance is just p(x). Useless for open-ended writing, excellent for summarizing, editing code, answering questions about a document: anything that quotes its input. c is essentially zero.
• Extra heads on the target. MedusaMedusaCai et al., 2024. Several small extra output heads on the target’s last hidden state, each predicting a token further ahead, trained while the base model stays frozen. drafts from the target’s own computation, and usually verifies a small tree of alternative continuations in one pass instead of a single chain.
• Feature-level drafting. EAGLEEAGLELi et al., 2024, then EAGLE-2 and EAGLE-3. A tiny draft network that reads the target’s hidden features, not just its tokens. guesses much better than a separate small model of similar cost, and is among the highest-acceptance practical methods. vLLM and SGLang support it.
• Self-speculation. Use the target itself, with some layers skipped, as its own draft (“Draft & Verify”, Zhang et al. 2023; LayerSkip, Elhoushi et al. 2024). No second model to store, same tokenizer by construction.

Some models now ship with this built in: DeepSeek-V3 was trained with a multi-token predictionmulti-token predictionA training objective where the model also predicts tokens 2+ positions ahead. The extra prediction module can later be used as a speculative draft. module that can serve as a draft.

The KV cache has to forget

There is one piece of bookkeeping that every implementation gets wrong at least once.

During verification the target writes KV cacheKV cacheThe stored keys and values for every token processed so far, in every layer, so the next step doesn’t recompute them. Each processed position adds one entry per layer. entries for every draft token, including the ones that will be rejected. The draft model’s cache likewise contains its own guesses. If a rejected token’s keys and values stay in the cache, the next step’s attention will look at a token that isn’t in the text. The model will be conditioning on a ghost.

So after each round, both caches must be cut back to the true text, a rollbackrollbackTruncating the KV cache back to a shorter length, discarding entries for tokens that were rejected. Cheap: it’s just forgetting the last few positions.. The cleanest rule: after a round, crop each cache to length len(text) − 1, and at the start of every forward pass, feed the model whatever tokens of the text its cache doesn’t cover yet. The fix token (or bonus token) is never in either cache, so it always gets fed at the start of the next round.

With a simple contiguous cache, cropping just reduces a length; Hugging Face’s DynamicCache has a crop method for it (check the docs for your version). In a paged cache like vLLM’s, rollback means rewinding block slots, which is fiddlier; you’ll see paged caches in Paged attention and prefix caching.

Why it helps less with many users

Speculation is a way to spend idle compute. It works because single-stream decode leaves the tensor cores idle. Batching spends the same idle compute a different way: serve B users per weight read.

On top of that, acceptance is ragged: in one round user A keeps 4 drafts and user B keeps 0, so the batch has to handle texts and caches that grow at different rates. So production engines treat speculation as a latency tool for light load: turn it on, or use a larger K, when few requests are running; dial it down or off as the batch fills. That’s also why your Spark is a good place to learn it. A one-user box is exactly the regime where speculation shines.

Build It

Add to your inference/ folder from the baseline lesson:

inference/
  spec/
    common.py          # load models, forward helper (given)
    verify_cost.py     # t2
    spec_greedy.py     # t3
    spec_sample.py     # t4
    dist_test.py       # t5
    sweep.py           # t6
    lookup.py          # t7
  server_spec.py       # t8

The pair: Llama 3.1 8B Instruct (target, the model your scoreboard tracks) and Llama 3.2 1B Instruct (draft). They share the Llama 3 tokenizer, 128,256 entries. In bf16 that’s about 16 GB plus 2.5 GB of weights, far under the Spark’s 128 GB, which also leaves room to load the 8B in fp32 (about 32 GB) for one test. If you’d rather use Qwen, Qwen2.5 7B Instruct with Qwen2.5 0.5B Instruct is the equivalent pair; one gotcha there is in t1.

# spec/common.py
import time, torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache

TARGET = "meta-llama/Llama-3.1-8B-Instruct"
DRAFT = "meta-llama/Llama-3.2-1B-Instruct"

tok = AutoTokenizer.from_pretrained(TARGET)
assert tok.get_vocab() == AutoTokenizer.from_pretrained(DRAFT).get_vocab(), "tokenizers differ"

def load(name, dtype=torch.bfloat16):
    m = AutoModelForCausalLM.from_pretrained(name, dtype=dtype).to("cuda")  # older transformers: torch_dtype=
    return m.eval()

target, draft = load(TARGET), load(DRAFT)
print(target.config.vocab_size, draft.config.vocab_size)

@torch.inference_mode()
def forward(model, cache, ids):
    """Feed new token ids (a Python list) on top of `cache` (updated in place).
    Returns fp32 logits, one row per fed token: shape [len(ids), vocab]."""
    x = torch.tensor([ids], device="cuda")
    out = model(input_ids=x, past_key_values=cache, use_cache=True)
    return out.logits[0].float()

def sync_time():
    torch.cuda.synchronize()
    return time.perf_counter()

def chat_ids(user_text):
    msgs = [{"role": "user", "content": user_text}]
    return tok.apply_chat_template(msgs, add_generation_prompt=True)

def time_verify(model, prompt_ids, n, trials=10):
    """Median seconds for one forward over n new tokens on top of a cached prompt."""
    # build the cache once, then each trial: feed n arbitrary token ids, time it, cache.crop(len(prompt_ids))
    ...

def spec_greedy(prompt_ids, max_new, K):
    seq = list(prompt_ids)
    tc, dc = DynamicCache(), DynamicCache()
    forward(target, tc, seq[:-1])   # invariant: each cache covers seq[:-1] or less
    forward(draft, dc, seq[:-1])
    stats = {"rounds": 0, "accepted": 0, "rejected": 0}
    while len(seq) - len(prompt_ids) < max_new:
        # 1. Draft K tokens greedily. First feed seq[dc.get_seq_length():] (the tokens the
        #    draft hasn't seen), take argmax of the last row, then feed each draft token.
        # 2. Verify: logits = forward(target, tc, seq[tc.get_seq_length():] + drafts).
        #    The LAST K+1 rows are the target's predictions after seq, after d1, ..., after dK.
        # 3. Walk left to right: keep d_i while d_i == argmax(row_i). At the first mismatch,
        #    append argmax(row_i) instead and stop. If all K match, append argmax(row_K) (bonus).
        # 4. Extend seq. Roll back: crop both caches to len(seq) - 1.
        ...
    return seq[len(prompt_ids):][:max_new], stats

def accept_or_fix(p, q, d):
    """p, q: fp32 probability vectors [vocab]. d: drafted token id.
    Returns (token_to_emit, kept: bool)."""
    ...

def lookup_draft(seq, K, max_ngram=3):
    """Find the most recent earlier occurrence of seq's last n tokens (try n = max_ngram down to 1).
    Return up to K tokens that followed it, or [] if no match."""
    ...

python bench.py --target spark-spec --base-url http://localhost:8000/v1 \
  --model meta-llama/Llama-3.1-8B-Instruct --precision bf16 \
  --prompts prompts/short.jsonl --concurrency 1,4,16 --n 200 --max-tokens 256 --ignore-eos \
  --change "spec decoding: Llama 3.2 1B draft, K=<yours>, greedy"

Measure

Verify cost (t2). Median of 10, target cache at ~512 tokens.

The K sweep (t6). Greedy, 20 short prompts, 256 tokens each.

Distribution test (t5).

Prompt lookup (t7).

What to expect. These are rough expectations to check, not results.

• Verify cost. Nearly flat from 1 to about 8 new tokens, then rising. If 8 tokens costs 2x a single token, something other than the weight read is dominating (look for a Python loop inside your forward).
• α̂ for 1B → 8B, greedy, chat prompts. Plausibly somewhere around 0.6 to 0.8: the two models were trained on similar data and agree on most easy tokens (punctuation, the second half of a word, boilerplate phrases). It will vary a lot by prompt: higher on code and lists, lower on creative writing. Sampling at temperature 1 usually lowers it.
• c. Between the bandwidth ratio (~0.16) and something like 0.4 in eager PyTorch. This is the number most likely to surprise you.
• Speedup. With α ≈ 0.7 and c ≈ 0.3, the calculator says best K ≈ 2-3 and about 1.4x. With c ≈ 0.15 you’d see closer to 1.75x at K ≈ 3-4. A naive implementation landing between 1.3x and 2x would be a reasonable result. If you get below 1x, check c first.
• Measured vs formula tokens per round. Close but not exact, because real acceptances aren’t independent: easy stretches accept long runs, hard ones reject early.
• Prompt lookup. Little to no gain on open chat (it often finds no match). Large gains on the rename task, where most of the output literally copies the input: many tokens per round with c ≈ 0.
• Scoreboard. Tok/s at c = 1 should rise by about your speedup. TTFT should barely change (the draft’s prefill is small). Throughput at c = 16 is still roughly your single-stream speed, since the lock still serializes requests: speculation made each request faster, not the server wider. Batching is still the next big lever.

Break It

Resample from p on rejection. The most tempting bug: “if the draft is wrong, let the target pick”. Make that change in accept_or_fix and rerun t5. On the toy vectors the output becomes min(p, q) + (1 − α)·p, which for the Figure 003 vectors is [0.32, 0.325, 0.145, 0.08, 0.078, 0.052] instead of p. The text still reads fine, which is the scary part: you’d never catch it by eye. Only the test catches it. Lesson: “plausible output” is not “correct distribution”.

Skip the rollback. Delete the crop calls. In greedy mode, the identity test now fails immediately after the first rejected draft in every prompt, and the output soon degrades: the target is attending to keys and values of tokens that aren’t in the text. Notice that the first round is always fine. Bugs in speculative decoding hide in round two.

Draft greedily, claim softmax. In spec_sample, draft with argmax but keep passing the softmax q into the accept rule. The distribution test fails. The rule’s correctness depends on q being the distribution the draft actually sampled from; a greedy draft’s real q is one-hot. Fix it either way (sample from q, or use the one-hot rule from t7) and the test passes again.

Use a draft from another family. In greedy mode, swap the draft for Qwen2.5 0.5B (bypass the tokenizer assert, and replace any draft id of 128,256 or above with a valid one, or the target’s embedding lookup crashes). Token ids now mean different strings, acceptance collapses toward zero, every round wastes K draft steps, and the speedup goes below 1. A bad draft can’t break correctness, only speed.

Explain It Back

Answer each out loud before opening it.

Go Deeper

• “Fast Inference from Transformers via Speculative Decoding”, Yaniv Leviathan, Matan Kalman, Yossi Matias, ICML 2023 (arXiv 2211.17192). The paper. Section 3 has the rule, the proof, and the α, K, c analysis that Figure 004 implements, including the optimal-K discussion.
• “Accelerating Large Language Model Decoding with Speculative Sampling”, Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean-Baptiste Lespiau, Laurent Sifre, John Jumper (DeepMind), 2023 (arXiv 2302.01318). The independent discovery. Short, with a clean statement of “modified rejection sampling” and results on Chinchilla 70B.
• “Blockwise Parallel Decoding for Deep Autoregressive Models”, Mitchell Stern, Noam Shazeer, Jakob Uszkoreit, NeurIPS 2018. The greedy-only ancestor. Read it to see exactly what sampling-exact speculation added.
• “Assisted Generation: a new direction toward low-latency text generation”, Joao Gante, Hugging Face blog, 2023. A practitioner’s walk through the same loop in transformers, with measured latency numbers and the cache-rollback details.
• “Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads”, Tianle Cai et al., 2024. Extra heads instead of a draft model, plus tree-shaped verification (checking several candidate continuations in one pass).
• “EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty”, Yuhui Li, Fangyun Wei, Chao Zhang, Hongyang Zhang, 2024, and its follow-ups EAGLE-2 and EAGLE-3. Why drafting on the target’s hidden features beats drafting on tokens. The state of the art you’ll meet in vLLM and SGLang.
• Prompt Lookup Decoding, Apoorv Saxena, 2023 (GitHub: apoorvumang/prompt-lookup-decoding). A few dozen lines; the whole idea fits on a screen. Now built into transformers as prompt_lookup_num_tokens.
• vLLM documentation, “Speculative Decoding”. How a production engine exposes draft models, n-gram, Medusa and EAGLE drafts, and what it says about when speculation helps or hurts under load.

Connections

Built on: Why ~17 tokens per second? (the idle compute this lesson spends) and Baseline: my Spark vs Fireworks (the server and scoreboard). The cache rollback is a direct use of the KV cache.

It competes with batching for the same idle compute, which is why it fades at high load. Quantization stacks with it: a 4-bit target makes every verify pass cheaper, and a 4-bit draft lowers c. Inside vLLM and SGLang shows how production engines combine speculation with paged caches and continuous batching.

The biggest lever left is α. In Train your own draft model you’ll distill a small model on your 8B’s own outputs, rerun this lesson’s K sweep with it, and see how far a draft trained to imitate its target can push the acceptance rate.