14. Safe Top-p Nucleus Sampling with a Hard Token Blocklist

Problem

A production inference server samples the next token from a language model using top-p (nucleus) sampling: keep the smallest set of highest-probability tokens whose cumulative probability reaches p, renormalize, and sample from just that set.

Layered on top of that, a safety system maintains a hard blocklist of token ids that must never be emitted — e.g. tokens that only ever appear inside disallowed content — regardless of how confident the model is in them. The blocklist has to be enforced before the nucleus is even computed, not filtered out after sampling, or a blocked token could still win the sample by chance.

Source: src/14_safe_nucleus_sampling.py

def safe_nucleus_sample(
    logits: torch.Tensor,
    p: float,
    blocked_token_ids: torch.Tensor,
    generator: torch.Generator,
) -> torch.Tensor: ...

>>> logits = torch.tensor([[0.0, 0.0, 100.0, 0.0, 0.0]])
>>> blocked = torch.tensor([2])
>>> gen = torch.Generator().manual_seed(0)
>>> safe_nucleus_sample(logits, p=0.9, blocked_token_ids=blocked, generator=gen)
tensor([...])  # never index 2, even though it has by far the highest logit

Step-by-Step Approach

1. Clone logits and set every blocked column to -inf first, before any softmax or sorting — this is what guarantees a blocked token can never be sampled, since softmax(-inf) == 0 exactly.
2. Softmax the (now blocklist-masked) logits into a probability distribution.
3. Sort probabilities descending per row, and compute the cumulative sum along that sorted order.
4. Mark a sorted position for removal once the cumulative mass before it already exceeds p — this keeps the first token that crosses the threshold so the nucleus is never accidentally empty.
5. Zero out the removed positions, scatter the filtered probabilities back into original vocabulary order (undoing the sort), then renormalize so the row sums to 1 again.
6. Sample one index per row with torch.multinomial, passing the provided generator for reproducibility.

The key insight is ordering: the blocklist must be applied to the raw logits before softmax, not to the final probabilities or the sampled result afterward — masking late leaves a window where a blocked token still has nonzero probability mass (or, worse, could still be the one sampled) before you get a chance to reject it.

Reference solution

import torch


def safe_nucleus_sample(
    logits: torch.Tensor,
    p: float,
    blocked_token_ids: torch.Tensor,
    generator: torch.Generator,
) -> torch.Tensor:
    logits = logits.clone()
    # enforce the blocklist BEFORE softmax/sorting: -inf guarantees zero
    # probability, so a blocked token can never enter the nucleus or win
    # the sample, even by floating-point fluke
    logits[:, blocked_token_ids] = float("-inf")

    probs = torch.softmax(logits, dim=-1)
    sorted_probs, sorted_idx = torch.sort(probs, dim=-1, descending=True)
    cumulative = torch.cumsum(sorted_probs, dim=-1)

    # drop a sorted token once the cumulative mass *before* it already
    # reached p; this keeps the first token that crosses the threshold,
    # so the nucleus is never empty
    remove_sorted = cumulative - sorted_probs > p
    sorted_probs = sorted_probs.masked_fill(remove_sorted, 0.0)

    # scatter filtered probabilities back to original vocab order, then renormalize
    filtered_probs = torch.zeros_like(probs).scatter(1, sorted_idx, sorted_probs)
    filtered_probs = filtered_probs / filtered_probs.sum(dim=-1, keepdim=True)

    return torch.multinomial(filtered_probs, num_samples=1, generator=generator).squeeze(-1)

Key Functions & Tricks

• logits[:, blocked_token_ids] = float("-inf") — advanced indexing to mask an arbitrary set of columns per row in one op; softmax turns -inf into exactly 0 probability.
• torch.sort(probs, dim=-1, descending=True) — returns both sorted values and the permutation indices needed to scatter back later.
• torch.cumsum(sorted_probs, dim=-1) — running total of probability mass in descending order, the core of the top-p cutoff decision.
• cumulative - sorted_probs > p — the "exclusive cumulative sum" trick: subtracting the current element gives the mass accumulated strictly before it, which is what determines whether it's still needed to reach the threshold.
• torch.zeros_like(probs).scatter(1, sorted_idx, sorted_probs) — undoes the sort, placing filtered probabilities back at their original vocabulary indices.
• torch.multinomial(probs, num_samples=1, generator=generator) — samples according to a probability distribution, with an explicit generator for reproducible tests.

How to Recognize This Pattern