Speculative Decoding - LLM Generation & Decoding

Speculative Decoding (Verification Step)

What you'll build

The accept/reject loop of speculative decoding: given the probabilities a small draft model and the large target model assign to K drafted tokens, decide how many of them to keep.

Background

Decoding a large model is memory-bound: each step reads every weight to produce one token. Speculative decoding (Leviathan et al., 2023; Chen et al., 2023) lets a cheap draft model propose K tokens, then runs the target model once on all of them in parallel. Each drafted token x_i is accepted with probability min(1, p_target(x_i) / p_draft(x_i)); at the first rejection, the rest are thrown away and the target resamples that position from the residual distribution max(0, p_target - p_draft) (normalized). This scheme produces exactly the target model's distribution, not an approximation, and typically gives a 2 to 3x speedup.

This problem implements the accept/reject loop only. Resampling the rejected position and the bonus token after a fully accepted draft need the full distributions; they are the follow-up.

Algorithm

np.random.seed(42)                       # once, at the start (in the starter)
accepted = []
for i in 0 .. K-1:
    a = 1.0 if draft_probs[i] == 0 else min(1, target_probs[i] / draft_probs[i])
    r = np.random.uniform()              # one draw per drafted token, in order
    if r < a:  accepted.append(draft_tokens[i])
    else:      break                     # reject token i and everything after it
return accepted, len(accepted)

Inputs

• draft_probs: list of K floats, p_draft(x_i) for each drafted token
• target_probs: list of K floats, p_target(x_i)
• draft_tokens: list of K ints
• target_dist_next: full target distribution after the last drafted token, or None. Not used in this problem (it is only needed for the bonus token).

Expected Return Format

A tuple (accepted_tokens, num_accepted):

• accepted_tokens: list of ints, always a prefix of draft_tokens
• num_accepted: int in [0, K], equal to len(accepted_tokens)

Edge cases and conventions

• target_probs[i] >= draft_probs[i] means certain acceptance (a = 1); target_probs[i] = 0 means certain rejection.
• Draw exactly one uniform per token you examine, in order, after the seed, so results are reproducible (with seed 42 the first draws are 0.3745, 0.9507, 0.7320, ...).

Constraints

NumPy only.

Follow-up

• Add the residual resampling on rejection and the bonus token from target_dist_next when all K are accepted. Prove that the output token is distributed exactly as p_target.
• If each token is accepted with probability alpha, the expected number of tokens per target pass is (1 - alpha^(K+1)) / (1 - alpha). How do you choose K?

Input: speculative_decode([0.5, 0.5, 0.5], [0.4, 0.4, 0.4], [10, 20, 30])
Output: ([10], 1)

Input: speculative_decode([0.3, 0.3, 0.3], [0.9, 0.9, 0.9], [1, 2, 3])
Output: ([1, 2, 3], 3)

import numpy as np

def speculative_decode(draft_probs, target_probs, draft_tokens, target_dist_next=None):
    """
    Verify draft tokens with the speculative decoding acceptance rule.

    Args:
        draft_probs: Draft-model probability of each drafted token (K floats)
        target_probs: Target-model probability of each drafted token (K floats)
        draft_tokens: The K drafted token ids
        target_dist_next: Full target distribution after the last drafted
                          token, or None (not used in this problem)

    Returns:
        accepted_tokens: Accepted prefix of draft_tokens (list of ints)
        num_accepted: len(accepted_tokens), between 0 and K
    """
    np.random.seed(42)
    # Your code here
    pass