AI Infrastructure series · Part 11 of 24

Part 9 said an LLM predicts the next token. This part opens the model and follows one question, “What is the capital of France?”, through Llama 3.1 8B: every matrix it passes through and the shape of each one, where the 8 billion parameters sit, which steps do maths and which mostly read memory, and what gets stored as the KV cache. Once you can see this walk-through, prefill, decode, KV cache sizing and multi-GPU splits stop being separate topics. They are all consequences of the same few matrices.

My post on the 2017 “Attention Is All You Need” paper covers where this design came from and why it changed infrastructure strategy. This one stays inside a running model, with the real shapes of a model you can download today.

What happens between pressing Enter and the first word?

A transformer language model is simpler than its reputation. It is a lookup table at the front, a stack of identical layers in the middle, and a scoring table at the back. Llama 3.1 8B has 32 layers; the 70B model has 80. Every token becomes a list of 4,096 numbers (8,192 in the 70B) that travels down through every layer, each layer adding a correction. At the end, the list for the last token becomes a score for every word the model knows.

Think of each row of the residual stream as a token’s working notes. At the start, the row for ” capital” only says “the word capital”; by the last layer, the row for “?” holds enough to predict ” Paris”. The rest of this post opens step 3.

What happens inside one layer?

Every layer has two halves, and both follow the same pattern: normalise the input, compute something, add the result back onto the input. The first half is attention, the only place where tokens look at each other. The second half is the MLP, where each token is processed on its own. The diagram shows one layer of the 8B model with the shapes for our 7-token question.

The table lists every step of that layer with its weight matrix and the shape of what comes out. The prefill column is our 7-token question. The decode column is the first answer token, which has 7 earlier tokens plus itself in its history, so its attention covers 8 positions.

Two things stand out. Only seven steps have weights, and all seven are plain matrix multiplications, which is exactly the work Tensor Cores are built for (part 3). And the decode column is the prefill column with 7 replaced by 1, except for attention, where the second dimension is the whole history. That one difference explains most of what follows.

What does RMSNorm do?

RMSNorm takes a token’s 4,096 numbers, computes their root mean square (square each, average, square root), divides every number by it, and multiplies each by a learned scale. The result has a predictable size however large the residual stream has grown, which keeps 32 layers of additions inside the range 16-bit arithmetic handles well. It costs almost nothing in FLOPs, but it reads and writes the whole activation, so inference engines fuse it into the neighbouring kernel. Llama applies it before each half, plus once after the last layer.

Why are there residual connections?

Each half adds its output onto its input: h = x + Attention(RMSNorm(x)), then out = h + MLP(RMSNorm(h)). A layer never replaces the token’s notes, it only writes a correction on top, which is what makes deep stacks trainable. For infrastructure, the residual stream is the one tensor that has to cross a boundary when a model is split by layers across GPUs: 4,096 numbers per token, 8 KB in 16-bit for the 8B model and 16 KB for the 70B.

How does attention work, step by step?

Attention is the only step in which one token reads information from other tokens. Each token’s normalised row is multiplied by three weight matrices to produce three vectors, split into heads of 128 numbers:

Then, in every head:

A worked example with three tokens

Real heads use 128 numbers per vector; two are enough to see the mechanism. Take the last three tokens before “?”: ” capital”, ” of”, ” France”. We compute what ” France” gathers. Its query is q = [1, 2]. The keys and values are made-up numbers chosen to show the pattern (illustrative). With 2 numbers per vector, the scale is √2 ≈ 1.414.

The mix is 0.85 × [4, 0] + 0.10 × [0, 2] + 0.05 × [2, 2] = [3.39 + 0.10, 0.20 + 0.10] ≈ [3.49, 0.30]. Without the scaling step the scores would stay 5, 2 and 1 and the weights would be 0.94, 0.05 and 0.02: sharper, with less of the other tokens getting through. With 128 numbers instead of 2, the difference is far larger, which is why the √128 matters.

At real scale, with 8,000 tokens of history, the newest token’s query is compared with 8,000 stored keys in each of 32 heads and 32 layers, and all those keys and values must be read from memory. That is the cost of long context.

Why do 32 query heads share only 8 KV heads?

Splitting attention into heads lets the model look for several kinds of relationship at once: one head may track which noun a pronoun refers to, another the previous word. Each head runs the steps above on its own 128-number slice. In the original transformer every head also had its own keys and values, which for Llama 3.1 8B would mean storing 32 heads of K and V per layer per token.

