← Inference From First Principles

Speculative Decoding

Letting a small model guess and a large one check: the acceptance rules that keep the output exact, what the guesses are worth on real text, and where the arithmetic says it should and should not pay.

The same 3-billion-parameter model, on the same GPU, produced tokens between 0.77 and 1.51 times as fast as it normally does, depending only on what it was writing and how many tokens a small helper model guessed ahead. On repetitive or structured text with six tokens guessed per step, it was 1.51 times faster. On paragraphs of prose with the same setting, it was slower than not guessing at all.

Speculative decoding is the first technique in this book that does not make one forward pass cheaper. It makes fewer of the expensive ones, by letting a cheap model guess the next few tokens and having the expensive model check all the guesses in a single pass. Whether that pays is not a property of the technique. It is a property of three numbers (how often the guesses are right, how much a guess costs, and how much a check costs) and the chapter measures each.

It also has to be correct, and the correctness is subtler than “the output is the same”. With greedy decoding the output matched in most runs and not in some, for a reason that is entirely about arithmetic. With sampling, there is an acceptance rule that provably keeps the target’s distribution, and two plausible-looking rules that do not. The chapter implements all three and tests them.

Where this sits

Chapter 5 built generation as a loop of forward passes, one per token, and noted that scoring a known sequence needs only one pass. Chapter 9 measured why a decode step is slow: it is limited by launches and by reading weights, not by arithmetic, so a pass over two tokens costs about what a pass over one does. Chapter 11 turned that into the three-ceiling model, and Chapter 10 made the cache that lets a new pass avoid redoing the past. Speculative decoding uses all of them: a pass over several new tokens is nearly free, so spend it on checking guesses.

The Inference Lab piece added here is the speculative loop: a draft model, a target model, two caches that are rolled back after every iteration, a greedy and a sampling acceptance rule, the arithmetic that predicts the speed-up from measured step times, and the experiments. It is not yet inside Chapter 15’s scheduler, which assumes one new token per sequence per step; Chapter 15’s batches would have to learn that a sequence can advance by several.

For the Pi series, the pattern that matters is what agents write. Tool calls, JSON, code and quoted text are the repetitive, structured outputs on which a small model guesses well; free prose is not. An agent’s latency budget is therefore partly decided by the shape of what it emits.

The idea, and why checking is cheap

Plain decoding runs the target model once for every token. The target’s work for one token is dominated by reading its weights and launching its kernels, which Chapter 9 showed is nearly independent of how many tokens the pass covers. Measured here, on Llama-3.2-3B-Instruct in fp16 with about 70 tokens of context:

New tokens in the pass 1 2 3 5 9
Target pass time 33.7 ms 31.6 ms 32.1 ms 31.9 ms 32.5 ms

Nine tokens cost no more than one. (The numbers are medians of 45 timings each, from a run in which the machine was in a steady state; I come back to why that qualification matters.) The pass that scores nine tokens produces a next-token distribution at every one of those positions, which is exactly what is needed to check nine guesses.

The guesses come from a draft model: a much smaller model with the same tokenizer. Here the target is Llama-3.2-3B-Instruct (3.21 billion parameters, 28 layers) and the draft is Llama-3.2-1B-Instruct (1.24 billion, 16 layers). Their vocabularies were compared at run time and are identical, which is the one hard requirement: a draft token has to mean the same thing to both.

One iteration, with \(k\) guesses:

  1. The draft model runs \(k\) times in sequence, each time extending its own continuation by one token, and the \(k\) tokens it chose are the proposals.
  2. The target runs once over the proposals, together with any tokens it has not yet seen. This pass yields \(k+1\) next-token distributions: one for each proposal’s position, and one for the position after the last.
  3. The proposals are checked in order. Each is accepted or rejected by a rule (below). At the first rejection the target supplies a token of its own and the rest of the proposals are discarded. If all are accepted, the target’s last distribution supplies a bonus token.
  4. Both caches are cut back to the accepted prefix.

Every iteration therefore produces at least one token, and up to \(k+1\), for the price of one target pass and \(k\) draft passes. Nothing is wasted in correctness terms. The question is only whether the extra draft work buys enough extra tokens.

Rolling back the cache

