Speculative Decoding: Accept Step

Start from brute force

The obvious first draft tests every drafted token on its own and counts how many pass: sum(r[i] <= min(1, p[i] / q[i]) for i in range(len(p))). Each individual test is the right one, and it costs O(k) time for k drafted tokens and O(1) space, so speed is not the issue: it is wrong. On the first example it returns 2, because token 2 passes on its own with a ratio of 2.0, but the answer is 1. Token 2 was drafted as the continuation of token 1, and token 1 was rejected, so token 2 continues a prefix the output will never contain. Only an unbroken run of acceptances from the start counts.

The approach

Read the draft as a prefix: token i only means something if every token before it was kept. So scan left to right and stop at the first rejection. Each test accepts token i with probability min(1, p[i] / q[i]): certainly when the target likes it at least as much as the draft did, and in proportion when the draft was over-confident. Combined with resampling at the rejected position, which happens outside this function, that rule makes every emitted token follow the target model’s distribution exactly.

The reference solution keeps one counter, accepted. For each i it computes threshold = min(1.0, p[i] / q[i]); if r[i] <= threshold it adds 1, otherwise it hits break, so the tokens after a rejection are never even examined. It returns accepted, the length of the accepted prefix. Supplying the uniform draws as r makes the result deterministic: a real serving loop draws them at random, and because r is uniform on [0, 1), r[i] <= threshold happens with probability exactly threshold.

Walk through an example

Trace the first example, spec_accept([0.9, 0.2, 0.8], [0.5, 0.8, 0.4], [0.1, 0.5, 0.2]), one drafted token per row.

Complexity

Optimal: O(a + 1) time, where a is the number accepted, because the loop stops at the first rejection, so at most O(k) for k drafted tokens; O(1) space for the counter. Brute force (count every token that passes) is O(k) time and O(1) space: wrong, not slow.

Edge cases to test

• A ratio of 1 or more always accepts: spec_accept([0.9, 0.9], [0.3, 0.3], [0.99, 0.99]) → 2, since every r is below 1.
• The boundary is inclusive: spec_accept([0.5], [1.0], [0.5]) → 1, because r equals the threshold and the test is <=.
• Rejecting the first token: spec_accept([0.1], [0.9], [0.5]) → 0. Nothing from the draft is kept, but the pass still produces one token, because the target resamples that position.
• An empty draft: spec_accept([], [], []) → 0.
• A zero draft probability: Python raises ZeroDivisionError, while in TypeScript 0.5 / 0 is Infinity, so the token is accepted, and 0 / 0 is NaN, so it is rejected. A token sampled from the draft always has q > 0, but the division-free test r[i] * q[i] <= p[i] avoids the question.
• A later token that would pass is still discarded: in the first example token 2 has a ratio of 2.0 but comes after a rejection, so the answer stays 1.

Common mistakes

• Accepting on the raw target probability instead of the p_target / p_draft ratio — that changes the output distribution.
• On rejection, sampling from the plain target distribution instead of the residual max(0, p − q), renormalized over the vocabulary.
• Keeping tokens that came after a rejected one.
• Using < instead of <=: harmless with continuous random draws, but it fails the inclusive boundary that the statement specifies.

Follow-ups an interviewer asks next

What happens after the first rejection?

The target model resamples that position from the residual distribution max(0, p − q), renormalized, which needs the full probability vectors from both models, not just the sampled token’s entries. If every drafted token is accepted, it samples one bonus token from its own next-position distribution. Either way, one target pass yields between 1 and k + 1 tokens.

Why does the output match the target distribution exactly?

The chance of emitting token x by acceptance is q(x) · min(1, p(x)/q(x)) = min(p(x), q(x)). Rejection happens with probability 1 − Σ min(p, q), which equals Σ max(0, p − q), so resampling from the normalized residual adds exactly max(0, p(x) − q(x)). The two parts sum to p(x).

How many tokens does one target pass produce on average?

If each drafted token is accepted independently with probability α and the draft has k tokens, the expected number of tokens per target pass is (1 − α^(k+1)) / (1 − α), the formula from Leviathan et al. With α = 0.8 and k = 4 that is about 3.36. A longer draft only helps while the draft model keeps agreeing with the target.

What does the rule become under greedy decoding?

With temperature 0 the target distribution puts all its mass on one token, so a drafted token is accepted exactly when it equals the target’s argmax at that position. Verification stops at the first mismatch and emits the target’s own argmax there. The output is then token for token identical to plain greedy decoding with the target model.

Where this shows up

This accept/reject rule is what makes speculative decoding a provably lossless speedup — the output distribution is identical to the target model’s. Leviathan et al., who introduced it, reported a 2–3× speedup on T5-XXL with identical outputs; the gain in any given deployment depends on how often the draft model agrees with the target. It is the core trick behind fast LLM serving.

Reference solution

def spec_accept(p, q, r):
    accepted = 0
    for i in range(len(p)):
        threshold = min(1.0, p[i] / q[i])   # accept probability for this token
        if r[i] <= threshold:
            accepted += 1                    # keep it
        else:
            break                            # first rejection stops the run
    return accepted


print(spec_accept([0.9, 0.2, 0.8], [0.5, 0.8, 0.4], [0.1, 0.5, 0.2]))  # -> 1
print(spec_accept([0.5, 0.5], [0.5, 0.5], [0.3, 0.9]))                 # -> 2

function specAccept(p: number[], q: number[], r: number[]): number {
  let accepted = 0;
  for (let i = 0; i < p.length; i++) {
    const threshold = Math.min(1, p[i] / q[i]);   // acceptance probability
    if (r[i] <= threshold) accepted++;            // keep it
    else break;                                   // first rejection stops
  }
  return accepted;
}

console.log(specAccept([0.9, 0.2, 0.8], [0.5, 0.8, 0.4], [0.1, 0.5, 0.2]));  // -> 1
console.log(specAccept([0.5, 0.5], [0.5, 0.5], [0.3, 0.9]));                 // -> 2