Skip to main content

speculative_decode_with

Function speculative_decode_with 

Source
pub fn speculative_decode_with<D: Drafter + ?Sized>(
    decoder: &Decoder,
    prompt_tokens: &[usize],
    kv_caches: &mut [KvCache],
    drafter: &D,
    options: &SpeculativeOptions,
) -> SpeculativeDecodeResult
Expand description

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

prompt_tokens is processed as one prefill batch (one forward_batch call for the whole prompt, not one per prompt token – itself a real saving independent of speculation).

§Cache state

kv_caches may be warm. options.start_pos states where prompt_tokens begins, and must equal every cache’s current seq_len – the caches hold exactly the context preceding the prompt, and this function appends to them. On return they hold that context plus the prompt plus every generated token except the last (whose KV is not computed until it is fed, which the next call does for free by passing it as the anchor).

§Output distribution

Identical to plain token-at-a-time sampling from decoder with options.sampling, at any temperature. See the module docs.