Both models cache keys and values (Chapter 10). Verification wrote entries for all \(k\) proposals into the target’s cache, and drafting wrote entries into the draft’s; after a rejection some of those belong to tokens that are not in the output. They must go. The library’s DynamicCache has a crop(n) that keeps the first \(n\) positions, and the loop calls it with the length of the accepted prefix. Two details are easy to get wrong. The target’s cache keeps the accepted proposals but not the token it supplied itself, which has not been through the model yet, so it goes in the next pass’s input. And the draft’s cache may be missing the last proposal, if every guess was accepted, because the draft never fed its own final proposal back in. The loop handles this by feeding the draft whatever tokens it has not seen at the start of each iteration. A test with random small models checks that the greedy loop reproduces plain decoding for several \(k\) and several prompts, which fails immediately if either cache is left misaligned. The rollback itself is cheap: about 0.2 ms per iteration, against tens of milliseconds for the model passes.

Greedy acceptance

With greedy decoding the rule is the obvious one. Proposal \(i\) is accepted if it is the target’s arg-max at that position:

def accept_greedy(proposals, target_logits):
    for i, x in enumerate(proposals):
        best = int(torch.argmax(target_logits[i]))
        if best != x:
            return i, best          # accepted i proposals; the target supplies `best`
    return len(proposals), int(torch.argmax(target_logits[len(proposals)]))

target_logits[i] predicts proposals[i]; the extra row at the end predicts the bonus token. The output is, token for token, what greedy decoding of the target alone would give. In exact arithmetic that is a theorem: each accepted token is the arg-max the target would have chosen, and so is each substituted one, given the same prefix.

In floating point it is not a theorem, and the experiment shows the difference.

Identical, up to arithmetic

I ran 12 prompts (six paragraphs of this book, three Python snippets, three repetitive or structured texts), 64 new tokens each, three rounds. Plain greedy decoding of the target repeated identically in all 24 repeats, so it is deterministic. Speculative decoding reproduced the plain output in all 36 runs at \(k = 1\) and \(k = 2\). At \(k = 3\), 4 and 6 it matched in 30 of 36. The six that did not were the same two prose prompts in every round, differing from plain decoding first at the 9th and the 27th new token, and only at those \(k\).

That looks like a bug and is not. At each divergence I recomputed the next-token logits three ways over the plain output’s context: a prefill followed by one single-token step, a prefill followed by a 5-token pass (what verification does), and one full prefill of the whole context. At the 9th token plain decoding chose " pattern" and the speculative loop " latency"; the two logits differed by 0.0078 on a top logit of 15.44, and all three recomputations favoured " latency", contradicting plain decoding’s own choice. At the 27th token plain decoding chose " parallel" and the speculative loop " model"; the logits differed by 0.0078 (two of the recomputations) or 0.0156 (the full prefill) on a top logit of 12.27, all favouring " parallel". In fp16 the spacing between adjacent representable numbers at those magnitudes is 0.0078, so the competing tokens differed by one or two units in the last place. The first case shows how little it takes: the same logits, computed from a context whose cache had been built in a different order, rounded the other way.

The conclusion is the one Chapters 2 and 10 reached about other changes of execution. A different sequence of operations on the same weights is a different computation, and where two tokens are tied to within rounding, a different computation can break the tie differently. Speculative decoding changes the shapes the target sees (a few tokens at a time, in varying groupings) and so can change which near-tied token wins. Neither output is wrong; they are two roundings of an undecided choice. The speculative output depends on \(k\), and a deployment that must be bit-reproducible against a plain baseline cannot use it without accepting this. What it can promise is that output is stable for a fixed configuration and equal to the baseline wherever the target has a margin larger than rounding.

Sampling acceptance

Sampling is where the guarantee needs real care, because “accept if it agrees” does not exist. Let \(p\) be the target’s distribution at a position and \(q\) the draft’s. The draft sampled \(x \sim q\). The rule from the two papers that introduced the technique (Leviathan et al., ICML 2023, and Chen et al., 2023, a preprint) is:

  • accept \(x\) with probability \(\min\left(1, \dfrac{p(x)}{q(x)}\right)\);
  • on rejection, discard \(x\) and sample instead from the residual distribution
