How LLMs work

Concept#

An LLM is a next-token predictor. Feed it a sequence of tokens and it outputs a probability distribution over its vocabulary for “what comes next”. You pick one token (sampling, covered in Sampling parameters and the model landscape), append it to the sequence, and repeat. That loop, called autoregressive generationAutoregressive generationProducing text one token at a time, where each new token is appended to the input before the next step. This is why output length drives latency.Open the glossary entryClose, is the whole inference process. Chat formatting, system prompts and tool calls are just structure fed into the same loop.

The engine inside the loop is the transformerTransformerThe neural network design behind modern LLMs. Stacked blocks of self-attention and feedforward layers process all positions of the input in parallel during training and prefill.Open the glossary entryClose, built on self-attentionSelf-attentionA step where each token looks at the visible tokens, including itself, and takes a weighted mix of their information. It lets the model link distant words and runs in parallel across positions.Open the glossary entryClose. Every token computes a Query (“what am I looking for”), a Key (“what do I contain”) and a Value (“what do I offer”). A token’s Query is compared against the Keys of the tokens it is allowed to see, which produces attention weights. Those weights blend the Values. All tokens do this in parallel, in a few large matrix multiplications.

Why It Matters#

RNNs and LSTMs read a sequence one token at a time through a hidden state. That has two costs. Training cannot be spread across a sequence on a GPU, because step 5 needs the result of step 4. And information from far back has to survive being squeezed through a fixed-size state again and again, so long-range links degrade.

Self-attention fixes both. During training every position is processed at once, and a token can read directly from any earlier token in one step, however far back it sits. This parallelism is a large part of why training at current scale became practical.

You will need this mechanism for almost every later cost and latency question in the course, so it is worth getting it right now.

Mental Model#

Picture a room of experts (the earlier tokens), each wearing a sign that says what they know (Key). You hold a note describing what you need (Query). You compare your note against every sign, then take a weighted blend of what the relevant experts say (Value). Every token does this at the same time.

One rule changes the picture for text generation: you may only ask the experts who spoke before you. Nobody can read ahead to tokens that do not exist yet. That rule is the causal maskCausal maskA rule in attention that blocks each token from looking at later tokens. Without it the model could copy the answer it is supposed to predict.Open the glossary entryClose, and it is what makes next-token prediction possible.

Example#

“The trophy didn’t fit in the suitcase because it was too big.” What does “it” refer to? Self-attention lets “it” connect directly to “trophy” regardless of distance. An RNN has to carry that link through every intermediate hidden state and often loses it.

Technical Deep Dive#

• Multi-head attention: several attention computations run in parallel (“heads”), each free to learn a different relationship (syntax, coreference, position). Their outputs are concatenated.
• Positional encoding: attention has no built-in notion of order, so position information is injected explicitly (learned, sinusoidal or rotary schemes). Without it, “dog bites man” and “man bites dog” look the same to attention.
• Causal masking: in a decoder-only LLM, a token may attend only to itself and the tokens before it. During training the mask lets the model learn next-token prediction at every position in parallel without peeking at the answer. During generation there is simply nothing later to look at.
• Residual connections and layer norm: keep training stable across dozens of stacked transformer blocks.

Code#

def attention(Q, K, V, causal_mask):
    d_k = Q.shape[-1]
    scores = (Q @ K.transpose(-2, -1)) / (d_k ** 0.5)   # (n x n)
    scores = scores.masked_fill(~causal_mask, float("-inf"))  # no peeking ahead
    weights = softmax(scores, axis=-1)
    return weights @ V

scores is an n x n matrix, where n is the sequence length. Doubling the prompt quadruples the number of scores. Dividing by the square root of d_k keeps the dot products from growing with the vector size, which would push softmax into near one-hot outputs and starve the training signal. The mask sets “future” positions to minus infinity so softmax gives them zero weight.

Prefill and decode#

Generation has two phases with different performance profiles.

• PrefillPrefillThe first phase of a request, in which the whole prompt is processed in parallel. It sets time to first token and is usually limited by compute.Open the glossary entryClose: the model processes your whole prompt in one pass. All prompt tokens are known, so they are handled in parallel as large matrix multiplications, which keeps the GPU busy. Prefill sets the time to first token.
• Decode: the model then generates one token per step. Each step has little new arithmetic to do, but it must read the model weights and the stored history from memory. Decode is typically limited by memory bandwidth, not compute (NVIDIA, 2023-11-17).

Mental model: prefill is reading the whole question once. Decode is writing the answer one word at a time while glancing back at your notes before each word.

