How a language model turns last-step logits into the next token. The model itself only produces a score per vocabulary id. Sampling decides which id is actually emitted.

Input: logits with shape [batch, vocab_size].
Output: next token ids with shape [batch, 1].

Typical knobs: temperature, top_k, top_p. Common defaults for greedy decoding: temperature = 0, top_k = 0 (off), top_p = 1.0 (off).

Pipeline

Each step below is optional except the first branch and the final sample. Greedy (temperature == 0) skips the rest.

logits
  → temperature == 0?  argmax  → done
  → logits / temperature
  → top-k mask  (optional)
  → top-p / nucleus mask  (optional)
  → softmax
  → multinomial sample 1 token

1. Greedy (temperature <= 0)

Pick the highest logit. Ranking is deterministic; no softmax, no randomness. z_i is the logit for token i, and V is the vocabulary.

$$ \hat{y} = \arg\max_{i \in V} z_i $$
return torch.argmax(logits, dim=-1, keepdim=True)

2. Temperature (logits = logits / temperature)

Write scaled logits as z_i / T with T > 0. Then apply softmax. Textbook softmax is the T=1 case:

$$ \mathrm{softmax}(z)_i = \frac{\exp(z_i)}{\sum_{j \in V} \exp(z_j)} $$

Temperature is the same function on z/T:

$$ p_i = \mathrm{softmax}(z / T)_i = \frac{\exp(z_i / T)}{\sum_{j \in V} \exp(z_j / T)} $$

Grey bars = softmax(z) (T=1, before). Colored bars = softmax(z/T) (after). Four logits z=(10,8,6,3). Ranking stays A > B > C > D; only the heights change. T>1 (right) flattens the distribution: tokens that were low-probability become more likely to be sampled. T<1 (left) does the opposite — mass piles on the top token.

Bar chart: colder T piles mass on A; hotter T lifts B/C/D

Grey dashed line = input logits z (same for every T). Colored curves = softmax(z/T). T only changes how peaked the distribution is: smaller T → sharper peak on the top token; larger T → flatter, so other tokens get more probability. Argmax stays index 10.

softmax(z/T) over a logit bump: colder T is a sharp peak, hotter T is flatter; argmax stays index 10

When T > 0 the sampler does not take argmax; it divides logits by T and later samples from softmax. Smaller T makes the distribution peakier (almost always the top token). Larger T makes it flatter (lower-ranked tokens get sampled more often). That is why different T can yield different text.

logits = logits / temperature
# softmax runs after optional top-k / top-p

3. Top-k (top_k > 0)

Temperature already ran. Top-k cuts by rank: keep the k largest logits, send the rest to -inf. z_(k) is the k-th largest value (a threshold, not the top-1). Anything strictly below it is dropped. Ties at the cutoff can keep more than k tokens.

$$ \tilde{z}_i = \begin{cases} z_i / T & \text{if } z_i \ge z_{(k)} \\ -\infty & \text{otherwise} \end{cases} $$

Then exp(z̃_i) = 0 for dropped tokens, so they get probability 0 after softmax. Survivors are renormalized to sum to 1. top_k=0 disables this step.

kth = torch.topk(logits, k, dim=-1).values[..., -1:]
logits = logits.masked_fill(logits < kth, -inf)

4. Top-p (0 < top_p < 1)

Top-p cuts by probability mass, not by count. Sort remaining tokens by probability descending: p_(1) ≥ p_(2) ≥ ⋯. Walk down the list; keep the prefix whose running sum is still ≤ p; drop the rest. How many tokens survive is dynamic: a peaky softmax fills p with few ids; a flat softmax needs many.

$$ V_p = \bigl\{ (1),\ldots,(m) \bigm| m = \max\{ m' : \textstyle\sum_{i=1}^{m'} p_{(i)} \le p \} \bigr\} $$

The implementation always keeps the top-1 token (mask[..., 0] = False), then sets z̃_(i) = -∞ for i \notin V_p. top_p=1.0 disables this step (V_p = V). Top-k and top-p can both be on: k first, then p.

sorted_logits, sorted_idx = torch.sort(logits, descending=True)
probs = softmax(sorted_logits)
cum = cumsum(probs)
mask = cum > top_p          # drop the tail after cumulative mass exceeds p
mask[..., 0] = False        # always keep the top-1 token
# scatter masked scores back to vocab order

5. Sample

After masking, renormalize with softmax and draw one token. This is not argmax. Greedy is only the T ≤ 0 branch.

$$ p_i = \frac{\exp(\tilde{z}_i)}{\sum_{j \in V} \exp(\tilde{z}_j)}, \qquad \hat{y} \sim \mathrm{Categorical}(p) $$
probs = softmax(logits)
return torch.multinomial(probs, num_samples=1)

Full sample funciton

A complete PyTorch sampler matching the pipeline above. Any causal LM can call this once per generated token.

def sample(
    logits: torch.Tensor,
    *,
    temperature: float,
    top_k: int,
    top_p: float,
) -> torch.Tensor:
    if temperature <= 0:
        return torch.argmax(logits, dim=-1, keepdim=True)
    logits = logits / temperature
    if top_k > 0:
        kth = torch.topk(logits, min(top_k, logits.size(-1)), dim=-1).values[..., -1:]
        logits = logits.masked_fill(logits < kth, torch.finfo(logits.dtype).min)
    if 0.0 < top_p < 1.0:
        sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)
        probs = torch.softmax(sorted_logits.float(), dim=-1)
        cum = torch.cumsum(probs, dim=-1)
        mask = cum > top_p
        mask[..., 0] = False
        sorted_logits = sorted_logits.masked_fill(mask, torch.finfo(logits.dtype).min)
        logits = torch.zeros_like(logits).scatter_(-1, sorted_idx, sorted_logits)
    probs = torch.softmax(logits.float(), dim=-1)
    return torch.multinomial(probs, num_samples=1)