Problem 104 · Tokenization & Sampling

Top-k, Top-p and Min-p Sampling with a Seeded RNG

HardAI engineering

01Problem

Implement `filter_probs(logits, top_k=0, top_p=1.0, min_p=0.0)` that converts logits to probabilities and keeps a token only if it survives all active filters: in the `top_k` most likely (0 disables), in the smallest prefix of the sorted distribution whose cumulative mass reaches `top_p`, and with probability at least `min_p` times the top probability. Renormalise the survivors. Then `sample(logits, rng, temperature, ...)` applies temperature first, filters, and draws with the given `random.Random`; temperature 0 is greedy.

02Key insight

Order of operations is part of the spec: temperature reshapes the distribution before truncation, so a high temperature widens the nucleus. The three filters compose by intersection, and the top-1 token must always survive or a strict setting produces an empty distribution. min-p scales with model confidence (tight when the model is sure, loose when it is not), which is why it holds up better than a fixed top-p at high temperature. Taking an explicit seeded RNG instead of the global one makes sampling reproducible per request, which is what you need to debug a bad generation.

03Solution

solution.pypython46 lines

import math
import random


def _softmax(logits: list[float]) -> list[float]:
    m = max(logits)
    exps = [math.exp(x - m) for x in logits]
    total = sum(exps)
    return [e / total for e in exps]


def filter_probs(logits: list[float], top_k: int = 0, top_p: float = 1.0, min_p: float = 0.0) -> list[float]:
    probs = _softmax(logits)
    order = sorted(range(len(probs)), key=lambda i: (-probs[i], i))
    keep = set(order[:top_k]) if top_k > 0 else set(order)
    if top_p < 1.0:
        nucleus: set[int] = set()
        cum = 0.0
        for i in order:
            nucleus.add(i)
            cum += probs[i]
            if cum >= top_p - 1e-12:
                break
        keep &= nucleus
    if min_p > 0.0:
        threshold = min_p * probs[order[0]]
        keep = {i for i in keep if probs[i] >= threshold}
    total = sum(probs[i] for i in keep)
    return [probs[i] / total if i in keep else 0.0 for i in range(len(probs))]


def sample(logits: list[float], rng: random.Random, temperature: float = 1.0,
           top_k: int = 0, top_p: float = 1.0, min_p: float = 0.0) -> int:
    if temperature == 0:
        return max(range(len(logits)), key=lambda i: (logits[i], -i))
    probs = filter_probs([x / temperature for x in logits], top_k, top_p, min_p)
    r = rng.random()
    cum = 0.0
    last = 0
    for i, p in enumerate(probs):
        if p > 0:
            cum += p
            last = i
            if r < cum:
                return i
    return last

Try it first.

TimeO(V log V) per step (sort); O(V + k log k) with a partial selection

SpaceO(V)

tests.py · 10 passing

tests.pypython19 lines

L = [math.log(p) for p in [0.5, 0.3, 0.15, 0.05]]
close = lambda a, b: all(abs(x - y) < 1e-9 for x, y in zip(a, b))
assert close(filter_probs(L, top_k=2), [0.625, 0.375, 0, 0])
assert close(filter_probs(L, top_p=0.8), [0.625, 0.375, 0, 0])
assert close(filter_probs(L, top_p=0.81), [0.5 / 0.95, 0.3 / 0.95, 0.15 / 0.95, 0])
assert close(filter_probs(L, min_p=0.2), [0.5 / 0.95, 0.3 / 0.95, 0.15 / 0.95, 0])
assert close(filter_probs(L, top_k=3, top_p=0.5), [1, 0, 0, 0])
assert close(filter_probs(L, top_p=0.0), [1, 0, 0, 0])
assert sample(L, random.Random(0), temperature=0) == 0
a = [sample(L, random.Random(42), top_p=0.9) for _ in range(5)]
b = [sample(L, random.Random(42), top_p=0.9) for _ in range(5)]
assert a == b
rng = random.Random(7)
n = 20000
counts = [0, 0, 0, 0]
for _ in range(n):
    counts[sample(L, rng, top_k=2)] += 1
assert counts[2] == counts[3] == 0
assert abs(counts[0] / n - 0.625) < 0.02

04Staff-level follow-up

On GPU you avoid the full sort: top-k uses a partial select, and top-p can be done by sorting only the top few hundred candidates since the tail rarely matters. Also discuss what breaks reproducibility even with a seed: batch composition changing kernel reduction order, speculative decoding, and different hardware. Structured output constraints are another mask applied before these filters.