Token sampling

Deep dive

How it works

In: logits l[0..V), the tokens so far, parameters (T, k, p, min_p, penalties), an RNG · Out: one token id · State: RNG state (one seed per request); token counts · Cost: O(V) plus a sort or selection

1. Penalties. For each seen token (Hugging Face: prompt included; llama.cpp: last 64 by default): repetition θ divides a positive logit and multiplies a negative one; frequency subtracts count·α_f, presence subtracts α_p once. CTRL divides by θ whatever the sign (θ ≈ 1.2 with greedy decoding).
2. Temperature. l ← l/T; T ≤ 0 returns the argmax.
3. Top-k. Keep the k largest logits.
4. Top-p. Softmax, sort descending, keep the smallest prefix with mass ≥ p (Holtzman et al.).
5. Min-p. Keep prob ≥ min_p · prob_max: tight when the model is sure, loose when not.
6. Draw. Renormalise, take u ~ U[0,1) from the seeded generator, walk the cumulative sum past u.

Order, read from source: Hugging Face runs repetition penalty, temperature, top-k, top-p, min-p, typical. llama.cpp's default chain is penalties, DRY, top-n-sigma, top-k, typical, top-p, min-p, XTC, temperature, then dist(seed) (greedy if temperature is last and 0). After T, a high T keeps more tokens in top-p; before T it cannot (reasoning from the formulas).

Code

rustc --edition 2021 -O. The chain, a splitmix64 draw, and asserts for seed determinism, argmax at T → 0, top-p, min-p, penalties and draw frequencies.

// logits -> token: penalties, temperature, top-k, top-p, min-p, draw (Hugging Face order)
struct Rng(u64); // splitmix64: the seed is the whole state
impl Rng {
    fn next(&mut self) -> u64 {
        self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15);
        let mut z = self.0;
        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
        z ^ (z >> 31)
    }
    fn unit(&mut self) -> f64 { (self.next() >> 11) as f64 / (1u64 << 53) as f64 }
}

#[derive(Clone, Copy)]
struct Params { temp: f64, top_k: usize, top_p: f64, min_p: f64, rep: f64, freq: f64, pres: f64 }

fn softmax(l: &[f64]) -> Vec<f64> {
    let m = l.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
    let e: Vec<f64> = l.iter().map(|x| (x - m).exp()).collect(); // exp(-inf) = 0 for dropped tokens
    let s: f64 = e.iter().sum();
    e.iter().map(|x| x / s).collect()
}

/// Processed logits; dropped tokens are -inf.
fn warp(logits: &[f64], history: &[usize], p: &Params) -> Vec<f64> {
    let mut l = logits.to_vec();
    let mut count = vec![0u32; l.len()];
    for &t in history { count[t] += 1; }
    for (i, &c) in count.iter().enumerate().filter(|(_, &c)| c > 0) {
        l[i] = if l[i] <= 0.0 { l[i] * p.rep } else { l[i] / p.rep }; // repetition, sign-safe
        l[i] -= c as f64 * p.freq + p.pres;                           // frequency, presence
    }
    for x in l.iter_mut() { *x /= p.temp; }
    let mut order: Vec<usize> = (0..l.len()).collect();
    order.sort_by(|&a, &b| l[b].partial_cmp(&l[a]).unwrap());        // descending
    let mut keep = vec![true; l.len()];
    for &i in order.iter().skip(p.top_k.max(1)) { keep[i] = false; } // top-k
    let masked = |keep: &[bool], l: &[f64]| -> Vec<f64> {
        l.iter().zip(keep).map(|(&x, &k)| if k { x } else { f64::NEG_INFINITY }).collect()
    };
    let pr = softmax(&masked(&keep, &l));
    let mut cum = 0.0;
    for &i in &order {
        if cum >= p.top_p { keep[i] = false; }                       // smallest prefix, mass >= p
        cum += pr[i];
    }
    let pr = softmax(&masked(&keep, &l));
    let thr = p.min_p * pr[order[0]];                                // min-p
    for &i in &order { if pr[i] < thr { keep[i] = false; } }
    masked(&keep, &l)
}

fn sample(logits: &[f64], history: &[usize], p: &Params, rng: &mut Rng) -> usize {
    if p.temp <= 0.0 { // greedy: temperature -> 0 keeps only the argmax
        return (0..logits.len()).rev().max_by(|&a, &b| logits[a].partial_cmp(&logits[b]).unwrap()).unwrap();
    }
    let pr = softmax(&warp(logits, history, p));                     // renormalised over survivors
    let (u, mut acc) = (rng.unit(), 0.0);
    for (i, &q) in pr.iter().enumerate() {
        acc += q;
        if u < acc && q > 0.0 { return i; }
    }
    pr.iter().rposition(|&q| q > 0.0).unwrap()                       // rounding: last survivor
}

