Sampling Strategies

The model proposes. The sampler decides.

In every decode step → Ch.06, the model produces one number for every word in its vocabulary — about 50,000 numbers, called logits. These are raw scores, not a decision. The job of the sampler is to turn those 50,000 scores into a single chosen token. How it makes that choice determines the entire character of the output.

This chapter walks through each "judging style" — what it does mathematically, when to use it, and how it appears in nano-vLLM's sampler. By the end, you'll understand exactly what happens when you set temperature=0.7 in an API call.

Logits and softmax — the foundation

What is a logit?

A logit is a raw, unbounded score the model assigns to each possible next token. It can be any number — positive, negative, large, small. A logit of 9.7 means the model thinks this token is very likely; a logit of -3.2 means it thinks it's very unlikely. But logits are not probabilities — they don't add up to 100%, and a logit of 9.7 doesn't directly tell you "97% likely". They're just relative scores.

What is softmax?

To make a decision, we need to convert those raw logits into proper probabilities — numbers between 0 and 1 that sum to exactly 100%. The function that does this is called softmax. It takes the whole list of logits and produces a probability for each, where higher logits get higher probabilities, and everything adds up to 1.

Once we have probabilities, the sampler can make its choice. The different sampling strategies are all about how to use these probabilities — and how to reshape them before choosing.

Greedy decoding — always pick the top

Greedy decoding is the simplest strategy: always pick the token with the highest probability. No randomness, no chance for second place. In our weather example, greedy always picks "sunny" because it has the highest score (40%). The technical name for "pick the highest" is argmax — the argument (position) of the maximum value.

Temperature — the creativity dial

Temperature controls how much randomness the sampler allows. It works by scaling the logits before softmax: each logit is divided by the temperature value. This simple division has a powerful effect on the shape of the probability distribution.

Reshape the distribution yourself

Adjust temperature, top-k, and top-p below and watch the probability distribution change in real time. Then click "Sample a token" to see which one gets picked — run it multiple times to feel how randomness changes with each setting.

Top-k and top-p — trimming the candidates

Temperature reshapes the entire distribution, but it never fully removes bad options — even at high temperature, a terrible token retains a small chance. Top-k and top-p solve this by cutting off the unlikely tokens entirely before sampling.

Top-k sampling — keep the best K

Top-k keeps only the K highest-probability tokens and discards everything else. With top-k=50, only the 50 most likely tokens can be selected; the other ~49,950 are removed (their probability set to zero). Then sampling happens among those 50. This guarantees the model never picks something wildly unlikely, while still allowing variety among reasonable options.

Top-p (nucleus) sampling — keep the smallest group covering P%

Top-p, also called nucleus sampling, fixes top-k's rigidity. Instead of a fixed count, it keeps the smallest set of tokens whose probabilities add up to at least P. With top-p=0.9, you sort tokens by probability and keep adding them to the shortlist until their combined probability reaches 90% — then stop. The shortlist size adapts automatically to the model's confidence.

The sampler in code

nano-vLLM's sampler lives in sampler.py → Ch.02. It runs on the GPU, taking the raw logits from the model's final layer and producing one token ID per sequence. Here it is, annotated:

def sample(logits: Tensor, params: SamplingParams) -> Tensor:
    # logits shape: [batch_size, vocab_size] — one row per sequence,
    # 50,000 raw scores per row

    # ── GREEDY SHORTCUT: temperature 0 means pure argmax ──
    if params.temperature == 0:
        return logits.argmax(dim=-1)   # pick the single highest-scoring token

    # ── STEP 1: temperature scaling ──
    # Divide all logits by temperature. Smaller temp = sharper distribution.
    logits = logits / params.temperature

    # ── STEP 2: top-k filtering ──
    # Keep only the K highest logits; set the rest to -infinity
    # (-inf becomes 0 after softmax, removing those tokens entirely)
    if params.top_k > 0:
        # Find the K-th largest value as a threshold
        kth_value = torch.topk(logits, params.top_k, dim=-1).values[..., -1, None]
        logits = torch.where(logits < kth_value, float('-inf'), logits)

    # ── STEP 3: convert to probabilities via softmax ──
    probs = torch.softmax(logits, dim=-1)

    # ── STEP 4: top-p (nucleus) filtering ──
    # Sort by probability, keep the smallest set summing to >= top_p
    if params.top_p < 1.0:
        probs = _apply_top_p(probs, params.top_p)

    # ── STEP 5: sample one token from the final distribution ──
    # multinomial picks an index weighted by probability — this is
    # where the actual randomness happens
    return torch.multinomial(probs, num_samples=1).squeeze(-1)

def _apply_top_p(probs: Tensor, top_p: float) -> Tensor:
    # Sort probabilities descending, track the cumulative sum
    sorted_probs, sorted_idx = torch.sort(probs, descending=True, dim=-1)
    cumulative = torch.cumsum(sorted_probs, dim=-1)

    # Mark tokens to remove: those beyond the point where cumulative >= top_p
    remove = cumulative - sorted_probs > top_p   # keep the token that crosses the threshold
    sorted_probs[remove] = 0.0

    # Restore original token order and renormalise so probs sum to 1 again
    probs = torch.zeros_like(probs).scatter(-1, sorted_idx, sorted_probs)
    return probs / probs.sum(dim=-1, keepdim=True)

Choosing the right strategy for the job

Things beginners get wrong about sampling

Quiz

Three questions on sampling. Wrong answers explain exactly where the reasoning went wrong.

What you now know