How LLM Inference Works

You type a prompt, hit enter, and a moment later words start showing up.

What happens in between is simpler than you'd think. The model doesn't write a sentence. It guesses one token - a word, or most of one, adds it to the end, and guesses again. Over and over, until it's done.

Here's what that loop looks like.

Every LLM you've used works like this. The rest of this post slows each step down.

Text becomes tokens

Models don't see letters or words. They see tokens. A token is a chunk of text from a fixed list, usually somewhere between 50,000 and 100,000 entries long.

Common words often fit into one token. Rare words get split into smaller pieces. Spaces can be part of the token too, which is why you'll see ·the instead of the.

This demo runs the real tokenizer GPT-4 uses. Type anything you like, up to 20 words.

Try a word that doesn't exist. It doesn't break. It gets chopped into smaller pieces that do exist, which is how a model copes with words it has never seen. Try a typo too, like inthe. You get two tokens, ·in and the.

You'll never write a tokenizer yourself. You'll use one like this:

import tiktoken
 
enc = tiktoken.get_encoding("cl100k_base")
ids = enc.encode("Large language models predict the next token.")
 
print(len(ids))         # how many tokens you're paying for
print(enc.decode(ids))  # back to the original text

Tokens are what you pay for, and what fills up the context window. So it's worth having a feel for how your text splits.

Through the model

Each token id gets swapped for a vector, which is just a long list of numbers. Then those vectors go through a stack of layers. Llama 3 8B has 32 of them. Llama 3.1 405B has 126.

The interesting part of each layer is attention. It's how a token pulls in information from the words before it. A token can only look back, never forward.

Click on it. It looks hardest at mat, because that's what "it" is about. Real models learn weights like these during training. Nobody writes them by hand.

Under the hood, attention is one line of maths:¹

Each token makes three vectors. A query $Q$: what am I looking for? A key $K$: what do I have? And a value $V$: what do I pass on if someone picks me? The query gets compared against every earlier key, softmax turns those scores into weights, and the token gets back a weighted mix of the values.

After the last layer, the model multiplies the final vector by one more big matrix. Out comes a score for every token in the vocabulary. These scores are called logits. With a 200k vocabulary, that's 200,000 numbers, one for every possible next token.

Picking the next token

Logits are just scores. To pick a token, you turn them into probabilities with softmax:

$T$ is the temperature. Turn it down and the top token wins almost every time. Turn it up and everything flattens out, so unlikely tokens get a real shot.

Then two filters cut off the long tail:

• Top-k keeps only the $k$ most likely tokens.
• Top-p keeps the smallest group of tokens whose probabilities add up to at least $p$.²

Here's the part that's easy to miss: sampling isn't picking the biggest number. You're throwing a dart at the probability bar.

Hit Throw 100 and most darts land on Paris, but not all of them, which is why the same prompt can give you different answers. Greedy hands Paris the whole bar, Wild shrinks it to about half, and Top-p 0.5 cuts everything else so only Paris survives.

Here's the same logic in code:

import numpy as np
 
 
def sample(logits, temperature=1.0, top_k=50, top_p=0.95, rng=np.random.default_rng()):
    logits = logits / max(temperature, 1e-5)
 
    if top_k:
        cutoff = np.sort(logits)[-top_k]
        logits = np.where(logits < cutoff, -np.inf, logits)
 
    probs = np.exp(logits - logits.max())
    probs /= probs.sum()
 
    order = np.argsort(probs)[::-1]
    cumulative = np.cumsum(probs[order])
    keep = order[: np.searchsorted(cumulative, top_p) + 1]
 
    mask = np.zeros_like(probs)
    mask[keep] = probs[keep]
    return rng.choice(len(probs), p=mask / mask.sum())

One token at a time

Now you have one token. To get the next one, you add it to the input and run the model again.

def generate(model, tokenizer, prompt, max_new_tokens=64):
    ids = tokenizer.encode(prompt)
    logits, cache = model(ids)
    out = []
    for _ in range(max_new_tokens):
        next_id = sample(logits[-1])
        if next_id == tokenizer.eos_id:
            break
        out.append(next_id)
        logits, cache = model([next_id], cache)
    return tokenizer.decode(out)

The first call, prefill, reads the whole prompt at once. Every call after that is decode: one token in, one token out.

But wait. If you only pass in one token, how does the model still know about the rest? That's what the cache is for.

The KV cache

Remember the keys and values from attention? Every new token needs the keys and values of all the tokens before it. Without a cache, you'd recompute them on every single step. Step 100 would redo the work for 99 tokens you've already seen.

The KV cache keeps them around. Each step works out the keys and values for the one new token and reads the rest from memory.

Turn the cache off and step to the end. The total goes from 6 tokens processed to 18. Seven tokens in, the work has already tripled. With seven thousand, the gap is enormous.

The catch is memory. Every token in the context costs this much:

Read it as one key and one value, for each of $L$ layers and $H_{kv}$ key-value heads of size $d_h$, at $b$ bytes per number.

For Llama 3 8B in fp16, that's $2 \times 32 \times 8 \times 128 \times 2 = 131{,}072$ bytes, or 128 KiB per token. An 8,192-token context needs 1 GiB of cache. For one request.

So long contexts and big batches run out of GPU memory way before they run out of compute.

Prefill and decode

Those two phases behave very differently on a GPU.

During decode, the GPU has to read every weight in the model just to make one token. So for a single request, there's a rough ceiling:

An 8B model in fp16 is about 16 GB. An H100 reads memory at about 3.35 TB/s. So one request tops out at roughly 200 tokens a second, before batching or quantization.

That's why batching, quantization and speculative decoding exist. They all go after the same thing: how many bytes you move per token.