10. Top-K Sampling Filter for Next-Token Generation

Problem

A lab's inference server needs to sample the next token from a language model instead of always taking the argmax (greedy decoding), because greedy decoding produces repetitive, low-diversity text. Top-k sampling restricts the candidate pool to only the k highest-scoring tokens before converting scores into a probability distribution to sample from, which keeps generation focused while still allowing variety among plausible tokens.

Given the raw logits for every token in the vocabulary, keep only the k highest-scoring ones and renormalize them into a probability distribution.

Source: src/10_top_k_sampling_filter.py

def top_k_filter(logits: list[float], k: int) -> list[float]:

>>> top_k_filter([2.0, 1.0, 0.1], 2)
[0.731059, 0.268941, 0.0]
>>> top_k_filter([1.0, 1.0, 1.0], 1)
[1.0, 0.0, 0.0]

Step-by-Step Approach

1. Handle the trivial cases first: an empty logits list or k <= 0 both mean "keep nothing" — return all zeros (or an empty list).
2. Clamp k to min(k, len(logits)) so a caller-supplied k larger than the vocabulary just means "keep everything."
3. Select the indices of the k highest-scoring logits. Sorting by (-logits[i], i) picks the top-k by value and breaks ties by lower index, in one pass.
4. Subtract the maximum kept logit from every kept logit before exponentiating (the standard "max-subtraction" trick) so math.exp never overflows on large logits.
5. Sum the exponentiated scores and divide each by that sum to get a softmax probability, restricted to just the kept indices.
6. Round each probability to 6 decimal places and place it back at its original index in a zero-initialized output list; every non-kept index stays 0.0.

The key insight is that top-k sampling is really two separate steps glued together: a selection step (which tokens survive) and a renormalization step (turning surviving scores into a valid probability distribution) — keeping them conceptually separate makes both the max-subtraction stability trick and the zero-fill for dropped tokens obvious.

Reference solution

import math


def top_k_filter(logits: list[float], k: int) -> list[float]:
    n = len(logits)
    if n == 0 or k <= 0:
        return [0.0] * n
    k = min(k, n)
    # pick the k highest-scoring indices, O(n log k) via a bounded heap-style
    # selection (sorted() here for clarity; heapq.nsmallest would be the
    # O(n log k) production choice) -- ties broken by lower index first
    top_indices = sorted(range(n), key=lambda i: (-logits[i], i))[:k]
    # subtract the max kept logit before exponentiating for numerical stability
    m = max(logits[i] for i in top_indices)
    exp_scores = {i: math.exp(logits[i] - m) for i in top_indices}
    total = sum(exp_scores.values())
    result = [0.0] * n
    for i in top_indices:
        result[i] = round(exp_scores[i] / total, 6)
    return result

Key Functions & Tricks

• sorted(range(n), key=...) — selects the top-k indices and their tie-break order in one call instead of a separate selection then sort.
• key=lambda i: (-logits[i], i) — negating the score sorts descending by value while the index acts as a deterministic tie-breaker.
• max-subtraction before exp() — prevents floating-point overflow when logits are large, a standard numerical-stability trick for softmax.
• math.exp — converts log-scale scores back to a linear scale so they can be summed and normalized into probabilities.
• round(..., 6) — keeps expected test outputs as clean, reproducible literals instead of raw floating-point noise.
• zero-initialized output list — makes "not in the top-k" an explicit, easy-to-check 0.0 rather than a missing key.

How to Recognize This Pattern