Module speculative 

Speculative decoding: propose several candidate next tokens with something cheap, then verify them all in a single Decoder::forward_batch call instead of one forward_token call per token.

Two halves, deliberately separated:

• Drafting is the Drafter trait. The only implementation in the tree is PromptLookupSpeculator, an n-gram match over the history with no model at all (the same idea as vLLM’s “prompt lookup decoding”), chosen because it needs no GPU, no second set of weights and no checkpoint to be useful. A model-based drafter (MTP head, EAGLE, dFlash) is a second impl of the same trait.
• Verification is speculative_decode_with, and it does not know or care which drafter proposed the block.

Losslessness is the property that matters most here

Speculative decoding is only worth having if it produces exactly the same distribution the target model would have produced on its own, just faster. This module implements the speculative-sampling rejection rule (Leviathan et al. 2023 / Chen et al. 2023): a draft token x proposed with draft probability q(x) is accepted with probability min(1, p(x)/q(x)), and on rejection the position is resampled from the normalised residual max(0, p - q). That rule is lossless at every temperature.

It is worth being precise about what the previous accept test – argmax(target_logits[i]) == guess – actually guaranteed, because it looks like the same thing and is not. Argmax matching is exactly the special case of the rule above at temperature = 0, where p is a point mass: p(x) is 1 when the guess is the argmax and 0 otherwise, so acceptance is certain or impossible and the residual collapses back onto the argmax. Above temperature 0 it is a different algorithm with a different output distribution – it silently biases generation toward the target’s argmax, because a draft token only survives if it happens to be the most likely one. accept_or_resample is therefore not an optimisation; it is the difference between “lossless” being true and being a claim.

The invariant is tested directly, not assumed: resampling_reproduces_the_target_distribution pushes two hundred thousand tokens through the accept/reject rule with deliberately bad draft distributions and asserts the empirical output matches the target distribution; speculative_decode_at_temperature_matches_plain_sampling compares a real decode at temperature 1.0 against the target’s own exactly enumerated per-position marginals; and speculative_decode_matches_greedy_token_for_token asserts token-for-token identity with a plain forward_token loop at temperature 0.

Structs

DraftBlock

A block of drafted tokens plus, per position, the distribution that position was drawn from.

DraftDist

The distribution one drafted position was sampled from, as its complete support: (token id, probability) pairs summing to 1.

PromptLookupSpeculator

Proposes candidate continuation tokens by looking for the longest available match of the most recent ngram_size tokens earlier in history, and returning up to max_draft_len tokens that followed that earlier occurrence. Returns an empty block if no match is found or history is too short to contain one.

SpeculativeDecodeResult

Result of a speculative decode run, with the counters that make its actual savings observable rather than just assumed.

SpeculativeOptions

Everything speculative_decode_with needs beyond the model, the prompt and the drafter.

Traits

Drafter

Proposes a block of candidate continuation tokens.

Functions

accept_or_resample

The speculative-sampling accept/reject decision for one drafted position.

speculative_decode

Greedy speculative decode over a fresh KV cache, with prompt-lookup drafting. Thin wrapper over speculative_decode_with, kept for callers that want the original no-options shape.

speculative_decode_with

Decodes options.max_new_tokens tokens, using drafter to propose candidate continuations and verifying each block in a single batched call.