\[ r(y) = \frac{\max(0,\; p(y) - q(y))}{\sum_z \max(0,\; p(z) - q(z))}. \]
The probability that this procedure outputs a particular token \(y\) is the probability that the draft proposed it and it was accepted, plus the probability that something was rejected and the residual produced \(y\):
\[ q(y)\min\left(1, \tfrac{p(y)}{q(y)}\right) \;+\; \Big(1 - \sum_z \min(p(z), q(z))\Big)\, r(y). \]
The first term is \(\min(p(y), q(y))\). The factor in front of \(r(y)\) is the total rejection probability, which equals \(\sum_z \max(0, p(z) - q(z))\), the same quantity that normalises \(r\), so the second term is \(\max(0, p(y) - q(y))\). The sum of the two is \(p(y)\), for every \(y\). The output is distributed exactly as the target, whatever the draft is. A bad draft costs speed (more rejections), never correctness.

Two natural simplifications break this, and the chapter tests both:

  • Resample from the target instead of the residual, after a rejection. It over-represents the tokens the draft already favours.
  • Keep the draft’s token and ignore \(p\). This samples the draft’s distribution, not the target’s.

The distributions \(p\) and \(q\) must be the ones actually sampled from. If sampling uses a temperature, top-\(k\) or top-\(p\) (Chapter 3), the same transformation has to be applied to both before computing the ratio, or the guarantee is about distributions nobody samples from.

Testing it

A guarantee about a distribution can be tested, and each of the tests below fails for the wrong rules:

  1. Synthetic distributions (unit tests). With known \(p\) and \(q\) over six tokens, 20,000 draws through the correct rule pass a chi-square test against \(p\); 20,000 through each wrong rule fail it. If the draft equals the target, every proposal is accepted.
  2. The whole loop (unit test). With two small random models over a vocabulary of eight tokens, 6,000 runs of the full sampling loop, with caches and rollback, produce first-two-token pairs whose frequencies match the target’s exact joint distribution. This is the test that would catch a cache misalignment in the sampling path, which a single-step test cannot.
  3. Real models. For each of the 12 prompts I took the real next-token distributions \(p\) (3B) and \(q\) (1B), drew 200,000 tokens through each rule, and compared the counts for the 60 most likely tokens, plus a pooled remainder, to \(p\) with a chi-square test. I repeated this for three seeds.

The real-model result, over 12 contexts and 3 seeds (36 tests for the correct rule):

Rule Tests with \(p < 0.01\) Total variation distance to \(p\), prose contexts, seed 0
Speculative sampling (min, residual) 1 of 36 0.017 to 0.038
Resample from the target 10 of 12 on every seed 0.14 to 0.18
Keep the draft’s token 10 of 12 on every seed 0.31 to 0.44

The correct rule’s one small \(p\)-value was 0.007, on a context that gave 0.54 and 0.76 with the other two seeds; with 36 tests, one below 0.01 is not surprising (three fell below 0.05, where 1.8 would be expected). For comparison, the draft and target distributions in the prose contexts differed by a total variation distance of 0.30 to 0.44, so the draft is a long way from the target, and the correct rule still reproduces the target to within sampling noise.

The two contexts where the wrong rules were not rejected are ones in which the draft and target distributions are almost identical (the probability that a proposal would be accepted is about 1.0). There the wrong rules are nearly right. The test has no power where the draft is already perfect, and a reader should not take a pass in such a context as evidence. A test of an acceptance rule needs contexts where the two models disagree.

This is a result about the first sampled token at twelve contexts; the loop-level test extends it to two tokens on a toy model. It is not a proof for the implementation at every position, and it does not cover top-\(k\) or top-\(p\) processing.

What a guess is worth

Let \(\alpha\) be the probability that a given proposal is accepted, assumed the same for every proposal. The number of tokens an iteration produces is the number of accepted proposals plus one, and the expected value follows from a geometric series:

\[ E[\text{tokens}] = 1 + \alpha + \alpha^2 + \dots + \alpha^k = \frac{1 - \alpha^{k+1}}{1 - \alpha}. \]
It rises with \(k\) but saturates at \(1/(1-\alpha)\). The estimate of \(\alpha\) from data is accepted proposals divided by proposals *considered*, where a rejected proposal counts as considered and the ones after it do not (they were never checked).

I ran the greedy loop over the 12 prompts at \(k\) of 1, 2, 4 and 8, 64 new tokens each, and counted:

Kind of text \(\hat\alpha\) at \(k=4\) Tokens per target pass
Prose (six paragraphs of this book) 0.74 2.95
Python code 0.95 4.56
Repetitive or structured text 0.98 4.88