Those notes are the KV-cacheKV cacheStored key and value vectors of earlier tokens, reused at each decode step so they are not recomputed. It saves compute but its memory grows with context length and concurrent requests. It is not a response cache.Open the glossary entryClose. Each past token’s Keys and Values never change once computed, so the server stores them and, at each new step, computes K and V only for the new token. Without the cache, every step would redo the projections for the entire history.

The cache is a trade: you save compute and pay in memory. Per token, it stores a Key and a Value for every layer:

KV-cache bytes = 2 x layers x KV-heads x head-dim x sequence length x batch x bytes per value

The NVIDIA article gives the same idea with hidden size in place of KV-heads times head-dim. Worked example, illustrative numbers, not a specific model: 32 layers, 8 KV-heads, head-dim 128, 16-bit values (2 bytes). That is 2 x 32 x 8 x 128 x 2 = 131,072 bytes (128 KiB) per token. An 8,192-token context is 1 GiB for one request. Ten concurrent requests at that length hold 10 GiB of cache before the model weights are counted.

Two consequences follow. Long contexts and many concurrent users compete for the same GPU memory, so the cache often caps how many requests a server can run at once. And techniques such as grouped-query attention (fewer KV-heads) exist mainly to shrink this cache. The KV-cache is not a response cache: it never returns a stored answer. It reuses stored Keys and Values so they are not recomputed, within one generation and, with prefix or prompt caching, across requests that share the same prompt prefix. We return to it, with serving costs, in Cost and latency optimization.

Diagrams#

Real-World Application#

Quadratic attention explains why “just use a huge context window” is not free. Doubling a prompt roughly quadruples the attention-score arithmetic during prefill. Total cost grows less steeply than that at short lengths, because the feedforward and projection layers scale only linearly with length. Memory is the second bill: the KV-cache grows linearly with context and with every concurrent request. When your API cost or latency jumps after users start sending longer prompts, look at both prefill time (time to first token) and cache memory (how many requests a server can hold).

Interview Perspective#

Expect: “explain attention”, “why not RNNs”, “what happens during inference”, “what is a KV-cache and why does it matter”. A strong answer names the loop (logits, softmax, sample, append), says what Q, K and V do, and separates prefill from decode. The weak answer is “it just predicts words” with no mechanism, which reads as memorized trivia. A common follow-up is “how would you make attention cheaper for long context?” Sparse, sliding-window and linear attention are the usual families, and grouped-query attention and FlashAttention are the usual serving answers (later lessons touch cost optimization).

Provenance: Common industry pattern. These are standard topics for ML and GenAI roles, not a verified list for any named company.

Practice Questions#

• Basic: What does a transformer output at each generation step?
• Intermediate: Why can’t attention alone tell “dog bites man” from “man bites dog”?
• Advanced: Why does the KV-cache not help with the first token of a response, and what does it help with?
• Product-Company: Your API cost tripled after users started sending longer prompts. Explain why in terms of attention cost and KV-cache memory, not just “more tokens = more cost”.

Activity#

Trace by hand what happens to the sentence “The cat sat because it was tired”. Identify which earlier token “it” should attend to most strongly and why. Then write down which tokens the word “cat” is allowed to attend to under a causal mask, and what would go wrong in training if it could see “tired”.

Homework#

Pick two real models you might use in production (for example a small open-weight model and a frontier API model). Find each model’s context window and, for the open model, its layer count, KV-head count and head dimension from its published config (stretch task; look for fields such as num_hidden_layers, num_key_value_heads and head_dim or hidden_size divided by num_attention_heads in the model’s config.json). Use the formula above to estimate KV-cache bytes per token and per full-length request. Write two sentences on what that means for how many users one GPU can serve. If a number is not published, say so instead of guessing.

Quick Revision#

• LLM = autoregressive next-token predictor; transformer = the architecture; self-attention = the mechanism.
• Self-attention: Q/K/V per token, parallel, causal mask, n x n scores so cost is quadratic in sequence length.
• Beats RNNs on parallel training and long-range links.
• Prefill is parallel and sets time to first token; decode is one token per step and is usually memory-bandwidth-bound.
• KV-cache saves recomputation and costs memory that grows with context and concurrency.

Check yourself#

• Walk-through: Without notes, describe the path from a prompt to the fifth generated token, and say what the KV-cache stores at that point.
• Variation: The prompt becomes four times longer. What grows (prefill compute, KV-cache memory, time to first token), and what mitigations exist (grouped-query attention, FlashAttention, prompt caching)?
• Common misconception: “It predicts the next word” with no distribution or decoding step; confusing the KV-cache with a response cache.

Related#

• Tokenization and embeddings
• Cost and latency optimization