10. Top-k Sampling From Scratch

Problem

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

Given raw logits of shape (batch, vocab_size), keep only the k highest-scoring logits per row, softmax just those into a probability distribution, and sample one token index per row from it.

Source: src/10_top_k_sampling.py

def sample_top_k(logits: torch.Tensor, k: int, generator: torch.Generator | None = None) -> torch.Tensor

>>> import torch
>>> torch.manual_seed(0)
>>> logits = torch.tensor([[2.0, 1.0, 0.1, -1.0, 5.0]])
>>> sample_top_k(logits, k=2, generator=torch.Generator().manual_seed(0))
tensor([4])

Step-by-Step Approach

1. Clamp k to min(k, vocab_size) so a caller-supplied k larger than the vocabulary just means "keep everything."
2. Use torch.topk(logits, k, dim=-1) to get the k highest logits per row and their original vocabulary indices in one call — this already sorts them descending.
3. Softmax only the kept top-k values with torch.softmax. Softmax subtracts the row max internally, so this stays numerically stable even with large logits — no manual shifting needed.
4. Sample a local index (0..k-1) per row from that softmax distribution with torch.multinomial, passing the seeded generator for reproducibility.
5. Map the local index back to a real vocabulary id with topk_idx.gather(-1, sampled_local), then squeeze the trailing dimension to get a flat (batch,) result.

The key insight is that top-k sampling is two decoupled steps: a hard selection step (which tokens survive) and a renormalization-plus-sampling step (turning survivors into a distribution to draw from) — and torch.topk plus gather is the idiomatic way to keep the mapping between "local softmax index" and "real vocabulary id" correct without ever materializing a full (batch, vocab) masked tensor.

Reference solution

import torch


def sample_top_k(logits: torch.Tensor, k: int, generator: torch.Generator | None = None) -> torch.Tensor:
    vocab_size = logits.shape[-1]
    k = min(k, vocab_size)
    # (batch, k): the k highest logits per row, already sorted descending
    topk_vals, topk_idx = torch.topk(logits, k, dim=-1)
    # softmax handles the max-subtraction stability trick internally, so no
    # manual logit shifting is needed even though topk_vals can be large
    probs = torch.softmax(topk_vals, dim=-1)  # (batch, k)
    # sample a *local* index (0..k-1) into the kept candidates per row
    sampled_local = torch.multinomial(probs, num_samples=1, generator=generator)  # (batch, 1)
    # map the local index back to the real vocabulary id via the topk indices
    sampled = topk_idx.gather(-1, sampled_local).squeeze(-1)  # (batch,)
    return sampled

Key Functions & Tricks

• torch.topk(logits, k, dim=-1) — returns both the top-k values and their original indices in one sorted call.
• torch.softmax(topk_vals, dim=-1) — numerically stable by construction (internally subtracts the row max before exponentiating).
• torch.multinomial(probs, num_samples=1, generator=...) — draws a categorical sample per row; the explicit generator argument is what makes sampling reproducible in tests.
• Tensor.gather(-1, index) — the standard trick for mapping a "local" (post-selection) index back to an "original" index without a Python loop.
• Tensor.squeeze(-1) — drops the trailing size-1 dimension left by multinomial's (batch, 1) output.

How to Recognize This Pattern