Pooled over all prompts, \(\hat\alpha\) was 0.85, 0.84, 0.85 and 0.85 for \(k = 1, 2, 4, 8\): it is a property of the text and the pair of models, not of \(k\). Tokens per target pass were 1.85, 2.54, 3.63 and 4.87, against 1.85, 2.54, 3.68 and 5.09 from the formula. The formula is close up to \(k = 4\) and over-estimates at \(k = 8\), because acceptance is not independent: a prompt on which the draft is bad is bad at every position.

The shape of the prose results shows this. Of the 133 iterations on prose at \(k = 4\), 37 accepted nothing and 41 accepted all four; the middle (1, 2 or 3 accepted) held 55. The draft is either aligned with the target for a stretch, as in a phrase it can complete, or off it. Code and structured text were almost all “accept all four” (35 of 43 and 37 of 40).

What a pass costs

Speed depends on what the iteration costs relative to \(k+1\) plain steps. In an iteration the draft runs \(k\) times and the target once (over up to \(k+1\) tokens), so the speed-up in tokens per second over plain decoding is

\[ S = \frac{E[\text{tokens}] \cdot t_{\text{target}}}{k\, t_{\text{draft}} + t_{\text{verify}}} \]
where \(t_{\text{target}}\) is one plain decoding step and \(t_{\text{verify}}\) the cost of the checking pass. If verification is as cheap as one step (\(t_{\text{verify}} \approx t_{\text{target}}\)) and drafting is free, \(S\) is just \(E[\text{tokens}]\). The more interesting question is the cost of the draft, which I expected to scale with the size of the model. It does not scale with its parameter count:
Target (3B) Draft (1B) Ratio
Parameters 3.21 B 1.24 B 0.38
Weights in fp16 5.98 GiB 2.30 GiB 0.38
Layers 28 16 0.57
One decoding step (healthy run) 33.7 ms 18.3 ms 0.54
One decoding step (mean over the speed run) 44.7 ms 24.8 ms 0.56

The draft has 38% of the target’s weights and takes 54 to 56% of its time per step. That ratio tracks the layer count (57%), not the bytes. This is consistent with the launch-bound picture of Chapters 7 and 9: at batch 1 on this GPU a step is mostly a fixed number of kernel launches per layer, so a model with fewer layers is faster in proportion to its depth, not its width. I have not isolated this (it would take a model pair with the same depth and different widths), so read it as consistent with the explanation and not a test of it. The practical content is not in doubt: a draft that is 2.6 times smaller in parameters was only 1.8 times faster, and the cost ratio, which sets the whole trade-off, has to be measured, not read off a parameter count.

With the ratio of 0.56, the formula gives a map of where speculation pays. Using the measured step times (target 44.7 ms, draft 24.8 ms, verification 33 to 42 ms, depending on the pass length):

\(\alpha\) \(k=1\) \(k=2\) \(k=3\) \(k=4\) \(k=6\) \(k=8\) Best \(k\)
0.50 1.01 0.86 0.72 0.61 0.49 0.38 1
0.74 1.17 1.13 1.04 0.95 0.83 0.69 1
0.85 1.24 1.27 1.23 1.17 1.11 0.99 2
0.95 1.31 1.40 1.43 1.43 1.48 1.43 6

(These are computed from the formula with the measured times, not measured speed-ups. The verification times were taken as measured per pass length, which varied by a few milliseconds from noise.)

The break-even \(\alpha\), the acceptance rate below which speculation loses, is 0.49 at \(k=1\), 0.63 at \(k=2\), 0.77 at \(k=4\) and 0.81 at \(k=6\). Every extra guess costs a full draft step (here more than half a target step) and is only worth it if it is likely to be accepted. With a cheaper draft the picture changes a lot. If the draft cost 0.38 of the target (the weight ratio) instead of 0.56, and verification cost one target step, the same formula would give 1.47 at \(\alpha = 0.85\), \(k=4\), against 1.15 with the 0.56 ratio. That figure is arithmetic, not a measurement, and it shows how much of the loss is the cost of the draft model.

Measured speed

The experiment above predicts; this one measures. For each of 12 prompts and three rounds I ran plain decoding and speculative decoding at \(k\) of 1, 2, 3, 4 and 6 back to back, in an order that rotated from run to run, and computed each speculative run’s tokens per second relative to the plain run on the same prompt in the same round. The ratios are of adjacent runs because the machine is not stable. The GPU is a display GPU shared with a desktop, and over the course of this experiment the target’s step time went from 33.2 ms to 56.2 ms (and the draft’s from 17.3 to 32.4). A first version of the experiment measured all the plain runs first and all the speculative ones later; the machine sped up part-way through, and that version reported speed-ups of up to 1.8 times that were an artefact of the drift. It is discarded. Plain decoding in the final run averaged 24.8 tokens per second.

