Top-k Sampling - LLM Generation & Decoding

Top-k Sampling

What you'll build

Sample the next token from only the k most likely tokens: apply temperature, keep the top k logits, renormalize, and draw one token. Return the token and the filtered distribution.

Background

Pure sampling from the full softmax occasionally picks a token from the long tail of thousands of unlikely tokens, and one bad token can derail the rest of the text. Top-k (Fan et al., 2018; GPT-2 used k = 40) truncates the tail. Its weakness is that k is fixed: when the model is confident, k = 40 still admits 39 poor tokens; when it is uncertain, 40 may cut off good ones. Top-p (next problem) adapts the cutoff instead.

Algorithm

1. z = logits / temperature.
2. Find the indices of the k largest values of z.
3. Set every other logit to -inf.
4. Softmax (stable) over the result: the kept tokens renormalize to sum to 1, the rest get exactly 0.
5. Sample one token index from that distribution.

Inputs

• logits: float array (vocab_size,)
• k: int, 1 <= k <= vocab_size
• temperature: float > 0, default 1.0 (applied before filtering)

Expected Return Format

A tuple (token, filtered_probs):

• token: Python int, an index whose filtered probability is > 0
• filtered_probs: float array (vocab_size,) with exactly k non-zero entries summing to 1

Edge cases and conventions

• Keep the np.random.seed(42) line from the starter and sample with np.random.choice(len(probs), p=probs), so runs are reproducible.
• k = 1 is greedy decoding: the argmax gets probability 1.
• Tests avoid ties at the k-th largest logit, so any tie-breaking is fine.

Constraints

NumPy only. np.argsort or np.argpartition for the top-k.

Follow-up

• np.argpartition finds the top-k in O(V) instead of O(V log V). Does that matter for a 128k vocabulary on a GPU?
• How do top-k and temperature interact if you apply temperature after filtering instead?

Input: top_k_sampling(np.array([5.0, 3.0, 2.0, 1.0, 0.1]), k=3)
Output: (token in {0, 1, 2}, filtered_probs = [0.8438, 0.1142, 0.042, 0.0, 0.0])

Input: top_k_sampling(np.array([2.0, 0.0, 4.0, 1.0, -1.0]), k=3, temperature=2.0)
Output: filtered_probs = [0.2312, 0.0, 0.6285, 0.1402, 0.0]

import numpy as np

def top_k_sampling(logits, k, temperature=1.0):
    """
    Sample from the top-k most likely tokens.

    Args:
        logits: Raw logits (vocab_size,)
        k: Number of top tokens to keep
        temperature: Temperature, applied before filtering

    Returns:
        token: Sampled token index (int)
        filtered_probs: (vocab_size,), k non-zero entries summing to 1
    """
    np.random.seed(42)
    # Your code here
    pass