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. Exact speculative decoding preserves the target model's output distribution; 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 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)Different sequences accept different numbers of tokens, and the batch must stay rectangular. This implementation advances every sequence by the batch minimum, 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)At a rejection, the exact algorithm resamples from the normalized residual distribution, the positive part of p minus q. This implementation uses an approximation: it samples from the target distribution directly. The repository keeps an earlier draft of the routine in test.py, so the two files show how the implementation developed.
Universal Assisted Generation: drafting across tokenizers
Standard speculative decoding uses a shared vocabulary for draft and target. Universal Assisted Generation supports different tokenizers by translating 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.
"""Re-tokenization can split the drafted text into a different number of target tokens, so the batch truncates to the shortest re-encoding. 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. When the tokenizations produce different spans, this implementation resolves the alignment by taking the first token. The translation layer handles these differences before verification.
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.