\(k\) Speed-up, median of per-run ratios Pooled Predicted \(\hat\alpha\) Tokens per target pass Output identical to plain
1 1.24 1.25 1.24 0.85 1.85 36 / 36
2 1.20 1.27 1.25 0.84 2.54 36 / 36
3 1.17 1.25 1.21 0.84 3.09 30 / 36
4 1.11 1.17 1.16 0.85 3.63 30 / 36
6 0.98 1.00 1.10 0.85 4.32 30 / 36

The prediction (the formula with the mean of the timings taken before and after the experiment) matches the pooled measurement to within about 0.05 up to \(k=4\) and over-estimates at \(k=6\) by 0.1. The best setting for the mixed workload is a small \(k\), 1 to 3, and the gain is about 1.2 to 1.3 times. Past \(k=4\) each additional guess costs more than it returns.

The mixed workload hides the structure. Splitting by the kind of text, with the median speed-up of each run over plain:

Kind of text \(k=2\) \(k=4\) \(k=6\)
Prose 1.09 0.94 0.77
Code 1.29 1.31 1.22
Repetitive or structured 1.44 1.49 1.51

On prose the best \(k\) is small and the gain is small, and at \(k=6\) speculation costs a quarter of the speed. On repetitive text, speed keeps rising with \(k\) and reaches 1.5 times. The table before predicted this: \(\alpha = 0.74\) at \(k=4\) gives 0.95 against the measured 0.94, and at \(k=6\) it gives 0.83 against 0.77. For code the prediction is less good: \(\alpha = 0.95\) gives 1.43 at \(k=4\) and 1.48 at \(k=6\), against 1.31 and 1.22 measured. I do not have an explanation for the gap on code. The \(\alpha\) for each kind came from the acceptance experiment and not from the timing runs, and \(\alpha\) at \(k=6\) may differ from \(\alpha\) at \(k=4\) for code; I did not record it.

Where does the time go? At \(k=4\) an iteration averaged 84 ms of drafting, 38 ms of verification and 0.2 ms of rollback, over a mixture of machine speeds. Drafting is two thirds of the iteration. That is the lever: a faster draft helps more than a better verifier, since verification is nearly free by construction.

When it loses

The experiments found one way to lose, and the arithmetic suggests others:

  • A draft that is wrong too often. Below the break-even \(\alpha\) for the chosen \(k\), speculation is slower than plain decoding; prose with \(k \ge 4\) here. This is a property of the text, so the best \(k\) is workload-dependent, and a fixed \(k\) tuned on code will hurt on prose. A system can adapt \(k\) per request from its recent acceptance. I did not implement that.
  • A draft that is too expensive. The cost ratio here is 0.56. As a draft’s step cost approaches the target’s, the break-even \(\alpha\) rises towards 1, and a draft slower than the target cannot pay at any \(\alpha\).
  • Memory. The draft’s weights (2.3 GiB) and its cache are in addition to the target’s. Both models together occupied 8.3 GiB of this 12 GiB card, and the next section shows what happens when that does not fit.
  • Verification that is no longer free. The flatness of the target’s pass time up to nine tokens is a statement about batch 1 on a launch-bound model. At a large batch the target is compute-bound and each extra token costs real arithmetic; the verification of \(k\) guesses for every sequence in the batch is then an expense. That is established background (Chapter 11’s roofline), not something measured here; batched speculation is among the things not measured.

A measurement that was not about the technique

In an earlier run of this experiment the draft model’s step took 285 to 560 ms, against 18 ms later. Nothing in the code had changed. Two of my own earlier diagnostic processes had been stopped but remained as defunct entries that still held about 4 GB of the GPU’s memory. With the 6 GiB target and the 2.3 GiB draft on top, the total no longer fitted in what the system would keep resident, and part of the draft’s weights seem to have been served from system memory across the PCIe link on every step (2.3 GiB at roughly 8 GB/s would take about 290 ms, which is the size of the delay). I did not observe the paging directly; the evidence is that the draft’s step took several times longer than the larger target’s, that the memory-query call reported zero free bytes after the target was loaded, and that everything returned to normal once the memory was released. An experiment run in that state would have concluded that the 1B draft is slower than the 3B target and that speculation can never pay. Both models were on the GPU, by every check the code made, the whole time. The timing runs that were affected are discarded, and the evidence register lists them. The lesson goes with Chapter 11’s: check where the memory is before trusting a time.

