Part I taught you the machine. Part II is about the few kernels that matter on it, so first you need to know what a neural network actually asks the GPU to do. At inference time a model is a fixed program: a list of operations on arrays whose shapes are known in advance, applied to weights that never change. Almost every operation falls into one of four classes. Matrix multiplies (every linear layer) do nearly all the arithmetic and can sit far right on the roofline. Elementwise operations (activations, bias and residual adds), row reductions (layernorm, softmax) and data movement (embedding lookups, transposes) do almost no arithmetic and are always memory-bound. So the work splits cleanly: make the matrix multiplies fast (chapters 7 and 8), and make everything else move as few bytes as possible, mostly by fusing it into its neighbours (chapter 9). This chapter builds that picture from an MLP block up to all of GPT-2 small, counts its parameters, FLOPs and bytes, and ends with the number formats (FP32 down to FP4) that decide how many bytes there are.
A neural network is a function with a very large number of adjustable constants. You give it an input x (a sentence, an image), it returns an output y (the next word, a label), and the constants, called parameters or weights, decide which function it is. Training searches for weights that make the outputs good, by running the model forward, measuring the error, and running backward to adjust every weight a little; that costs about three times the arithmetic of the forward pass alone, over trillions of examples. Inference uses the finished weights: forward only, weights frozen. This course is about inference, because that is where models spend their lives and where the money goes.
The function is built from layers: small, standard operations applied one after another, each one taking arrays of numbers in and giving arrays out. A modern language model uses maybe ten distinct kinds of layer, repeated dozens of times. Seen from the GPU, three facts make this workload unusual:
So from here on, “running a model” means: launch a known list of kernels, each one an operation on arrays. The rest of this chapter is about what those operations are and what each one costs.
The arrays a model works on are called tensors: an array with a shape (the size of each dimension) and a data type (FP32, FP16…, section 7). A vector is a tensor of shape [768], a matrix [768, 3072]. Whatever the shape, memory is still one long line (chapter 2), so every tensor also has a layout: the rule that turns an index into an offset. Frameworks use row-major by default: the last dimension is contiguous.
A language model reads text as tokens, pieces of words, each one an integer ID from a fixed vocabulary (GPT-2’s has 50,257). The first layer turns each token into a vector of C numbers (768 in GPT-2 small); every later layer takes those vectors and returns new ones of the same size. So the data flowing through the model, the activations, has three dimensions: B sequences in the batch, T tokens per sequence, C channels per token. Shape [B, T, C].
The step in the middle of the figure is the most useful idea in this chapter. Almost every layer (linear layers, activations, layernorm) acts on each token on its own: token 5 of sequence 2 does not care which sequence it is in. So the kernels see just a matrix of M = B·T rows, one per token in flight, each C numbers wide. Only attention (chapter 10) mixes tokens, and only within a sequence. Keep M in mind: it is the number that moves every operation along the roofline.
Viewing [B, T, C] as [B·T, C] changes no bytes: it is the same line of memory with a different indexing rule. Frameworks call that a view. Swapping two dimensions, for example [B, T, H, D] → [B, H, T, D] to separate attention heads (chapter 10), changes which elements are neighbours in memory, so producing it contiguous means copying the whole tensor. Good kernels avoid the copy by reading the data in its original layout, with a different index formula.
The workhorse layer takes each token’s vector of K numbers and returns N numbers, each one a weighted sum of all the inputs plus a constant:
y[m, n] = b[n] + ∑k x[m, k] · W[k, n]
For all tokens at once that is Y = XW + b: an [M, K] matrix times a [K, N] weight matrix, the operation BLAS libraries call GEMM (general matrix multiply). Every output element is a dot product of length K: K multiplies and K adds.
Count it for the first MLP layer of GPT-2 small, which maps 768 channels to 3,072 (K = 768, N = 3,072), in FP32:
| M = 1 (one user, decode) | M = 512 (a prompt) | |
|---|---|---|
| FLOPs, 2MKN | 4.7 M | 2.42 G |
| Bytes: X, W, Y once each | 9.46 MB (almost all W) | 17.3 MB |
| Intensity | 0.5 FLOP/B | 140 FLOP/B |
| Roofline time on your card | 16 µs, memory-bound | 78 µs, compute-bound |
Same layer, same weights, a factor of 280 in intensity. At M = 1 each weight is used for one multiply-add and thrown away; at M = 512 it is used 512 times. (The roofline times use 88% of the 672 GB/s spec for memory and the 30.9 TFLOPS FP32 peak for compute.)
The second, and by a lot. One thread per output means the first layer launches 3,072 threads (12 blocks of 256) and the second only 768 threads: 3 blocks. On a 48-SM GPU, 3 blocks occupy 3 SMs and leave 45 idle, and each of those 768 threads walks a 3,072-long dot product one load after another. Chapter 3’s rule is that bandwidth needs bytes in flight: 3 SMs cannot keep enough requests outstanding to pull more than a small fraction of 672 GB/s. Expect the second layer to take several times longer than the first even though it moves the same bytes, and both to be far from the 16 µs roofline. It is chapter 2’s “many more blocks than SMs” rule, broken by a shape. Decode kernels fix it by splitting the K loop across many blocks (split-K, chapter 15). Lab 5 measures both.
The formula above stores W as [K, N] = [in, out], and that is how GPT-2’s original checkpoints store it. PyTorch’s nn.Linear stores the same numbers as [out, in] and computes XWT. Nothing about the math changes, but the kernel does: in one layout a column of W is contiguous, in the other a row is. BLAS libraries name the variants by which operand is transposed (NN, NT, TN), and a fast kernel for one is not automatically fast for the other. Checking the layout before writing the kernel is the cheapest bug you will ever avoid; lab 5 and the rest of Part II use [K, N], so Y = XW with no transpose.
The + b is a separate operation only on paper. Adding one number per output column costs nothing if the matmul kernel does it while the result is still in registers, just before writing it. That final step of a matmul kernel is called its epilogue, and bias, activation and residual add are all classic epilogue passengers (chapter 9). Lab 5 keeps them as separate kernels on purpose, so that you can see what they cost when they travel alone.
Stacking linear layers alone buys nothing: a linear function of a linear function is still one linear function, and all those weights would collapse into one matrix. Everything that makes a network more than a single matrix lives in the other layers. They do very little arithmetic, but there are many of them.
An activation function is applied to every number independently, which is what breaks the linearity. GPT-2 uses GELU, in its tanh approximation:
gelu(x) = ½x (1 + tanh(√(2/π) (x + 0.044715x3)))
Roughly: pass positive values through, squash negative ones towards zero, smoothly. Other elementwise operations in a transformer: bias adds, the residual add (out = x + layer(x), which lets each layer add a correction to its input instead of replacing it), scaling, and conversions between number formats.
The cost model is the simplest one in the course. Each element is read once and written once, 8 bytes in FP32, for a handful of FLOPs: intensity around 1 FLOP per byte at best, fifty times left of the ridge. Every elementwise kernel is memory-bound, and its time is its bytes divided by bandwidth. GELU on a [512, 3,072] FP32 activation moves 12.6 MB, so about 21 µs on your card, however few or many FLOPs the tanh takes. Unless it is tiny: at M = 1 the same GELU moves 25 KB, which takes 0.04 µs of DRAM time and a few µs of launch overhead. Small elementwise kernels are overhead-bound.
Layer normalization rescales each token’s vector to mean 0 and variance 1, then applies a learned scale γ and shift β per channel. It keeps the numbers in a well-behaved range as they pass through dozens of layers:
μ = mean(x), σ2 = mean((x − μ)2), y = γ · (x − μ) / √(σ2 + ε) + β
Softmax turns a row of arbitrary scores into positive weights that sum to 1, and it is at the heart of attention (chapter 10):
softmax(x)i = exi − max(x) / ∑j exj − max(x)
(Subtracting the maximum changes nothing mathematically and prevents overflow: in FP32, ex is infinite for any x above about 88.7.)
Both have the same shape: reduce a row to one or two numbers, then apply an elementwise formula that uses them. The reduction is chapter 5’s, one per row, and like elementwise ops both are firmly memory-bound. Their special problem is that the formula needs the reduction’s result before it can start, so a naive implementation reads each row two or three times: once for the max or mean, once for the sum or variance, once to write the output. Chapter 9 gets that down to one read (Welford’s algorithm, online softmax), and chapter 14 uses online softmax to build FlashAttention.
The first layer is an embedding lookup: token ID t selects row t of a [50,257, 768] table. That is chapter 3’s gather: rows scattered across the table, but each row contiguous, so one warp per row is coalesced. Zero FLOPs, 3 KB per token. The other movers are the transposes and head splits of attention, which good kernels fold into their index arithmetic instead of running as copies.
Each separate kernel reads and writes the whole tensor: 2 × 6.3 MB = 12.6 MB, about 21 µs each at ~590 GB/s, so roughly 64 µs for the three (the conversion writes FP16, a bit less). A fused kernel reads the FP32 input once and writes the FP16 output once: 6.3 + 3.1 = 9.4 MB, about 16 µs. Four times faster, and the arithmetic is unchanged. Fusing it into the matmul’s epilogue removes even that read. This is the whole argument for fusion, and the reason “count the bytes, not the FLOPs” is the first rule for everything that is not a matmul.
| Class | Examples | FLOP per byte (FP32) | Bound | Where we optimize it |
|---|---|---|---|---|
| Matrix multiply | linear layers; attention’s QKT and ·V | ≈ M/2 for small M, up to hundreds | memory at small M, compute at large M | chapters 7, 8, 12 |
| Elementwise | GELU, bias, residual add, scale, cast | 0.08–1 | memory (or launch overhead when tiny) | fusion, chapter 9 |
| Row reduction | layernorm, RMSNorm, softmax | ≈ 0.5–1 | memory; extra passes over the row | chapter 9; attention: chapter 14 |
| Data movement | embedding gather, transpose, concat | 0 | memory | fold into the neighbours’ indexing |
Two consequences shape the rest of Part II. First, the matmuls are the only operations where the kernel’s arithmetic skill matters: for everything else, a perfect kernel is one that runs at the DRAM ceiling, and the only way to beat that is to move fewer bytes. Second, the small operations are cheap only in FLOPs. In a large-M forward pass they are a real slice of the time, and at M = 1 they are dominated by launch overhead. The next section puts numbers on both.
GPT-2 small is the model your spine project will run. Its whole configuration fits in one line: vocabulary V = 50,257 tokens, context up to 1,024 tokens, L = 12 layers, width C = 768, 12 attention heads of 64 channels, MLP width 4C = 3,072. Every layer (a transformer block) has the same two halves, each wrapped in a residual connection:
| Piece | Shape | Parameters |
|---|---|---|
| Token embedding (also the LM head) | 50,257 × 768 | 38,597,376 |
| Position embedding | 1,024 × 768 | 786,432 |
| Per block: Q, K, V projection | 768 × 2,304 + 2,304 | 1,771,776 |
| Per block: attention output | 768 × 768 + 768 | 590,592 |
| Per block: MLP up | 768 × 3,072 + 3,072 | 2,362,368 |
| Per block: MLP down | 3,072 × 768 + 768 | 2,360,064 |
| Per block: two layernorms | 4 × 768 | 3,072 |
| 12 blocks | 12 × 7,087,872 | 85,054,464 |
| Final layernorm | 2 × 768 | 1,536 |
| Total | 124,439,808 |
That is the “124M” in GPT-2 small’s name. In FP32 it is 498 MB, which is chapter 1’s “0.5 GB”: the number of bytes every decode step must read.
Each token passes through every linear layer once, for 2 FLOPs per weight. The four linear layers of a block hold 7,077,888 weights, so 12 blocks cost 170 MFLOPs per token. The LM head multiplies by the whole 38.6M-entry embedding table: 77 MFLOPs more, nearly a third of the total, spent on scoring every word in the vocabulary. Attention adds 4 × 768 FLOPs per layer for every earlier token the new one looks at, 37K per token of context over the 12 layers: 38 MFLOPs at a full 1,024-token context, small next to the rest. Total: about 250 MFLOPs per token, which is chapter 1’s “2 FLOPs per parameter” (2 × 124M = 249M). The rule works because the embedding table, used as a free gather at the input, is paid for in full as the LM head at the output.
About 1,024 × 247 MFLOPs for the linear layers and head, plus attention, where the average token looks back over half the prompt: roughly 0.27 TFLOPs in total. At 30.9 TFLOPS that is about 9 ms, compute-bound (the linear layers sit at intensity in the hundreds). Reading the 498 MB of weights once would take 0.8 ms, so memory is not the limit. That is why prefill is measured in thousands of tokens per second and decode in hundreds: prefill uses every byte of weights 1,024 times, decode once.
Here is decode as a machine. The weights are a plate cylinder that has to turn once for every step, and a batch is several sheets fed through on the same turn.
Now put the pieces together. The census below lists every operation of one forward step of GPT-2 small (bias adds folded into their matmuls), with its bytes, its FLOPs, and its roofline time: the larger of FLOPs/peak and bytes/bandwidth. It models one decode step for a batch of M sequences, each with t tokens of context already cached. Start at M = 1, then raise it, switch to FP16, and turn on the launch floor.
Three things to take away from playing with it:
Every byte count above assumed FP32, 4 bytes per number. The simplest way to make a memory-bound workload faster is to make the numbers smaller, and ML has been doing exactly that for a decade. A floating-point number is a sign, an exponent (which sets the range) and a mantissa (which sets the precision): value = ±1.mantissa × 2exponent. Formats differ in how they split their bits:
| Format | Largest | Spacing just above 1 | Typical role |
|---|---|---|---|
| FP32 | 3.4 × 1038 | 1.2 × 10−7 | accumulators, references, your first engine |
| TF32 | 3.4 × 1038 | 9.8 × 10−4 | tensor-core inputs for FP32 matmuls (chapter 8) |
| BF16 | 3.4 × 1038 | 7.8 × 10−3 | training, weights of most open models |
| FP16 | 65,504 | 9.8 × 10−4 | inference weights and activations |
| FP8 E4M3 | 448 | 0.125 | inference, with scale factors |
| FP8 E5M2 | 57,344 | 0.25 | gradients in FP8 training |
| FP4 E2M1 | 6 | 0.5 | weights, with per-block scales (Part IV) |
Neural networks tolerate this because they are statistical: a weight of 0.0213 stored as 0.0214 changes the output by less than the noise the model was trained to ignore. But there is one place where precision cannot be cut, and you have already met it in chapter 5’s reduction: long sums. A dot product of length 768 adds 768 products into one running total. Once the total is large, each new small product is below the spacing of the format and simply vanishes. Try it:
Hence the rule every fast kernel follows: store and multiply in a narrow format, accumulate in FP32. The inputs are rounded once, which the model tolerates; the sum is never rounded to the narrow format until the very end. Tensor cores are built around exactly this: FP16, BF16 or FP8 inputs, FP32 accumulators (chapter 8). Your spine engine starts in FP32 because it must match PyTorch’s outputs to about 10−5 before anything gets clever, and a mismatch in FP32 is a bug, not rounding.
Range. Gradients and some activations span many orders of magnitude, and FP16 overflows above 65,504 and loses precision below 6 × 10−5, so FP16 training needed loss scaling: multiply the loss by a large factor to keep gradients representable, and back it off when anything overflows. BF16 has FP32’s 8-bit exponent, so anything FP32 can represent, BF16 can too, only more coarsely. Losing 3 bits of precision turned out to hurt less than babysitting the range. For inference, where values are better behaved and known in advance, FP16 is still common, and the formats below 16 bits bring back the range problem in a new form: that is what the scale factors are for.
Put the census and the roofline side by side and the next five chapters write themselves:
| Chapter | What it does to the census |
|---|---|
| 7. Matmul I | Takes the linear layers from the naive kernel (a few percent of peak) to shared-memory tiling, and finds out why tiling alone stops near 10%. |
| 8. Matmul II | Register blocking, warp tiling and tensor cores: the linear layers within reach of cuBLAS, in FP32 and FP16. |
| 9. Softmax, layernorm, fusion | The small operations at the DRAM ceiling in one pass, and fused into their neighbours so their bytes disappear. |
| 10. The transformer from the GPU’s view | Attention in detail, the KV cache, and the full block as kernels. |
| 11. Lab: GPT-2 forward pass | The spine project’s milestone: all of it assembled, matching PyTorch, benchmarked. |
The order is not arbitrary: the matmul is where the FLOPs are and where the hardest kernel skills are learned, and those skills (tiling, register reuse, warp-level cooperation) are what every other kernel in Part III is built from.
One program, mlp, runs the MLP half of a GPT-2 block with one kernel per operation: layernorm, the naive matmul (768 → 3,072), bias, GELU, the naive matmul (3,072 → 768), bias, residual add, plus a fused bias_gelu. It checks all of them against a double-precision CPU reference, then times each one at M = 1, 64 and 512 tokens and prints its bytes, FLOPs, GB/s, GFLOP/s and share of the total.
cd ~/unibo-lessons/mlsys/labs/lab05-mlp-block
make test # every kernel vs the CPU reference: all PASS
make bench # the op-by-op table at M = 1, 64, 512
make ncu-table # optional: DRAM / SM throughput and grid size per kernelBefore you run make bench, write down your predictions:
| Question | What the chapter predicts |
|---|---|
| M = 512: share of the time in the two matmuls | above 95%: they are the only compute-heavy ops, and the naive kernel runs them at a few percent of peak |
| M = 512: GELU, bias, residual GB/s | near your lab 0 ceiling (~590 GB/s) |
| M = 512: bias+GELU fused vs bias then GELU | fused ≈ half the sum of the two |
| M = 1: elementwise ops | a few µs each whatever their size: launch-bound |
| M = 1: matmul fc vs matmul proj | proj several times slower (3 blocks vs 12, section 3) |
| M = 512: naive matmul GFLOP/s | a few percent of the 30.9 TFLOPS peak: chapter 7’s starting point |
Paste back the PASTE: lines (one from make test, three from make bench). Your numbers replace the predictions here, and the M = 512 matmul line becomes chapter 7’s baseline.
Write bias_gelu_residual for the second half of the sublayer (out = x + v + b) and measure it against bias then residual. Then predict, before measuring, how long all of layernorm would take if you fused it with nothing but made it read each row only once instead of twice. (Hint: the second read of a 3 KB row mostly hits L1. What does that do to your prediction?)
Linear layers, activations and normalizations act on each token’s vector independently: the output for one token does not depend on any other token. So which sequence a token belongs to does not matter, and the kernels see one [B·T, C] matrix. Only attention mixes tokens, within a sequence.
About 0.5 FLOP/B at M = 1 (each weight, 4 bytes, used for one multiply-add): memory-bound. About 140 FLOP/B at M = 512 (each weight used 512 times): past the FP32 ridge of ~46, compute-bound.
Neither by much, and the residual add is slightly slower: both are memory-bound, so time is bytes over bandwidth, and the residual add moves 12 bytes per element (two inputs, one output) against GELU’s 8. FLOPs do not enter.
The elementwise formula needs the row’s reduction (max, sum, mean, variance) before it can start, and computing that reduction requires the whole row. So the kernel reads the row to reduce it, then again to normalize it. For short rows the second read hits L1 and costs little; for long rows (softmax over thousands of attention scores) it goes back to L2 or DRAM and doubles the traffic. Chapter 9 removes it.
Each weight of a linear layer is used in one multiply-add per token. It fails for parameters that are not matmul weights: the embedding table used as an input lookup costs no FLOPs, layernorm parameters cost a few. GPT-2 happens to balance it out because the tied embedding table is also the LM head, a full matmul at the output, so the rule gives 249M against a detailed count of ~250M.
Each operation is a kernel launch, and a launch costs a few microseconds no matter how little work it does. GPT-2 small runs over a hundred small kernels per token, each far below the launch floor, so they cost launches, not bytes. The fixes are fusion (fewer kernels) and CUDA graphs (cheaper launches).
A long sum in a narrow format stops growing once the running total is so large that each new term is smaller than half the format’s spacing at that magnitude: the terms round away. Rounding the inputs costs a little accuracy once; rounding the running total costs accuracy at every step. So inputs are narrow, accumulators are FP32.
BF16 keeps FP32’s 8 exponent bits (same range) and has 7 mantissa bits (spacing ~0.8% near 1). FP16 has 10 mantissa bits (~0.1%) but only 5 exponent bits, so it overflows above 65,504 and starts losing precision below ~6×10−5 (the smallest normal number; below it only subnormals remain, down to ~6×10−8).
train_gpt2.c is the clearest listing of every operation and shape. Also his “Let’s reproduce GPT-2 (124M)” lecture.