Skip to main content

sample_with_scratch_pool

Function sample_with_scratch_pool 

Source
pub fn sample_with_scratch_pool(
    logits: &[f32],
    config: &SamplerConfig,
    past_tokens: &[u32],
    rng: &mut SplitMix64,
    scratch: &mut SamplerScratch,
    pool: Option<&Pool>,
) -> u32
Expand description

The same chain with the whole-vocab passes spread over the CPU pool.

WHY: at Qwen3.8’s 248 320-entry vocab the serial sampler is ~14 passes over a megabyte plus 248k exp and a select_nth over a second copy — measured on the RTX 5090 pod as the gap between bench --core and the production loop (50.9 against 46.8 tok/s with penalty+confidence alone; the temperature path pays the softmax and the partition on top). The GPU graph owns the token, so during decode the pool sits idle — this is free work.

WHAT IS PRESERVED: every value. The parallel passes are elementwise (each output depends on its own input only), the sums that feed divisions stay sequential in index order, and top-k’s threshold is the k-th largest VALUE — the same number select_nth returned. The sampled token is bit-identical to the serial chain for the same seed.