Speculative decoding buys latency with a second, smaller model: the draft model proposes gamma tokens autoregressively, the target model scores all of them in a single forward pass, and an accept-reject rule keeps the longest prefix that the target model agrees with. Done correctly, the output distribution is exactly the target model's; the speedup comes from accepting several tokens per target pass. I implemented it from scratch on the repository's own models, then extended it to the case nobody's tutorial covers: a draft model with a different tokenizer.

The accept-reject core

For each drafted token, compare the target probability against the draft probability and accept with probability min(1, p_target divided by p_draft). A token drafted with higher confidence than the target would give gets accepted proportionally less often. Acceptance must be a prefix, so a cumulative product turns the per-token accept mask into the accepted run length.

# verify all gamma draft tokens in one target pass
full_target_logits, _ = target_model(draft_context)
target_logits = full_target_logits[:, start_pos-1:-1, :]
target_probs = F.softmax(target_logits, dim=-1)
p_target = torch.gather(target_probs, -1, draft_tokens.unsqueeze(-1)).squeeze(-1)

alpha = (p_target / p_draft).clamp(max=1.0)
u = torch.rand_like(alpha)

accepted_mask = (u <= alpha)
prefix_mask = accepted_mask.int().cumprod(dim=1).bool()
number_of_accepted_tokens = prefix_mask.sum(dim=1)

Condensed from speculative_decoding.py

Batching adds one decision the single-sequence papers skip: different sequences accept different numbers of tokens, but the batch must stay rectangular. This implementation advances every sequence by the batch minimum, conservative but simple, then repairs at the first disagreement point by sampling from the target distribution there, or takes the free bonus token when every draft was accepted.

min_n_accepted = number_of_accepted_tokens.min().item()

if min_n_accepted > 0:
    idx = torch.cat([idx, draft_tokens[:, :min_n_accepted]], dim=-1)

if min_n_accepted < gamma:
    p_target_last = target_probs[:, min_n_accepted, :]
    corrected_token = torch.multinomial(p_target_last, num_samples=1)
    idx = torch.cat([idx, corrected_token], dim=-1)
else:
    bonus_probs = F.softmax(full_target_logits[:, -1, :], dim=-1)
    bonus_token = torch.multinomial(bonus_probs, num_samples=1)
    idx = torch.cat((idx, bonus_token), dim=1)

Condensed from speculative_decoding.py

One deliberate simplification to be aware of: at a rejection, the paper resamples from the normalized residual distribution, the positive part of p minus q; this version samples from the target distribution directly, which keeps the code short at the cost of exactness in that branch. The repository keeps an earlier draft of the whole routine in test.py, and comparing the two files is a fair record of how the implementation got refined.

Universal Assisted Generation: drafting across tokenizers

Standard speculative decoding assumes draft and target share a vocabulary. The more useful case is often the opposite: the fast model you have does not speak the target's token space. Universal Assisted Generation bridges it through text. Decode the current sequence to a string, re-encode with the draft tokenizer, draft gamma tokens, decode the drafted snippet back to text, re-encode into the target vocabulary, and verify as usual.

def universal_speculative_decode(
    target_model, target_decode, target_encode,
    draft_model, draft_decode, draft_encode,
    prompt_strings, max_new_tokens, gamma, eps=1e-12
):
    """
    Performs Universal Assisted Generation (UAG) for models with different
    tokenizers on a batch of prompts, using probabilistic acceptance.
    """

speculative_decoding.py

The costs are real and worth stating. Re-tokenization can split the drafted text into a different number of target tokens, so the batch truncates to the shortest re-encoding; and getting a draft probability for a target-space token requires walking the draft model along its own tokenization of the same text, token by token. Tokenizations do not always align one-to-one, and the code says so in a comment where the alignment is resolved by taking the first token. The acceptance mathematics is unchanged; everything hard about UAG lives in the translation layer.

Training the draft to agree

Acceptance rate is the whole game: every rejection wastes drafted tokens. The file includes a temperature-scaled KL-divergence loss module for distilling a draft model toward the target's distributions, softmax against log-softmax with log_target=True and the usual batchmean reduction. A draft model distilled on the target's outputs accepts longer prefixes, which is what turns the mathematics above into wall-clock speedup.