15. Speculative Decoding: Draft-and-Verify Acceptance (Simplified)

Problem

Speculative decoding speeds up autoregressive generation by having a small, cheap "draft" model propose several tokens ahead, then having the large "target" model verify all of them in a single parallel forward pass. Each draft token is accepted with probability min(1, p_target(x) / p_draft(x)); the first rejection resamples from the residual distribution max(0, p_target - p_draft) (renormalized), which keeps the overall sampling distribution mathematically identical to sampling from the target model alone. If every draft token is accepted, a bonus token is sampled from the target model's next-position distribution "for free."

Given draft_probs (shape (K, vocab)), target_probs (shape (K+1, vocab)), and draft_tokens (shape (K,)), walk the K positions accepting or rejecting each draft token, resample on the first rejection, or sample a bonus token if all K are accepted.

Source: src/15_speculative_decoding_acceptance.py

def speculative_decode_step(draft_probs, target_probs, draft_tokens, generator=None) -> tuple[torch.Tensor, int]

>>> import torch
>>> torch.manual_seed(0)
>>> draft_probs = torch.tensor([[0.1, 0.6, 0.1, 0.1, 0.1]])
>>> target_probs = torch.tensor([[0.1, 0.6, 0.1, 0.1, 0.1], [0.2, 0.2, 0.2, 0.3, 0.1]])
>>> draft_tokens = torch.tensor([1])
>>> speculative_decode_step(draft_probs, target_probs, draft_tokens, torch.Generator().manual_seed(0))
(tensor([1, 1]), 1)

Step-by-Step Approach

1. Loop over each of the K draft positions (a Python loop is fine here — the acceptance decision at position i depends on whether position i-1 was accepted, so this is inherently sequential, not vectorizable across positions).
2. At position i, look up p_target = target_probs[i, token] and p_draft = draft_probs[i, token] for the draft's proposed token, and compute accept_prob = min(1, p_target / p_draft) with torch.clamp(..., max=1.0).
3. Draw a uniform random number with torch.rand((), generator=generator) and accept the draft token if it's below accept_prob; otherwise this is the first rejection.
4. On rejection, build the residual distribution torch.clamp(target_probs[i] - draft_probs[i], min=0.0), renormalize it to sum to 1, sample a replacement token from it with torch.multinomial, and stop — return immediately with everything accepted so far plus this resampled token.
5. If the loop completes without any rejection, every draft token was accepted: sample one bonus token from target_probs[K] (the extra verification position) and append it.
6. Return the finalized token sequence and how many of the K draft tokens were accepted before any rejection (or K if all were accepted).

The key insight is why the residual distribution is max(0, p_target - p_draft) rather than just p_target: since the draft's proposal already "used up" some of the probability mass at accepted tokens, resampling from raw p_target after a rejection would double-count that mass and bias the output distribution away from the target model — the residual is exactly the mass the draft under-proposed relative to the target, which is what makes the whole accept-or-resample process provably unbiased.

Reference solution

import torch


def speculative_decode_step(draft_probs, target_probs, draft_tokens, generator=None):
    K = draft_tokens.shape[0]
    accepted = []
    for i in range(K):
        token = int(draft_tokens[i].item())
        p_target = target_probs[i, token]
        p_draft = draft_probs[i, token]
        # the core speculative-decoding accept rule: accept with probability
        # min(1, p_target/p_draft) so the accepted stream is exactly
        # distributed as if sampled from the target model alone
        accept_prob = torch.clamp(p_target / p_draft, max=1.0).item()
        r = torch.rand((), generator=generator).item()
        if r < accept_prob:
            accepted.append(token)
            continue
        # rejection: resample from the residual mass the target assigns
        # above what the draft already accounted for -- this is what keeps
        # the combined accept-or-resample process unbiased
        residual = torch.clamp(target_probs[i] - draft_probs[i], min=0.0)  # (vocab,)
        residual = residual / residual.sum()
        new_token = int(torch.multinomial(residual, num_samples=1, generator=generator).item())
        accepted.append(new_token)
        return torch.tensor(accepted, dtype=torch.long), i
    # every draft token accepted -> sample one bonus token "for free" from
    # the target model's next-position distribution
    bonus_token = int(torch.multinomial(target_probs[K], num_samples=1, generator=generator).item())
    accepted.append(bonus_token)
    return torch.tensor(accepted, dtype=torch.long), K

Key Functions & Tricks

• torch.clamp(p_target / p_draft, max=1.0) — implements the min(1, ratio) acceptance probability without a Python min() call on a tensor scalar.
• torch.rand((), generator=generator) — a scalar-shaped uniform draw for the accept/reject coin flip, seeded for reproducibility.
• torch.clamp(target_probs[i] - draft_probs[i], min=0.0) — the residual distribution; clamping at 0 is essential since the raw difference can be negative wherever the draft over-proposed relative to the target.
• torch.multinomial(residual, num_samples=1, generator=...) — same categorical-sampling primitive used for the bonus token and the resampled rejection token.
• Early return inside the loop — encodes that acceptance is a sequential process that stops at the first failure, unlike the fully vectorized sampling problems earlier in this set.

How to Recognize This Pattern