Sampling

The decode loop from Chapter 1 has one step still unexplained. Every iteration it runs the model, reads the logits at the final position — one score for each of the 151,936 tokens in the vocabulary — and has to turn that row of scores into a single next token. Choosing it is the sampler, the last closed box in the loop, and the only part of generation allowed to be random.

There are three ways to choose, and forge has all of them: take the highest-scoring token every time, reshape the scores with a temperature and draw from them, or restrict the draw to the top of the distribution. This chapter reads each one from forge’s real code and then proves the one most likely to hide a mistake.

8.1 Greedy, temperature, and the fp32 softmax

Both choices live in one branch of the decode loop.

        if temperature:
            # softmax in fp32: over a 151,936-wide vocabulary bf16 lifts the tail
            probas = torch.softmax(logits.float() / temperature, dim=-1)
            probas = top_p_filter(probas, top_p)
            # multinomial has diverged and crashed on accelerator backends for
            # large draws; the round trip through cpu is cheap
            nxt = torch.multinomial(probas.cpu(), num_samples=1, generator=generator).to(device)
        else:
            nxt = torch.argmax(logits, dim=-1, keepdim=True)

Listing 8.1 · the two ways forge turns a row of logits into a token: sample under a temperature, or take the argmax.

The else branch is greedy decoding: torch.argmax returns the single highest-scoring token. It is deterministic — the same logits give the same token every run — and it is what temperature = 0 selects. That is all greedy means: no randomness, always the top of the list.

Everything above the else runs when a temperature is set. The logits are divided by the temperature and turned into probabilities by a softmax, then torch.multinomial draws one token from that distribution. The randomness enters at that draw and nowhere else, seeded by a generator so a run can be reproduced exactly. Temperature reshapes the distribution before the draw: below one it sharpens toward the top token — as it approaches zero the top token’s probability approaches one and sampling collapses back into greedy — at one it leaves the trained distribution untouched, and above one it flattens, handing low-probability tokens a real chance to be picked.

One detail in that softmax is not incidental: logits.float(). The model runs in half precision, but over a vocabulary 151,936 tokens wide the exponentials in a bf16 softmax underflow the small probabilities in the tail to zero — the very tokens temperature and top-p exist to reach. Lifting the row to fp32 for the softmax keeps that tail intact. The tiny model in this chapter’s driver already runs in fp32, so the cast is inert there; it earns its place on the real 0.6B model’s bf16 weights, the same fp32-guards-a-wide-reduction pattern RMSNorm used in Chapter 3.

8.2 The nucleus

Temperature makes the whole tail reachable, and most of that tail is nonsense — thousands of tokens each holding a sliver of probability that, summed, can outweigh the handful of good ones. Top-p, or nucleus, sampling cuts it: keep only the smallest set of tokens whose probabilities add up to p, throw the rest away, and sample within what remains.

def top_p_filter(probas: torch.Tensor, top_p: float | None) -> torch.Tensor:
    if top_p is None or top_p >= 1.0:
        return probas
    sorted_p, idx = torch.sort(probas, dim=-1, descending=True)
    # compare the exclusive prefix, not the inclusive sum, so the boundary token stays
    prefix = torch.cumsum(sorted_p, dim=-1) - sorted_p
    keep = prefix < top_p
    keep[..., 0] = True
    kept = torch.where(keep, sorted_p, torch.zeros_like(sorted_p))
    out = torch.zeros_like(probas).scatter(-1, idx, kept)
    return out / out.sum(dim=-1, keepdim=True).clamp_min(1e-12)

Listing 8.2 · forge’s nucleus filter. Sort, keep the running total under p, renormalize.

Sort the probabilities from high to low, walk down the sorted list accumulating a running total, and keep tokens while that total is still under p. The subtle line is the prefix. It sums the probabilities that come before each token, not including the token itself — torch.cumsum(sorted_p, dim=-1) - sorted_p — so the token whose own mass carries the total across p is itself kept, not dropped. Comparing the inclusive sum instead would cut that boundary token, and the nucleus would come up one token short on every row. keep[..., 0] = True covers the case where the single most likely token already exceeds p on its own: the top token is always kept, so the model can never be left with nothing to draw from. The final division renormalizes the survivors back to a proper distribution that sums to one.

8.3 Correct, to the bit

Greedy needs no defense — it is argmax, and there is nothing in it to get wrong. Temperature and top-p are where mistakes hide, and top-p is the one with a real function and a boundary that is easy to miscount. So the check targets it: forge’s top_p_filter against an independent implementation of the same rule, run on a batch of sixty-four seeded logit rows, and a second check that a fixed generator, drawing from each filtered distribution, lands on the same token. Sampling is random, so the driver fixes every seed and verifies only what is deterministic — the filtered distributions themselves, and the tokens a seeded draw takes from them.

The nucleus filters agree exactly — 0.00e+00 — because both apply the identical rule to the identical probabilities: sort, exclusive prefix, keep, renormalize. And all sixty-four seeded draws land on the same token, so the randomness itself is reproducible: same seed, same distribution, same choice. A correct sampler is not one that always picks the same word — it is one whose randomness you can pin down and replay.

The decode loop from Chapter 1 now runs on nothing but code from this Part. The forward pass builds the logits, block by block, in Chapters 2 through 6; the cache makes every step after the first cheap in Chapter 7; and the sampler turns the final row into a token here. Read that loop again and not one call in it is a black box. What Part I still adds is not a new step in the loop but the scaffolding around it — loading a real released checkpoint into the model, the tokenizer that turns text into ids and back, and the roofline arithmetic that predicts how fast the whole thing can run before a single kernel is written. Then Part II moves from one model on one machine to serving many.