Grouped-query attention (GQA, Ainslie et al. 2023) keeps 32 query heads but only 8 key-value heads. Each group of 4 query heads shares one set of keys and values. The query heads still ask different questions; they just look them up in the same index. The K and V matrices shrink from 4,096 × 4,096 to 4,096 × 1,024, and so does everything stored per token.

The cache size per token is 2 (K and V) × layers × KV heads × 128 × 2 bytes. For the 8B model that is 2 × 32 × 8 × 128 × 2 = 131,072 bytes, the 128 KB used throughout this series. The 70B model has 80 layers and also 8 KV heads, which gives 320 KB. Multi-query attention, with a single KV head, saves more but costs quality; 8 heads is the compromise most current open models chose.

16-bit K and V, one sequence, binary units (1 GB = 1,024 MB). The full multi-head rows are hypothetical, for comparison.

Notice what GQA does not change: there are still 32 query heads, so the score maths is the same. GQA cuts memory and memory reads, which is what limits decode. Part 14 sizes the cache for real workloads.

How does the model know word order?

Attention on its own ignores order: shuffle the keys and values and the weighted sum is identical. Something has to tell the model that ” France” came right after ” of”.

Llama uses rotary position embeddings (RoPE, Su et al. 2021). Before scores are computed, each query and key is rotated by an amount that depends on its position. The 128 numbers of a head are treated as 64 pairs, and each pair is rotated like the hand of a clock: some pairs turn quickly as position increases, others very slowly. When a rotated query meets a rotated key, the dot product depends on how far apart the two positions are, so attention sees relative distance without any extra weights. Llama 3 sets the base of these rotation speeds (rope_theta) to 500,000 so the slowest pairs turn slowly enough to tell apart positions far into a long document, and Llama 3.1 adds a rope_scaling setting that stretches positions trained at 8,192 tokens out to 131,072.

RoPE has no parameters and only touches Q and K, never V. It has one practical consequence: keys are stored in the KV cache already rotated, so a cached key belongs to one position. That is why prefix caching in serving engines can reuse cached blocks only when the same tokens sit at the same positions, which in practice means an identical prefix at the start of the prompt.

What does the MLP do?

After attention, each token’s row goes through the MLP (also called the feed-forward block). Unlike attention, the MLP never looks at other tokens. It does the same thing to every row, which makes its cost per token fixed, whatever the length of the conversation.

The 4,096 numbers are multiplied by two separate matrices, gate and up, each producing 14,336 numbers. The gate goes through SiLU (x times the sigmoid of x), a smooth function that passes large positive values and squashes negative ones towards zero. The two 14,336-number vectors are then multiplied number by number, so the gate decides how much of each “up” feature gets through. The down matrix projects the result back to 4,096. This gated design is called SwiGLU (Shazeer 2020); it replaced the simpler two-matrix MLP of the original transformer because it trained better for the same compute.

Research that probes trained models (for example Geva et al. 2021) finds MLP neurons that respond to specific patterns and push the output towards specific words, which is why the MLP is often described as where much stored knowledge lives. It also holds most of the weights: 176.2 million per layer, against 41.9 million for attention. That is the block Mixture of Experts models (part 12) replace with many smaller expert MLPs, of which each token uses only a few.

Where do the 8 billion parameters sit?

Every number in this post comes from the model’s config.json, the file that ships with the weights. These are the fields that set the shapes for Llama 3.1 8B:

"hidden_size": 4096,            # width of the residual stream
"intermediate_size": 14336,     # MLP width
"num_hidden_layers": 32,
"num_attention_heads": 32,      # query heads
"num_key_value_heads": 8,       # KV heads (GQA)
"head_dim": 128,
"vocab_size": 128256,
"hidden_act": "silu",           # SwiGLU MLP
"rms_norm_eps": 1e-05,
"rope_theta": 500000.0,
"max_position_embeddings": 131072,
"tie_word_embeddings": false    # input and output tables are separate

Multiply out the matrix shapes and you get the full count. The 70B model uses hidden size 8,192, MLP size 28,672, 80 layers and 64 query heads with the same 8 KV heads and vocabulary.

Grouped into blocks, the 8B model has 5.64B in the MLPs, 1.34B in attention and 1.05B in the two token tables, 8.03B in all. The 70B has 56.4B in MLPs, 12.1B in attention and 2.1B in token tables, 70.55B in all.

