Transformers from First PrinciplesPost 9 of 18 · 18 of 18 written
- How to use this guide
- 01The machinery: how a neural network computes and learns
- 02Making deep networks trainable
- 03Turning text into numbers
- 04Attention from first principles
- 05The Transformer block
- 06Position: telling the model where each token is
- 07Architecture families and training objectives
- 08Pretraining at scale
- 09Inference and serving
- 10Efficiency variants: MoE, long context and beyond attention
- 11Post-training: from text predictor to assistant
- 12Adapting models: prompting, fine-tuning, LoRA and distillation
- 13Embeddings, retrieval and RAG
- 14Beyond text: vision, speech and multimodal models
- 15Inside the model: interpretability
- 16Evaluation, failure modes and security
- 17From model to system
- 18Fluency kit
9. Inference and serving
Generating text runs the model once per output token, and that loop has two very different phases. Prefill processes the whole prompt in parallel and is limited by compute; decode produces one token at a time and is limited by memory bandwidth. Almost every serving optimization - the KV cache, GQA, FlashAttention, batching, quantization, speculative decoding - attacks one of these two bottlenecks.
The generation loop
- Tokenize the prompt.
- Run the model and take the logits at the last position.
- Choose the next token (decoding strategies below).
- Append it. Stop on an end-of-sequence token or a length limit; otherwise go back to step 2.
Decoding strategies
| Strategy | How it chooses | Use |
|---|---|---|
| Greedy | always the most likely token | deterministic tasks; can loop or repeat |
| Beam search | keeps the k most likely partial sequences | translation, short constrained outputs; bland for open text |
| Temperature sampling | samples from softmax(z / τ) | τ < 1 sharper, τ > 1 more random |
| Top-k | samples only among the k most likely tokens | cuts the long tail of nonsense |
| Top-p (nucleus) | samples among the smallest set with cumulative probability ≥ p | adapts the candidate set to the model’s confidence |
| Min-p | keeps tokens with probability ≥ p_min × the top token’s | robust at higher temperatures |
Repetition and frequency penalties lower the logits of tokens already used. Note that temperature 0 is not always perfectly reproducible in production: batching changes the order of floating-point operations, which changes rounding.
The KV cache
Under the causal mask, the keys and values of past tokens never change when a new token arrives. So cache them: at each decode step, compute q, k and v only for the new token, append k and v to the cache, and attend over the cache. Each step then processes one token instead of the whole sequence.
Its size per token:
The 2 counts K and V. For a 7B-class model with L = 32, 32 KV heads, d_h = 128 and FP16: 2 × 32 × 32 × 128 × 2 bytes = 0.5 MB per token. A 4k-token conversation needs 2 GB; 32 such conversations need 64 GB, far more than the 14 GB of weights. In serving, the KV cache, not the weights, usually decides how many users fit on a GPU.
Prefill vs decode
- Prefill pushes all prompt tokens through the model in one pass. It is large matrix multiplies, compute-bound and efficient on GPUs, and it sets the time to first token (TTFT).
- Decode produces one token per pass. Each step must read every weight and the whole KV cache from GPU memory to do very little math, so it is memory-bandwidth-bound. It sets the time per output token (TPOT, or inter-token latency).
Why decode is bandwidth-bound: a GPU needs hundreds of FLOPs per byte read to stay busy, but decoding one sequence does about one FLOP per byte (2 FLOPs per weight, 2 bytes per BF16 weight). Back-of-envelope speed at batch size 1: tokens per second ≈ memory bandwidth ÷ bytes read per token. A 7B model in BF16 is 14 GB; with about 3.3 TB/s of bandwidth that caps out near 240 tokens/s. Batching many users shares each weight read across sequences, which is where serving throughput comes from.
Shrinking the KV cache: MQA, GQA, MLA
- Multi-query attention (MQA): all query heads share one key head and one value head. The cache shrinks by a factor of h, with a small quality cost.
- Grouped-query attention (GQA): query heads are split into groups that share KV heads, for example 8 KV heads serving 64 query heads (8× smaller cache). The default in most current open models.
- Multi-head latent attention (MLA): introduced by DeepSeek; stores one compressed latent vector per token and reconstructs keys and values from it, shrinking the cache further.
FlashAttention: exact attention, less memory traffic
Naive attention writes the T × T score matrix to GPU main memory (HBM), reads it back for the softmax, writes it again and reads it once more to multiply by V. Memory traffic, not math, dominates. FlashAttention splits Q, K and V into tiles that fit in fast on-chip SRAM and computes the softmax incrementally with an “online softmax”: it keeps a running maximum and running sum per row and rescales as each new tile arrives.
The T × T matrix is never stored, memory becomes linear in T, and speed improves severalfold. It is exact, not an approximation: the same output up to floating-point rounding. In PyTorch it sits behind scaled_dot_product_attention.
Serving-level techniques
- Continuous batching: instead of waiting for a whole batch to finish, add new requests and drop finished ones at every decode step. Large throughput gains.
- PagedAttention: store the KV cache in fixed-size blocks, like virtual-memory pages, so no memory is reserved but unused and blocks can be shared. Introduced with vLLM.
- Prefix (prompt) caching: requests that share a prefix - a long system prompt, a document, few-shot examples - reuse its cached KV instead of recomputing it. This cuts TTFT and cost, and many APIs bill cached input tokens lower. Put the stable part of a prompt first.
- Chunked prefill and disaggregation: split long prefills into chunks interleaved with decode steps, or run prefill and decode on separate GPU pools, so one long prompt doesn’t stall everyone’s token stream.
Speculative decoding
A small, fast draft model proposes k tokens; the large target model checks all k in one parallel forward pass, which is cheap because it looks like prefill. An accept/reject rule based on both models’ probabilities keeps the output distribution exactly equal to the target model’s. When the draft usually agrees, you get several tokens per expensive pass, often a 2–3× speedup (approximate). Variants replace the draft model with extra prediction heads (Medusa, EAGLE) or n-gram lookups.
Quantization
Store weights - and sometimes activations and the KV cache - in fewer bits. Because decode is bandwidth-bound, halving the bytes nearly doubles speed and halves memory.
| Precision | Memory for 70B weights | Typical quality impact |
|---|---|---|
| BF16 | ≈ 140 GB | baseline |
| INT8 / FP8 | ≈ 70 GB | usually negligible |
| INT4 (GPTQ, AWQ and others) | ≈ 35 GB | small, task-dependent |
| 3 bits or fewer | < 27 GB | noticeable degradation |
Methods differ in how they choose roundings to minimize error: GPTQ uses second-order information layer by layer; AWQ protects the weights that matter most for large activations. A complication is outlier features - a few dimensions with huge activations - which naive quantization destroys; LLM.int8() handled them separately. GGUF is the common file format for local inference with llama.cpp.
Serving stacks and metrics
Engines you’ll meet: vLLM, SGLang, TensorRT-LLM and Hugging Face TGI on servers; llama.cpp and Ollama locally. Track TTFT, TPOT, throughput (tokens/s per GPU), cost per million tokens, and goodput (requests served within latency targets).
Questions you should now be able to answer
- Compute the KV cache per token for L = 80, 8 KV heads, d_h = 128 in BF16. How many 32k-token sequences fit in 40 GB?
- Why is decode memory-bound while prefill is compute-bound?
- Why does batching raise throughput but not single-user speed?
- How does FlashAttention compute an exact softmax without seeing a whole row at once?
- Why does speculative decoding leave the output distribution unchanged?
- Your chatbot has a 6k-token system prompt. Which optimization helps most, and how should you order the prompt?