Other ways to propose

A separate draft model is one way to get guesses. Medusa (Cai et al., 2024, a preprint at the time of reading) attaches extra decoding heads to the target itself, so that one pass proposes several future tokens without a second model. Its abstract reports speed-ups of over 2.2 times when the heads are trained on a frozen backbone and 2.3 to 3.6 times when the backbone is fine-tuned together with the heads. Those are the authors’ figures, on their models and hardware; I have not run it. The trade is different from the one above: no second model in memory and no draft step to pay for, but a training step, and a changed model artefact (Chapter 2’s manifest would change). The README of DeepSeek’s DeepInfer repository lists “DSpark and confidence-scheduled verification” as a feature that decodes with draft tokens and schedules verification using confidence estimates, with a configurable draft length. That is vendor documentation of a feature; I infer nothing about its design or its speed. The three, draft model, extra heads and vendor-specific schemes, are one family of ideas with different costs, not interchangeable settings of one algorithm.

What the numbers do not show

  • One pair of models. Llama-3.2-3B with Llama-3.2-1B, fp16. The cost ratio and \(\alpha\) belong to this pair; other families, sizes and quantised drafts (Chapter 12) differ.
  • Batch 1. No batching, which is where the technique is least certain and where Chapter 15’s scheduler works.
  • A shared, drifting machine. The speed-up ratios are of adjacent runs, but the model’s inputs (step times) are means of two measurements 1.7 times apart.
  • Twelve prompts, 64 tokens. Six of them are paragraphs of this book. The code and structured kinds have three prompts each, so their medians have little support; the ranking prose < code < structured is consistent with the acceptance rates and with the formula but rests on a small sample.
  • Greedy speed only. The sampling rule’s exactness was tested; its wall-clock speed-up was not. A sampling draft’s acceptance is lower than a greedy draft’s for the same pair, because a sampled draft token is less likely to match.
  • Divergences from two prompts. The near-tie explanation fits the logit gaps, and the same prompts diverged in every round. Why those prompts: not investigated.
  • No adaptive \(k\), no tree-shaped drafts, no integration with the scheduler or the prefix cache.
  • The code gap. The model over-predicts code speed-ups at \(k=4\) and 6, and I do not know why.

Inference Lab increment

The Lab now has:

  • speculative_generate: the draft/target loop with DynamicCache.crop rollback, greedy or sampling, with per-iteration statistics (accepted counts, time in drafting, verification and rollback).
  • plain_greedy: the baseline, with per-step times.
  • accept_greedy and accept_sample: the two acceptance rules, as pure functions of proposals and logits.
  • expected_tokens_per_iteration, predicted_speedup and break_even_alpha: the cost model.
  • chi_square: the test used to check acceptance rules against a known distribution.
  • timing, acceptance, speed, divergence and distribution commands.

The rule this adds: speculation changes how often the expensive model runs, not what it computes, provided the acceptance rule is exact; whether it is faster is a measurement of three numbers, the acceptance rate, the cost of a guess and the cost of a check, and it is measured per workload. The first half of that sentence is a property of the rule and can be tested against a distribution. The second is a property of the deployment and can only be measured on it.

What you should now be able to do

  • Describe one iteration of speculative decoding and say what each cache must hold at the end of it.
  • Implement greedy and sampling acceptance, and say why the sampling rule needs the residual distribution, using the three-line derivation.
  • Explain why greedy speculative output is equal to plain output in exact arithmetic but may differ by a near-tied token in floating point.
  • Test an acceptance rule against a known distribution, and recognise a test with no power.
  • Estimate the expected tokens per iteration from an acceptance rate and \(k\), and the speed-up from step times.
  • Find the break-even acceptance rate for a draft/target pair and choose \(k\) for a workload.
  • Say when speculation will not pay: low acceptance, an expensive draft, or a target that is no longer launch- or bandwidth-bound.
  • Check that a timing is not secretly about memory residency.