Attention, the part the paper is named after, holds about a sixth of the weights. The vocabulary tables matter in small models (13% of the 8B) and fade in large ones (3% of the 70B), because they grow with the width while the layers grow with width squared times depth. In 16-bit, the 8B weights are 16.1 GB and the 70B weights 141 GB, which is the starting point for part 7’s memory sizing.

Where do the FLOPs go for each token?

Every weight in a matrix multiplication is used once per token, for one multiply and one add. That gives the standard rule: about 2 × parameters floating-point operations per token. The input embedding is a lookup, not a multiplication, so it does not count; the output head does. For the 8B model, that is 2 × (8.03B − 0.525B) ≈ 15.0 billion operations (GFLOPs) per token. For the 70B it is about 139 GFLOPs.

Attention adds work that is not in the parameter count: the scores and the weighted sum of values. For a token with n positions of history, each layer does 2 × n × 4,096 operations for the scores (32 heads × 128 = 4,096) and the same again for the mix, so 4 × n × 4,096 × 32 layers in total. That grows with context, while the weight work does not.

GFLOPs for one new token at that position. Attention share = attention ÷ (weights + attention).

GFLOPs per new token for Llama 3.1 8B at each history length.

The worked number to remember: at 32,768 tokens of history, the attention maths for one new token in the 8B model (17.2 GFLOPs) is larger than the maths for all 8 billion weights (15.0). For prefill the effect is quadratic: a prompt of n tokens does roughly half of n × n score work per layer because of the triangular mask. A 2,000-token prompt on the 8B costs about 2,000 × 15.0 = 30.0 TFLOPs for the weights plus about 1.0 TFLOPs for attention.

What reads memory and what does maths, in prefill and decode?

Counting FLOPs is half the story. A GPU can only use a weight after reading it from memory, and in 16-bit each weight is 2 bytes. In one pass with T tokens, each weight does 2 × T operations. So the ratio of maths to bytes read is roughly T operations per byte: the number of tokens going through the layer together.

An H100 SXM does 989 TFLOPS in BF16 and reads 3.35 TB/s, a ratio of about 295 operations per byte (part 4). Any pass with fewer than about 300 tokens in flight is waiting on memory; any pass with more is waiting on maths. Our 7-token question is so short that even its prefill is memory-bound. A 2,000-token prompt is not.

Ideal peaks, illustrative: real kernels reach a fraction of both, but the bottleneck is the same.

Prefill uses each weight it reads 2,000 times, so the Tensor Cores are the bottleneck and a longer prompt means a longer wait for the first token. Decode reads all 16 GB to do 16 GFLOPs of maths, spending about 300 times longer moving data than calculating. That is why batching many users into one decode step is nearly free until the batch reaches a few hundred (part 17), and why part 13 treats prefill and decode as two different workloads.

What is the KV cache in this walk-through?

It is exactly the K and V outputs from the attention half: 8 heads × 128 numbers for K and the same for V, in each of 32 layers, for every token so far. That is 2 × 32 × 8 × 128 × 2 bytes = 128 KB per token, and 896 KB for our 7-token question. During decode, the new token computes its own K and V, appends them, and reads every earlier token’s K and V to compute its attention. Nothing else from earlier tokens is needed, which is why the cache is enough to skip recomputing them.

The cache read grows with every token. At 2,000 tokens it adds 262 MB to the 16.1 GB weight read, under 2%. At 131,072 tokens it is about 17 GB, more than the weights, so each decode step takes roughly twice as long as it did at the start of the conversation. Multiply by the number of users in a batch and the cache, not the model, becomes what fills the GPU. Part 14 works through that sizing.

How are the layers split across GPUs?

Llama 3.1 70B in 16-bit is 141 GB of weights, more than one 80 GB H100 holds, before any KV cache. The walk-through above shows the two natural ways to cut it: across the width of each layer, or across the depth of the stack.

