Skip to article frontmatterSkip to article content
Site not loading correctly?

This may be due to an incorrect BASE_URL configuration. See the MyST Documentation for reference.

ch04 · Generating Tokens Correctly

Chapter header
TierTier 1 — CPU (any laptop)
Prerequisitesch01
ScorecardNo performance change. Establishes the correctness baseline every later optimisation is tested against.

The problem

We are about to spend twenty-five chapters making this engine faster. Every one of those changes is an opportunity to make it wrong — and wrong in the particular way that is hardest to catch, where the output is still fluent, still plausible, just no longer what the model would have said.

Before optimising, we need a decode path we own completely and can hold fixed. That means taking back the decisions model.generate() was making on our behalf. There are more of them than you would expect, and one of them is a bug that will otherwise ship to production and look like the model’s fault.

The idea

Sampling is a pipeline, and the order matters

Turning logits into a token is four decisions applied in sequence:

  1. Penalties adjust the logits of tokens already seen.

  2. Temperature rescales the whole distribution: below 1 sharpens it, above 1 flattens it.

  3. Truncation — top-k, top-p, min-p — removes the tail.

  4. Draw one token from what survives.

The order is not arbitrary, and getting it wrong is a classic source of two implementations disagreeing while both look correct. Apply temperature after truncation and you change which tokens survived the cut — high temperature will have widened a distribution that was already narrowed, so the nucleus you kept was computed against different probabilities than the ones you sample from.

A note on the repetition penalty, which looks like a typo when you first meet it:

sampling.py
def apply_repetition_penalty(
    logits: torch.Tensor, prev_ids: torch.Tensor, penalty: float
) -> torch.Tensor:
    """Divide the logit of every already-seen token, or multiply if it is negative.

    The asymmetry looks odd but is the standard formulation: dividing a negative logit would make
    the token *more* likely, which is the opposite of a penalty.
    """
    if penalty == 1.0 or prev_ids.numel() == 0:
        return logits
    seen = logits.gather(-1, prev_ids)
    seen = torch.where(seen > 0, seen / penalty, seen * penalty)
    return logits.scatter(-1, prev_ids, seen)

Dividing a negative logit makes it larger, and therefore the token more likely — the opposite of a penalty. The sign has to be handled explicitly, and implementations that forget it quietly encourage exactly the repetition they were meant to suppress.

The truncation filters are worth distinguishing, because they fail differently:

Stopping is not one condition

A request ends for one of several reasons, and a serving engine has to distinguish them because the caller needs to know which happened:

Only the first is the model finishing a thought. Reporting a budget exhaustion as a normal completion is how truncated JSON reaches a caller that had no reason to suspect it.

Stop strings are harder than stop tokens, because a string may span several tokens, and may be only partially emitted when you check. That is the same boundary problem as the next section, which is the one that actually bites.

The bug: streaming detokenisation

Here is a serving bug that looks exactly like a model bug.

The obvious way to stream is to decode each token to text as it is produced and send it. For ASCII this works perfectly, so it survives testing.

Then someone types a prompt in Japanese. A character like 日 is three bytes; with a byte-level tokenizer that is three tokens, and with a BPE tokenizer it is frequently split too. Decode the first of those bytes on its own and it is not valid UTF-8, so the decoder emits �. The user watches replacement characters appear mid-word and concludes the model is broken.

The fix is to hold back bytes that do not yet form a complete character:

tokenizer.py
class IncrementalDetokenizer:
    """Turns a stream of token ids into a stream of *safe to emit* text fragments.

    The contract: feed tokens one at a time, and receive only text that will never be revised.
    When a token completes a multi-byte character, the whole character is emitted at once; until
    then the partial bytes are held back.

    The subtlety that makes a hand-rolled version wrong is telling a *truncated* sequence apart
    from an *invalid* one. Both raise ``UnicodeDecodeError``, but the first should be held back
    and the second emitted as a replacement character. Python's incremental codec already draws
    that line correctly, so we use it rather than re-deriving the UTF-8 state machine.

    Without this, a streaming endpoint emits replacement characters mid-word for any non-ASCII
    text, which looks like model corruption and is actually a serving bug.
    """

The subtlety — and the reason our first attempt at this was wrong — is telling a truncated sequence from an invalid one. Both raise the same exception. The first should be held back until more bytes arrive; the second should be emitted as a replacement character immediately, because waiting for it to become valid means waiting forever. Python’s incremental codec already draws that line correctly, so we use it rather than re-deriving the UTF-8 state machine and getting it subtly wrong.

The build

Everything above lives in two modules. The engine now owns its decode path end to end:

sampling.py
def sample(
    logits: torch.Tensor,
    params: SamplingParams,
    *,
    prev_ids: torch.Tensor | None = None,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    """Pick one token per row of ``logits`` ([batch, vocab]).

    Order matters and is the conventional one: penalties act on raw logits, temperature rescales,
    then the truncation filters narrow the field. Applying temperature after truncation would
    change which tokens survive, which is a subtle way for two implementations to disagree while
    both looking correct.
    """

And the stopping logic is explicit about why a request ended:

request.py
    def check_finished(self) -> bool:
        """Stop on token budget or a stop token. Returns True if this call ended the request."""
        params = self.request.params
        if self.finished:
            return False
        if self.output_token_ids and self.output_token_ids[-1] in params.stop_token_ids:
            self.finished, self.finish_reason = True, "stop_token"
            return True

The measurement

This chapter’s measurement is not a latency number — it is a test suite. From here on, every optimisation must produce identical output to the path it replaces, and that claim is checked rather than asserted.

The pattern that matters most:

test_engines.py
def test_kv_cache_does_not_change_output(model):
    """Chapter 5's correctness claim: the cache is an optimisation, not a different model."""
    assert drain(NaiveEngine(model)) == drain(CachedEngine(model))

Two requirements this places on the engine, both easy to lose later:

The streaming contract gets its own tests, because the failure is silent:

test_tokenizer.py
def test_partial_multibyte_is_held_back_not_emitted(tok):
    """The heart of it: no output until a character is complete."""
    ids = tok.encode("日")  # three bytes
    det = IncrementalDetokenizer()
    assert det.append(ids[0]) == ""
    assert det.append(ids[1]) == ""
    assert det.append(ids[2]) == "日"

The cost

We now own code that HuggingFace used to own, and that has a price:

Key takeaways

Looking ahead

The decode path is correct and ours. It is also doing an enormous amount of redundant work: ch01 · What an Inference Server Actually Does’s engine re-runs the entire prefix for every single token. ch05 · The KV Cache caches what it already computed, measures the difference, and immediately runs into the constraint that shapes all of Part III.

Further reading

Nucleus sampling was introduced in Holtzman et al., The Curious Case of Neural Text Degeneration, which is also the clearest explanation of why pure greedy decoding produces degenerate repetition. For stop-string handling and chat-template edge cases, ch24 · The API Surface [DRAFT] returns to the topic once there is an API surface to break.

Part I — The Serving Problem
ch03 The Arithmetic of Inference
Part II — The Decode Loop
ch05 The KV Cache