Module sampler 

Token sampling — temperature, top-p, top-k, min-p, repetition penalty.

Randomness comes from an explicit SplitMix64 PRNG carried by the caller: reproducible with a seed, unbiased across the whole CDF (the v1 subsec_nanos source could never pick past ~23% of it).

Structs

SamplerConfig

Sampling configuration.

SamplerScratch

Reusable per-pipeline sampling workspace. The epoch table lets the repetition penalty visit each token id once without allocating a HashSet or clearing a vocab-sized boolean vector on every decode step.

SplitMix64

SplitMix64 — tiny, fast, statistically solid for sampling.

Constants

SPARSE_TOPK_MAX

Largest top-k the sparse chain serves. Past this the dense chain is the better tool anyway.

Functions

argmax

Greedy: index of the maximum value.

argmax_penalized

Greedy over the PENALIZED logits without the working copy: one pass that applies the repetition / presence penalty and the suppress list on the fly (a membership table over the vocab, built from the past tokens) and keeps the argmax with the same tie rule as argmax (highest index among equal maxima). Bit-identical to chain + argmax for temperature 0 — the values compared are the same expressions — and it is what a greedy decode with penalties pays per token, and what a speculative round pays per draft and per verified row (nine such passes a round at k=4).

distribution_into

The distribution the sampler would draw from — the whole chain minus the draw — as a normalized vector over the vocab, in out. Greedy configs (and the filtered-out fallback) come back as a one-hot, so a caller can treat every configuration uniformly. This is what speculative SAMPLING needs from both the draft head and the verify: accept-with-min(1, p/q), correct from max(0, p − q).

draw

Draw from a normalized distribution with the caller’s RNG.

draw_sparse

Draw from a sparse distribution: inverse CDF in id order — the same walk the dense categorical_sample makes over the vocab, so a seed lands on the same token when the survivor set and probs agree.

sample

Sample next token from logits. Chain order is fixed: rep-penalty → temperature → softmax → min-p → top-k → top-p → sample.

sample_with_scratch

Sampling entry point for hot decode loops with reusable scratch storage.

sample_with_scratch_pool

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

sparse_distribution_into

The sampler chain’s distribution as a SPARSE list — the same distribution chain builds over the whole vocab, for configs with a top-k, at a fraction of the cost. The dense chain copies the vocab, exponentiates it, selects, filters and normalises it — six or seven passes over 248k floats — and every one of them past the selection touches only the k survivors. Here: penalties on a copy ONLY when there are penalties, one pooled pass that selects the top-k penalized logits, one pooled pass for the vocab-wide softmax denominator (top-p is defined against the FULL normalisation, so the denominator must see every token), and the rest over k entries.

sparse_ok

Whether config can go through the sparse chain: a real temperature and a top-k in 1..=256. Qwen’s recommended instruct settings (0.7 / top-p 0.8 / top-k 20 / presence 1.5) do.

spec_accept_or_correct

One step of speculative sampling (Leviathan et al. / Chen et al.): the draft d was drawn from q; the target distribution at the same position is p. Returns None when d is accepted (with probability min(1, p[d]/q[d])) and Some(c) when it is rejected, c drawn from the residual max(0, p − q) renormalized — which is exactly what makes the emitted token stream distributed as p, draft or no draft. When the residual is empty (p ⊆ q, so p == q on the support) the correction falls back to a draw from p itself. scratch holds the residual; the pool spreads the vocab-wide pass.

spec_accept_or_correct_sparse

spec_accept_or_correct over sparse distributions: accept the draft d with min(1, p[d]/q[d]); on rejection draw the correction from the residual max(0, p − q) over p’s support (q’s support outside p contributes nothing to the residual). Empty residual → a draw from p.

top1_prob_pool

Top-1 probability of id under a softmax at temperature temp — the per-token confidence — with the exp pass over the pool and the sum sequential in index order (bit-identical to the serial fold). Uses the scratch’s partition buffer, idle now that top-k streams.

Type Aliases

Sparse

A distribution over at most SPARSE_TOPK_MAX tokens: (id, prob) sorted by id, probs summing to 1.