Part 2 · The Transformer, Mechanically

Part 1 ended on a hard fact: LLM inference is memory-bound, not compute-bound. That conclusion is only useful if you know which memory traffic dominates and why. To know that, you have to open the box. This part is the box.

Why mechanics before optimization

Every optimization in the rest of this book — KV caches, FlashAttention, continuous batching, quantization, Mixture of Experts — is a targeted attack on one specific line in the transformer’s cost. You cannot evaluate an attack you don’t understand. If you don’t know that the attention matrix grows as the square of sequence length, “FlashAttention saves memory” is a slogan, not knowledge. If you don’t know that most parameters live in the feed-forward network, “Mixture of Experts” sounds like magic instead of arithmetic.

So this part is deliberately mechanical. We trace a single request from raw text to a generated token, stopping at each component to ask the throughline question: what does this cost — in latency, memory, and dollars — and where is the lever to make it cheaper? By the end you’ll have a mental cost model precise enough that the Part 3 optimizations read as obvious consequences rather than tricks.

Note

We care about inference mechanics — running a trained model. Training reuses the same components but has its own cost regime, covered in Part 7. Keep the two separate in your head.

The roadmap

Here is the path, in dependency order. Each page is one box in this pipeline.

  text
   │  ┌─────────────────────────────────────────────┐
   ▼  │           one transformer layer (×N)         │
 [tokenizer] ──► [embeddings] ──►│  [attention] ──► [FFN / MLP]  │──► ... ──► [logits] ──► next token
   IDs            vectors        │   mixes tokens   most params  │
                                 └─────────────────────────────────┘

1. Tokenization & Embeddings — text becomes integer token IDs, then each ID is looked up as a vector. This is where the bill starts: you pay per token, in and out.
2. Attention From Scratch — how tokens share information via Q, K, V and softmax. This is the component whose cost grows as O(n²) in sequence length — the reason long context is expensive.
3. The FFN: Where the Parameters Live — the per-token MLP that holds most of the model’s weights, and therefore most of its weight-memory and matmul FLOPs.
4. Prefill vs Decode: The Two Phases — the key insight of the part. Processing the prompt and generating tokens have opposite bottlenecks. Almost every serving technique exists because of this split.
5. Why the KV Cache Exists — the optimization that turns naive O(n²) decode into O(n) by spending memory to save compute. Its memory footprint becomes the dominant runtime constraint that Part 3 spends its whole length managing.

How the pieces stack

A real model is just this block repeated. Llama-3-8B has 32 of these layers; GPT-3-scale models have ~96. Each layer does the same two things — attention, then FFN — on every token position. That repetition is why a single architectural cost, multiplied by layer count and sequence length, becomes a dollar figure on an invoice.

The one number to carry forward

By the end of this part you should be able to say, for any model: its parameters set the cost of decode (memory-bound weight loading), and its sequence length sets the cost of attention and the KV cache (which grows linearly per token but quadratically in attention compute). Hold those two levers and the rest of the book is detail.

Read the pages in order — each genuinely depends on the one before. When you reach prefill-vs-decode, slow down; it’s the hinge the entire serving stack swings on.

It actually happened — “Attention Is All You Need” (2017)

The architecture this whole part dissects arrived in one 2017 paper: Vaswani et al.’s Attention Is All You Need. Its radical move was deletion — it threw out the recurrence (RNNs and LSTMs) that sequence models had relied on and kept only attention, so every token could see every other token in a single parallel pass instead of being processed left-to-right one step at a time. That parallelism is exactly what makes prefill compute-bound and lets training scale across thousands of GPUs — and the very same all-pairs attention is what saddles inference with the O(n²) cost the rest of this book spends its length taming. The design that unlocked modern LLMs and the design that makes them expensive are one and the same.

Check your understanding

1. Why does the book insist on understanding transformer mechanics before studying optimizations like FlashAttention or MoE?
2. Name the two repeated sub-components inside a single transformer layer.
3. Which mechanical fact about attention makes long context expensive, and how does its cost scale with sequence length?
4. Where in the transformer do most of the parameters live, and why does that matter for inference cost?
5. The throughline question asks three things about every technique’s cost. What are the three?

Show answers

1. Because each optimization is a targeted attack on one specific line of the transformer’s cost; you can’t judge whether an attack helps unless you understand the cost it targets. Without the mechanics, the techniques are slogans, not knowledge.
2. Attention (which mixes information across token positions) and the feed-forward network / MLP (applied per token).
3. The attention mechanism builds an n×n score matrix over the sequence, so its compute and memory grow as O(n²) — doubling the sequence roughly quadruples that cost.
4. In the FFN/MLP plus the projection matrices, which is why weight-memory and matmul FLOPs are dominated by the FFN; it also sets up Mixture of Experts.
5. Latency, memory, and dollars — and how the technique makes each cheaper.