11. Nucleus (Top-p) Sampling From Scratch

Problem

Top-k sampling uses a fixed candidate count, which breaks down when the model's confidence varies token to token: a peaked distribution wastes diversity on k-1 near-zero-probability tokens, while a flat distribution truncates too aggressively. Nucleus (top-p) sampling instead keeps the smallest set of highest-probability tokens whose cumulative probability reaches a threshold p, so the candidate pool adapts to how confident the model actually is at each step.

Given raw logits of shape (batch, vocab_size), sort each row's probabilities descending, keep the smallest prefix whose cumulative probability is >= p (always keeping at least the top token), zero out and renormalize the rest, then sample one token index per row.

Source: src/11_nucleus_sampling.py

def sample_top_p(logits: torch.Tensor, p: float, 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_p(logits, p=0.9, generator=torch.Generator().manual_seed(0))
tensor([4])

Step-by-Step Approach

1. Sort each row descending with torch.sort(logits, descending=True, dim=-1), keeping both the sorted values and the permutation indices — sorting makes "smallest prefix reaching p" a contiguous run instead of a scattered subset.
2. Softmax the sorted logits to get sorted probabilities, then torch.cumsum them to get the running cumulative probability at each rank.
3. Build the removal mask as (cum_probs - sorted_probs) > p: a token is dropped only if the cumulative mass before it already exceeded p. This guarantees the token that crosses the threshold is kept, and the top-ranked token (whose "cumulative mass before it" is exactly 0) is never dropped.
4. Zero the masked-out probabilities with masked_fill, then renormalize the remainder so it sums back to 1 — this is the "nucleus."
5. Sample a rank via torch.multinomial on the renormalized sorted probabilities, then map that rank back to a real vocabulary id by gathering into the sort's permutation indices.

The key insight is that the removal mask must compare against the cumulative probability excluding the current token, not including it — comparing cum_probs > p directly would sometimes drop the very token that pushes the cumulative sum past the threshold, silently shrinking the nucleus by one token and occasionally collapsing it to nothing when p is very small.

Reference solution

import torch


def sample_top_p(logits: torch.Tensor, p: float, generator: torch.Generator | None = None) -> torch.Tensor:
    # sort descending so the "smallest prefix reaching p" is a contiguous run
    sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)  # both (batch, vocab)
    sorted_probs = torch.softmax(sorted_logits, dim=-1)  # (batch, vocab)
    cum_probs = torch.cumsum(sorted_probs, dim=-1)  # (batch, vocab)
    # a token is dropped only if the cumulative mass *before* it already hit
    # p -- this guarantees the token that crosses the threshold is kept, and
    # the top-ranked token (cum_probs - probs == 0) is always kept
    remove_mask = (cum_probs - sorted_probs) > p
    sorted_probs = sorted_probs.masked_fill(remove_mask, 0.0)
    sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True)  # renormalize the nucleus
    sampled_sorted = torch.multinomial(sorted_probs, num_samples=1, generator=generator)  # (batch, 1)
    # map the sampled *rank* back to the original vocabulary id
    sampled = sorted_idx.gather(-1, sampled_sorted).squeeze(-1)  # (batch,)
    return sampled

Key Functions & Tricks

• torch.sort(logits, descending=True, dim=-1) — returns both sorted values and the permutation needed to map ranks back to original indices.
• torch.cumsum(sorted_probs, dim=-1) — running cumulative probability, the core quantity the nucleus threshold is compared against.
• cum_probs - sorted_probs — recovers "cumulative probability strictly before this token" from a cumulative-inclusive tensor without a second cumsum pass.
• Tensor.masked_fill(mask, 0.0) — zeroes out excluded tokens in place of a Python-level filter.
• Tensor.gather(-1, index) — maps a sampled sort-rank back to the real vocabulary id, mirroring the top-k problem's local-to-global index mapping.

How to Recognize This Pattern