fn main() {
    let logits = [3.0, 2.5, 2.0, 1.0, 0.5, 0.0, -1.0, -2.0];
    let p = Params { temp: 0.8, top_k: 6, top_p: 0.9, min_p: 0.05, rep: 1.0, freq: 0.0, pres: 0.0 };
    // same seed => same sequence
    let run = |seed: u64, p: &Params| { let (mut r, mut h) = (Rng(seed), vec![]); for _ in 0..50 { h.push(sample(&logits, &h, p, &mut r)); } h };
    assert_eq!(run(7, &p), run(7, &p));
    assert_ne!(run(7, &p), run(8, &p));
    // T -> 0 is argmax
    let mut r = Rng(1);
    assert!((0..1000).all(|_| sample(&logits, &[], &Params { temp: 1e-4, ..p }, &mut r) == 0));
    // top-p: smallest prefix with mass >= p
    let q = Params { temp: 1.0, top_k: 8, min_p: 0.0, ..p };
    let base = softmax(&logits);
    let w = warp(&logits, &[], &q);
    let kept: Vec<usize> = (0..8).filter(|&i| w[i].is_finite()).collect();
    let mass: f64 = kept.iter().map(|&i| base[i]).sum();
    assert!(mass >= q.top_p && mass - base[*kept.last().unwrap()] < q.top_p);
    // min-p 0.2 drops token 3: exp(-2) = 0.135 < 0.2
    let m = warp(&logits, &[], &Params { temp: 1.0, top_k: 8, top_p: 1.0, min_p: 0.2, ..p });
    assert!(m[..3].iter().all(|x| x.is_finite()) && m[3..].iter().all(|x| *x == f64::NEG_INFINITY));
    assert_eq!(sample(&logits, &[], &Params { temp: 0.0, ..p }, &mut r), 0);
    // 400k draws match the renormalised distribution
    let target = softmax(&warp(&logits, &[], &p));
    let mut hits = [0f64; 8];
    for _ in 0..400_000 { hits[sample(&logits, &[], &p, &mut r)] += 1.0; }
    for i in 0..8 { assert!((hits[i] / 400_000.0 - target[i]).abs() < 0.005); }
    // penalties (token 1 seen twice, 7 once)
    let w = warp(&logits, &[1, 1, 7], &Params { temp: 1.0, top_k: 8, top_p: 1.0, min_p: 0.0, rep: 2.0, freq: 0.5, pres: 0.25 });
    assert!((w[1] - (2.5 / 2.0 - 1.0 - 0.25)).abs() < 1e-12 && (w[7] - (-4.0 - 0.5 - 0.25)).abs() < 1e-12 && w[0] == 3.0);
    println!("all checks passed");
}

Complexity and constants

• Penalties O(V + H) over history H; sort O(V log V); select_nth_unstable O(V + k log k); one u64 of RNG state.
• Measured on this Mac (Apple M4 Pro), one thread, f64, T = 0.8, k = 40, p = 0.9, min-p 0.05, mean of 200 draws: V = 32,000 takes 680 µs (page code), 44 µs (top-k selected first), 11-16 µs (greedy); V = 128,000 takes 3.3 ms, 155 µs, 50 µs. CPU code, not a GPU kernel.

Efficiency

• Cost: 0.7 ms per step at V = 32,000 naive, 44 µs with top-k first (measured). Bound: memory passes over V logits plus the sort.
• Levers:
  • top-k before sort and softmax — 678 → 43 µs at V = 32,000, 3,340 → 153 µs at 128,000 (measured) — k ≪ V
  • min-p instead of top-p — one max, no sort (Hugging Face source) — sort-bound loops
  • count table for penalties — O(H) updates (llama.cpp token_count) — long contexts
• Naive → standard → state of the art:
  • full sort, softmax, cumulative scan · O(V log V), 680 µs at V = 32,000 (measured)
  • top-k selection first · O(V + k log k), 44 µs (measured)
  • Gumbel-max on the device: argmax(l/T + g), g = −ln(−ln u), exact draw in one pass, no softmax · O(V) (candle has GumbelSoftmax; not measured)

Implementation notes

• Subtract the max before exp; keep at least one token; one generator per request, never shared by a batch.

Pitfalls

• Same numbers, different order, different distribution across Hugging Face and llama.cpp.
• Dividing a negative logit by θ > 1 raises its probability; multiply (llama.cpp source comment).
• Penalty scope differs: Hugging Face's repetition penalty covers the prompt; vLLM's presence/frequency cover output only.
• A seed is not reproducibility. Kernels change reduction order with batch size, so logits differ in the last bits and one flipped token changes everything after it. He et al.: 1,000 temperature-0 completions from Qwen3-235B-A22B-Instruct-2507 gave 80 distinct outputs; batch-invariant kernels made all identical. vLLM requires batch invariance and the same hardware and version; OpenAI calls seed best effort.
• Min-p's gain is disputed: Schaeffer et al. found none at matched hyperparameter counts.

Variants and alternatives

In production

• Hugging Face generate(), vLLM SamplingParams (parameters differ per request inside one continuous batch), llama.cpp sampler chains (std::mt19937 behind dist(seed)).
• OpenAI API: frequency_penalty, presence_penalty (-2..2), seed.

Implementations

Sources

• Holtzman et al., The Curious Case of Neural Text Degeneration, ICLR 2020. paper
• Fan et al., Hierarchical Neural Story Generation, ACL 2018 (top-k sampling). paper
• Keskar et al., CTRL, 2019, section 4.1. paper
• Nguyen et al., Turning Up the Heat: Min-p Sampling for Creative and Coherent LLM Outputs, 2024. paper
• Schaeffer et al., Min-p, Max Exaggeration, 2025. paper
• Meister et al., Locally Typical Sampling, 2022. paper
• He et al. (Thinking Machines Lab), Defeating Nondeterminism in LLM Inference, 2025. blog
• Kirkpatrick et al., Optimization by Simulated Annealing, Science 1983. paper
• Sutton, Barto, Reinforcement Learning: An Introduction, 1st ed., section 2.3. book
• Hinton et al., Distilling the Knowledge in a Neural Network, 2015. paper
• Code and docs read: transformers · llama.cpp · vLLM params · vLLM reproducibility · OpenAI spec