Speculative decoding: draft-and-verify loop, accept/reject against a larger model

The problem, from first principles

[03-language-model-assembly/07-greedy-decoding]'s generation loop runs one FULL forward pass through the model for every SINGLE new token, an inherent bottleneck: a large model is expensive to run, and generating 100 tokens means paying that expense 100 separate times, sequentially, with no way to parallelize across steps (each step genuinely needs the previous step's output before it can even start). Speculative decoding (Leviathan et al. 2023, Chen et al. 2023) sidesteps this with a clever trick: use a much SMALLER, cheaper "draft" model to quickly propose SEVERAL tokens ahead, then run the large "target" model just ONCE, in parallel, to check ALL of those proposed tokens at once (a single forward pass over the whole draft sequence is exactly as expensive as one ordinary step, since a Transformer processes every position in a sequence simultaneously anyway).

The verification rule is what keeps this mathematically sound: a proposed token is ACCEPTED only if the target model's own greedy choice, at that exact position, agrees with what the draft proposed. The moment they disagree, the draft is wrong from that point on (the target's own choice there is used as the correction), and everything the draft speculated AFTER that point is discarded, since it was built on a now-abandoned assumption. When the draft and target happen to agree often (a good, well-matched draft model), MULTIPLE tokens get accepted from a SINGLE expensive target forward pass, real, measurable speedup, and critically, the FINAL output is always exactly what the target model's own greedy decoding would have produced on its own, speculative decoding changes nothing about WHAT gets generated, only how efficiently it gets generated.

From theory to code

Implement speculative_decode_step(token_ids, <draft model args>, <target model args>, num_draft_tokens): draft num_draft_tokens tokens greedily with the draft model ([07-greedy-decoding]), verify them all in ONE target-model forward pass, and accept the longest PREFIX of agreement, correcting the first disagreement (if any) to the target's own choice.

Constraints

• The draft model's proposals come from [07-greedy-decoding]'s greedy_decode, run for exactly num_draft_tokens steps.
• The target model runs EXACTLY ONCE, on the ENTIRE drafted sequence at once (never once per drafted token), via [04-full-forward-pass]'s full_lm_forward with a causal mask.
• Acceptance is checked IN ORDER, starting from the first drafted token: a token is accepted only if the target's own argmax at that position matches the draft's proposal there, and checking STOPS at the first disagreement.
• The first disagreement is corrected to the TARGET's own choice (never the draft's rejected proposal), and nothing after it is included in the result.
• If every drafted token is accepted, the result is identical to running the TARGET model's own greedy_decode directly for num_draft_tokens steps.

Hints

draft_sequence = greedy_decode(token_ids, ...draft args..., num_draft_tokens)
mask = build_causal_mask(draft_sequence.shape[-1])
target_logits = full_lm_forward(draft_sequence, ...target args..., mask=mask)

accepted = draft_sequence[:, :seq_len]
for i in range(num_draft_tokens):
    draft_token = draft_sequence[0, seq_len + i]
    target_prediction = np.argmax(target_logits[0, seq_len - 1 + i, :])
    if draft_token == target_prediction:
        accepted = concat(accepted, draft_token); num_accepted += 1
    else:
        accepted = concat(accepted, target_prediction); break   # correct and stop