Tensor parallelism (from Megatron-LM, Shoeybi et al. 2019) cuts each layer’s matrices. Q, K and V are split by heads: with 4 GPUs, each holds 16 of the 64 query heads and 2 of the 8 KV heads, and computes its heads without talking to the others. Wo is split by rows, so each GPU produces a partial sum of the full 8,192-number row, and an all-reduce adds the four together. The MLP is split the same way (gate and up into slices of 7,168 columns, down by rows), followed by another all-reduce. That is 2 all-reduces per layer and 160 per forward pass for the 70B. Each carries only 16 KB per token, but no GPU can continue until it finishes, so tensor parallelism stays inside a server on NVLink (part 6). The KV cache splits too: each GPU stores only its 2 KV heads, 80 KB per token instead of 320 KB. With 8 KV heads, a tensor-parallel size above 8 forces engines to copy KV heads onto several GPUs.

Pipeline parallelism gives each GPU a block of whole layers. Only the residual stream crosses between stages, 16 KB per token at each of 3 boundaries, so it tolerates slower links and can span servers. The catch is that layers run in order: one request visits GPU 1, then 2, 3 and 4, and the others wait unless the engine keeps several batches in flight. Pipeline parallelism adds memory capacity, not speed for a single request.

In vLLM both are a single flag. The first command splits the 70B across 4 GPUs in one server; the second uses two 8-GPU servers, tensor parallel inside each and two pipeline stages across them:

vllm serve meta-llama/Llama-3.1-70B-Instruct --tensor-parallel-size 4

vllm serve meta-llama/Llama-3.1-405B-Instruct \
  --tensor-parallel-size 8 --pipeline-parallel-size 2

Part 5 covers multi-GPU basics and part 20 works out how many GPUs a given model and workload needs.

What do people get wrong about transformers?

• “Attention is where the parameters are.” It holds 17% of them. The MLP blocks hold 70% in the 8B and 80% in the 70B. Attention’s cost is in the KV cache and in long-context maths, not in weights.
• “A longer context only costs memory.” It also costs compute. At 32,768 tokens of history, attention maths for one token of the 8B is larger than the maths for all its weights.
• “GQA makes attention cheaper to compute.” It makes it cheaper to store and read. There are still 32 query heads doing the same score maths.
• “The model re-reads the whole conversation for every word.” Within one response, no: the KV cache holds the earlier tokens. Across requests, often yes: chat APIs are stateless, so each new turn sends the whole history and it is prefilled again unless the server’s prefix cache still holds it.
• “Splitting layers across more GPUs makes one answer faster.” Layer 2 needs layer 1’s output, so pipeline stages run in turn. Only tensor parallelism shortens a single request’s step, and only while its all-reduces stay fast.
• “The model scores the whole vocabulary for every prompt token.” In training it does. When generating, engines compute logits only for the last position, because that is the only one used to pick the next token.

Short answers to common questions

Why is the “8B” model 8.03 billion parameters?

Model names are rounded. Adding up the published shapes gives 8.030 billion, and the 70B gives 70.55 billion.

Does a longer prompt make the model use more memory?

The weights stay the same size. The KV cache grows by 128 KB per token for the 8B (320 KB for the 70B), and prefill needs temporary working memory that grows with prompt length.

Are the input embedding and the output head the same table?

Not in Llama 3.1 8B or 70B. The config sets tie_word_embeddings to false, so there are two separate tables of 0.525B (8B) or 1.051B (70B) parameters each.

Do all LLMs use these same parts?

Most current open decoder models (Llama, Qwen, Mistral) use RMSNorm, RoPE, SwiGLU and grouped-query attention with different sizes. Some change one block: Mixture of Experts models replace the MLP (part 12), and DeepSeek-V3 replaces GQA with a compressed attention (MLA) that stores about 69 KB per token.

Sources: shapes and settings from Meta’s published Llama 3.1 configurations (8B: 32 layers, hidden size 4,096, MLP size 14,336, 32 query heads, 8 key-value heads, head size 128, vocabulary 128,256, rope_theta 500,000, 131,072 positions, untied embeddings; 70B: 80 layers, hidden size 8,192, MLP size 28,672, 64 query heads, 8 key-value heads). Parameter, FLOP and KV cache figures are counted from those shapes. Methods: Vaswani et al. 2017 (transformer, scaled dot-product attention); Zhang and Sennrich 2019 (RMSNorm); Su et al. 2021 (RoPE); Shazeer 2020 (SwiGLU); Ainslie et al. 2023 (grouped-query attention); Geva et al. 2021 (MLP layers as key-value memories); Shoeybi et al. 2019 (Megatron-LM tensor parallelism). H100 SXM figures from NVIDIA’s datasheet. Attention weights in the worked example are illustrative.