Key terms

  • Draft model: a cheaper model that proposes tokens.
  • Target model: the model whose output is wanted.
  • Proposal: a token guessed by the draft.
  • Acceptance rate (\(\alpha\)): the probability that a proposal is kept.
  • Bonus token: the token the target supplies when all proposals are accepted.
  • Residual distribution: the normalised positive part of \(p - q\), sampled from after a rejection.
  • Rollback: cutting a cache back to the accepted prefix.
  • Break-even acceptance rate: the \(\alpha\) at which speculation matches plain decoding.
  • Near-tie: two tokens whose logits differ by about the numerical precision, so that rounding can pick either.

Reader exercises

  1. Add adaptive \(k\): after each iteration set \(k\) from a running estimate of \(\alpha\) and the cost model. Does it match the best fixed \(k\) on prose and on code with one setting for both?
  2. Use a 4-bit or int8 draft (Chapter 12). What happens to its cost ratio and to \(\alpha\)?
  3. Run the sampling loop end to end on the real pair and measure its speed-up at temperature 0.7 and 1.0. How much lower is the acceptance rate than greedy’s?
  4. Find out why the model over-predicts speed-ups on code at \(k=6\): record \(\hat\alpha\) per \(k\) per kind, and compare iteration times to the formula.
  5. Reproduce the near-tie divergence in float32 on a small model. What is the smallest pair of logit gaps it appears at?
  6. Implement prompt-lookup drafting: propose the tokens that followed the last occurrence of the current n-gram in the prompt. Where is it good and where useless?
  7. Extend the loop to a batch of 4 sequences. How does the verification time change with the batch, and where does the flatness of the pass time end?
  8. Make the distribution test powerful: find contexts where draft and target disagree most and compare the three rules there.

References

Evidence in this chapter

Claim Category
Speculative sampling with the min/residual rule yields exactly the target distribution Established background (the two papers); the derivation is in the text
Greedy speculation equals greedy decoding in exact arithmetic Established background
A target pass over 1 to 9 tokens takes the same time (31.6 to 33.7 ms) Book result, 45 timings per point, steady-state run
Greedy speculative output equalled plain output in 36 of 36 runs at k = 1, 2 and 30 of 36 at k = 3, 4, 6; the six differing runs were two prompts, at logit gaps of 1 to 2 fp16 steps Book result, 12 prompts x 3 rounds
\(\hat\alpha\) was 0.74 on prose, 0.95 on code, 0.98 on structured text, and 0.85 pooled whatever \(k\) Book result, 12 prompts
The draft cost 54 to 56% of the target’s step with 38% of its weights; this tracks layer count Book result for the cost ratio; the explanation is a hypothesis consistent with Chapters 7 and 9
Measured speed-up 1.25x at k=1, 1.17x at k=4, 1.00x at k=6 (pooled); prose 0.94x and structured 1.49x at k=4 Book result, interleaved runs on a drifting machine
The cost model predicts the pooled speed-up to within about 0.05 for k <= 4, and over-predicts at k=6 and for code Book result
The correct sampling rule is rejected in 1 of 36 chi-square tests at 0.01; two wrong rules in 10 of 12 contexts per seed Book result, first sampled token only
A 4 GB leftover allocation made the draft step 16 to 31 times slower (285 to 560 ms against 18 ms); the paging mechanism was inferred, not observed Book result (a discarded measurement; recorded in the evidence register)
Medusa’s reported 2.2x to 3.6x; DeepInfer’s draft-token feature Reported by the authors’ abstract and a vendor README; not reproduced
Batched speculation, adaptive k, tree drafts, sampled wall-clock speed-up, other model pairs Not measured

Implementation status

  • Manuscript: content/books/inference/17-chapter.md
  • Runnable examples: examples/ch17/speculative.py (original toy greedy demo), examples/ch17/speculative_lab.py (real-model loop, acceptance rules, cost model and experiments)
  • Tests: tests/test_ch17_speculative.py, tests/test_ch17_speculative_lab.py (cost-model limits and break-even, greedy and sampling rules on known distributions, wrong rules detected, the greedy loop equal to plain decoding for several k, a perfect draft accepting everything, a useless draft still correct, the sampling loop’s joint distribution; all on small random models on CPU, no GPU needed)
  • Evidence: research/ch17/sources.md, research/ch17/results.json, research/evidence-register.md
  • Verification command: python -m unittest discover -s tests -v

The chapter is a draft with draft = false. It makes no claim about any model pair other than the one measured, about batched serving, or about the internals of any vendor’s scheme. Chapter 18 moves from changing how many passes run to changing how each one is executed: compilation and execution graphs.