The hardware underneath

Every lesson so far has counted operations: so many multiplies per token, so many parameters. This lesson looks at the machine that performs them, because the machine explains things the maths alone never will: why a model that "needs" a tenth of a millisecond of arithmetic takes five milliseconds per token, why training runs are spread across thousands of chips in a particular way, and why everyone is shrinking numbers from 32 bits to 8.

Three facts carry the whole lesson:

1. A GPU is thousands of simple arithmetic units doing the same step on different numbers. Neural networks are mostly matrix multiplies, which are exactly that kind of work.
2. Arithmetic is cheap; moving data is expensive. The chip can multiply far faster than its memory can feed it, so speed is usually decided by how many bytes move, not how many operations run.
3. Fewer bits per number helps everywhere at once: more numbers per byte moved, more numbers in memory, and smaller, more numerous multipliers on the chip.

Sections 5 and 6 then apply those facts to many GPUs working together, and to the question every deployment starts with: will the model fit?

Everyday picture A CPU is a few master chefs. Each can cook anything, improvise, and follow a recipe full of "if the sauce splits, do this instead". A GPU is a kitchen of thousands of line cooks who all do the same step at the same moment on different ingredients: "everyone, chop your carrot now." That kitchen is useless for inventing a menu and unbeatable at ten thousand identical salads. A neural network is ten thousand identical salads: nearly all of its work is matrix multiplication, the same multiply-and-add done billions of times on different numbers.

Tiny worked example Multiply a 2 × 3 matrix by a 3 × 2 matrix:

A = [[1, 2, 3], [4, 5, 6]]
B = [[7, 8], [9, 10], [11, 12]]
m, k, n = len(A), len(B), len(B[0])
# each cell: row i of A dotted with column j of B
[[sum(A[i][p] * B[p][j] for p in range(k)) for j in range(n)] for i in range(m)]  # → [[58, 64], [139, 154]]

That count generalises into the most useful formula in this lesson. A FLOP (floating-point operation) is one multiply or one add, and multiplying an m × k matrix by a k × n matrix costs:

m, n, k = 2, 2, 3
# 2 FLOPs per step, k steps per cell, m·n cells
2 * m * n * k  # → 24
flops = 2 * 4096 * 4096 * 4096
flops  # → 137438953472
# milliseconds at 10¹⁵ FLOPs per second
round(flops / 1e15 * 1000, 2)  # → 0.14

Hardware usually does "multiply, then add to a running total" as a single instruction, the fused multiply-add, which is why FLOPs come in pairs.

flowchart LR
  subgraph CPU["CPU: a few master chefs"]
    direction TB
    c1["core: big control unit,<br/>big cache, runs any code"]
    c2["core"]
    c3["core"]
  end
  subgraph GPU["GPU: thousands of line cooks"]
    direction TB
    g1["group of cores:<br/>one instruction, many numbers"]
    g2["group of cores"]
    g3["... about a hundred groups"]
    g4["matrix units: a small<br/>tile multiply per instruction"]
  end
  W["matrix multiply:<br/>m × n independent cells"] --> GPU
  BR["branchy code:<br/>if this then that"] --> CPU

How far does the independence go? Suppose each core takes whole cells and performs one multiply-add per round:

import math
m = n = k = 64
def rounds(cores):
    # ⌈m·n / cores⌉ cells per core, then k multiply-adds per cell
    return math.ceil(m * n / cores) * k
rounds(1), rounds(64), rounds(4096), rounds(8192)  # → (262144, 4096, 64, 64)

In code: counted_matmul multiplies with plain loops and counts every multiply and add; matmul_flops is the 2·m·n·k formula; parallel_rounds is the rounds formula above.

Why it matters in practice. A GPU is fast only when it is given a lot of independent work at once: large matrices, and many sequences processed together. That is why serving systems batch requests together (primer.ml.inference) and why a small model answering one user at a time uses a sliver of the chip. It is also why neural networks look the way they do: architectures that turn into a few big matrix multiplies (the transformer, primer.ml.transformer) won partly because they suit this hardware, while step-by-step recurrences (primer.ml.cnn_rnn) do not.

Everyday picture Back in the kitchen. A cook's hands hold one or two things (the registers). The cutting board holds a few more (the on-chip SRAM, fast memory built into the chip itself). The fridge in the kitchen holds the day's ingredients (HBM, "high-bandwidth memory", the GPU's main memory). The storeroom down the hall is the CPU's memory (host memory). The warehouse across town is the disk. And other restaurants' pantries, reached by courier, are other machines over the network. Every step further out holds more and takes longer to reach. The cooks are fast; what slows the kitchen down is fetching.

Tiny worked example Round, illustrative numbers for one datacenter GPU and the machine around it (orders of magnitude, not any product's specification):

Registers and SRAM never hold a whole gigabyte; the column shows their rate. Notice the jump from HBM to host memory: about 67 times slower. A GPU that has to reach past its own HBM is a cook walking to the storeroom for every carrot.

HBM, host = 3.35e12, 50e9
weights = 16e9
# t = bytes / bandwidth, in milliseconds
round(weights / HBM * 1000, 2)  # → 4.78
round(weights / host * 1000)  # → 320
# how many times slower the storeroom is than the fridge
round(HBM / host)  # → 67

flowchart LR
  ALU["arithmetic units"] <--> R["registers<br/>~20 MB, ~100 TB/s"]
  R <--> S["on-chip SRAM<br/>~50 MB, ~20 TB/s"]
  S <--> H["HBM<br/>~80 GB, ~3.35 TB/s"]
  H <--> D["host memory<br/>~1 TB, ~50 GB/s"]
  D <--> K["local disk<br/>~10 TB, ~10 GB/s"]
  H <--> N["network: other machines<br/>~50 GB/s per GPU"]

In code: MEMORY_HIERARCHY lists the six levels with their capacity and bandwidth; transfer_seconds is the formula above.

Why it matters in practice. A model runs at full speed only if everything it touches every step (weights, activations, KV cache) lives in HBM. Spilling to host memory ("offloading") makes a model fit, at the price of each step waiting tens of times longer. And because HBM itself is slow compared with the arithmetic, the fastest code is the code that makes each trip to HBM count, which is the next section.

Everyday picture A sandwich shop. If the cook walks to the storeroom for each slice of bread for each sandwich, the cook spends the day walking. If the cook carries a tray of bread and a tray of fillings to the bench and makes a batch of sandwiches from them, every trip feeds many sandwiches. Same sandwiches (FLOPs), far fewer trips (bytes). The ratio of the two is the arithmetic intensity: operations done per byte fetched.

Our imaginary GPU does 10¹⁵ FLOPs per second but reads only 3.35 × 10¹² bytes per second from HBM, so it breaks even at about 299 FLOPs per byte (the ridge point of the roofline in primer.ml.inference). Any work doing fewer operations than that for each byte it fetches leaves the arithmetic units waiting on memory.

Tiny worked example Multiply two 4 × 4 matrices: 2 × 4³ = 128 FLOPs. Count the numbers fetched from slow memory (HBM) into fast memory (SRAM):

The arithmetic is identical in all three rows. Only the traffic changes. This trick is tiling: bring a small block of each matrix into fast memory and do every multiplication that block takes part in before throwing it away.

n = 4
# reads = 2n³ / T, for tiles of 1, 2 and 4
[2 * n**3 // T for T in (1, 2, 4)]  # → [128, 64, 32]
def intensity(n, T, b):
    flops = 2 * n**3
    # bytes moved: every read, plus one write per answer cell
    moved = b * (2 * n**3 / T + n**2)
    return flops / moved
round(intensity(4096, 128, 2), 1)  # → 63.0
round(intensity(4096, 128, 1), 1)  # → 126.0

flowchart LR
  subgraph HBM["HBM: big, slow"]
    A["A, in T × T tiles"]
    B["B, in T × T tiles"]
    C["C, the answer"]
  end
  subgraph SRAM["on-chip SRAM: small, fast"]
    a["one tile of A"]
    b["one tile of B"]
    acc["running total for<br/>one T × T tile of C"]
  end
  A -- "load" --> a
  B -- "load" --> b
  a --> mm["multiply-add:<br/>T³ steps, no traffic"]
  b --> mm
  mm --> acc
  mm -. "next pair of tiles along k" .-> A
  acc -- "write once, at the end" --> C

The same idea runs through the rest of this primer:

• FlashAttention (primer.ml.attention) tiles attention: blocks of queries, keys and values are loaded into SRAM, and the n × n score matrix is never written to HBM at all. Same answer, a fraction of the traffic.
• Kernel fusion: adding a bias or applying an activation does about one FLOP per number it reads, hopelessly below 299, so these steps are done inside the matrix-multiply kernel while the tile is still on chip.
• Decode (primer.ml.inference): generating one token for one user reads every weight to do just 2 FLOPs with it, an intensity of about 1. Batching users together is tiling across requests: one read of a weight serves every sequence in the batch.

In code: tiled_matmul runs the tiled multiply and returns a Traffic count of reads, writes, FLOPs and peak fast-memory use; matmul_reads and matmul_intensity are the formulas; primer.ml.inference.ridge_point is the break-even.

Why it matters in practice. Before asking how many FLOPs a piece of work needs, ask how many bytes it moves and how often each byte is reused. Most large speed-ups in modern AI systems (FlashAttention, fused kernels, batching, quantization) change the bytes, not the FLOPs.

Everyday picture Scientific notation on a form with a fixed number of boxes: 6.02 × 10²³. One box holds the sign. A few boxes hold the power of ten, which sets how big or small the number can be: its range. The rest hold the digits, which set how finely it is measured: its precision. With a fixed number of boxes, moving a box from the digits to the power buys range and costs precision. Computers do the same with bits and powers of two, and call the parts the sign, the exponent and the mantissa (the stored digits).

Tiny worked example Store −6.5 in bf16 ("brain float 16": 1 sign bit, 8 exponent bits, 7 mantissa bits).

1. Sign: negative, so the sign bit is 1.
2. Power of two: the largest power of two not above 6.5 is 4 = 2², so 6.5 = 1.625 × 2².
3. Exponent: stored with a bias of 127 added, so that negative powers need no sign of their own: 2 + 127 = 129 = 10000001 in binary.
4. Mantissa: the leading 1 of 1.625 is always there, so it is not stored. The fraction 0.625 = ½ + ⅛ is 0.101 in binary, padded to seven bits: 1010000.
5. The 16 bits: 1 10000001 1010000.

Most numbers are not so lucky. 0.1 has no finite binary expansion, so it is rounded to the nearest value each format can hold: 0.10000000149 in fp32, 0.1000977 in bf16, 0.0999756 in fp16 and 0.1015625 in 8-bit E4M3.

x = -6.5
M, bias = 7, 127
# s: 1 for a negative number
s = 1 if x < 0 else 0
# e: the power of two below |x| is 2², stored with the bias added
e = 2 + bias
e, format(e, "08b")  # → (129, '10000001')
# f: the fraction after the hidden 1 of |x| / 2², as a 7-bit whole number
f = round((abs(x) / 2**2 - 1) * 2**M)
f, format(f, "07b")  # → (80, '1010000')
# decode: (-1)^s × 2^(e - bias) × (1 + f / 2^M)
(-1)**s * 2**(e - bias) * (1 + f / 2**M)  # → -6.5

Two corners of the formula matter in practice. When the stored exponent is 0, the hidden 1 is dropped and the number is a subnormal: it lets values fade gradually towards zero instead of dropping off a cliff, at the cost of fewer significant bits. And IEEE-style formats reserve the all-ones exponent for infinity and NaN ("not a number"), which is where overflowing values go.

flowchart LR
  X["x = −6.5"] --> S["sign: negative<br/>s = 1"]
  X --> P["largest power of two<br/>not above 6.5: 2² = 4"]
  P --> E["exponent: 2 + bias 127<br/>e = 129 = 10000001"]
  P --> F["6.5 / 4 = 1.625<br/>drop the leading 1: .625"]
  F --> R["round .625 to 7 bits<br/>f = 1010000"]
  S --> B["1 | 10000001 | 1010000"]
  E --> B
  R --> B

E4M3 bends the IEEE rules: it has no infinity, and spends that exponent on ordinary numbers instead, which is how 8 bits reach 448 rather than 240.

What the picture means for training (primer.ml.pretraining covers mixed precision in full):

• Range failures. A gradient of 10⁻⁸ becomes exactly 0 in fp16 (below its smallest subnormal) but survives in bf16 as 1.0012 × 10⁻⁸. fp16 training therefore multiplies the loss by a large constant (loss scaling) to lift gradients into range; bf16 training does not need to.
• Precision failures. In bf16, 1 + 0.001 rounds back to exactly 1: a small weight update simply vanishes. So training keeps a master copy of the weights in fp32, adds its sums up in fp32, and uses 16-bit (or 8-bit) only for the big multiplies. That split is mixed precision.

Why smaller formats multiply throughput

Fewer bits per number pays three times:

1. Bytes. Half the bytes per number moves twice the numbers per second through every level of section 2, and fits twice the parameters in HBM. Memory-bound work speeds up directly: the lower bound on time per token for a 70-billion-parameter model (primer.ml.inference) is 41.8 ms at 16-bit, 20.9 ms at 8-bit and 10.4 ms at 4-bit.
2. Intensity. The same tile does twice the FLOPs per byte (section 3: 63 becomes 126), pushing more work past the break-even point.
3. Silicon. A multiplier is the expensive part of the chip, and its size grows with the square of the number of significand bits.

Everyday picture Long multiplication by hand: multiplying two 3-digit numbers means writing a 3 × 3 grid of single-digit products; two 6-digit numbers need a 6 × 6 grid, four times the work. A chip's multiplier is that grid built in wires.

# stored mantissa bits for fp32, fp16, bf16 and fp8 E4M3
mantissa_bits = [23, 10, 7, 3]
# p = M + 1, and cells = p²
[(M + 1) ** 2 for M in mantissa_bits]  # → [576, 121, 64, 16]

In code: FloatFormat describes a format by its exponent and mantissa bits, with FloatFormat.max_value, FloatFormat.min_normal, FloatFormat.min_subnormal and FloatFormat.epsilon derived from them; FP32, BF16, FP16, FP8_E5M2 and FP8_E4M3 are the five formats; encode rounds a number into its three fields with plain arithmetic, decode turns fields back into a number, round_to does both, and bit_string prints the bits; multiplier_cells is the p² formula. The tests check round_to against NumPy's own float16 and float32.

Why it matters in practice. Picking a number format is picking a point on the range-versus-precision curve for each kind of number in a model. Weights and activations tolerate coarse formats; gradients need range; the running sums of long dot products need precision. Modern training and serving use a different format for each, and the savings are among the largest in the field.

Everyday picture A group project. Four people each work through a quarter of the exercises, and then must agree on one combined answer sheet. Sitting at the same table they can compare notes in seconds; living in different cities they must post letters. The more often a team must compare notes, the more it matters who sits at the same table.

Large models are trained on many GPUs because no single one has the memory or the speed. The simplest split is data parallelism: every GPU holds a copy of the model and works on different examples, and after each step their gradients must be added up, so every copy takes the same step. That "add up, and give everyone the total" operation is an all-reduce.

Tiny worked example: a ring all-reduce. Four GPUs each hold a gradient of 8 numbers. Each splits its gradient into 4 chunks of 2 numbers and they sit in a ring, each passing to its right-hand neighbour.

1. Reduce-scatter, 3 steps: each GPU sends one chunk to its neighbour, which adds it to its own copy of that chunk. After 3 steps, each GPU holds one chunk that contains the sum from all four.
2. All-gather, 3 more steps: the finished chunks travel round the ring, overwriting the stale copies.

Each GPU sent 6 chunks of 2 numbers, 12 numbers, which is 2 × 3/4 × 8. The remarkable part: with 400 GPUs instead of 4, each would still send just under twice its gradient. The work per GPU barely grows.

flowchart LR
  G0["GPU 0<br/>chunks a b c d"] -- "one chunk per step" --> G1["GPU 1<br/>chunks a b c d"]
  G1 -- "one chunk per step" --> G2["GPU 2<br/>chunks a b c d"]
  G2 -- "one chunk per step" --> G3["GPU 3<br/>chunks a b c d"]
  G3 -- "one chunk per step" --> G0

N, S = 8, 14e9
in_machine, between_machines = 500e9, 50e9
# 2(N - 1)/N × S / B, in milliseconds
round(2 * (N - 1) / N * S / in_machine * 1000)  # → 49
round(2 * (N - 1) / N * S / between_machines * 1000)  # → 490
P, D, F = 7e9, 16_384, 1e15
# 6·P·D / F, in seconds
round(6 * P * D / F, 2)  # → 0.69

flowchart TB
  subgraph M1["machine 1: fast links, ~500 GB/s"]
    a1["GPU"] <--> a2["GPU"] <--> a3["GPU"] <--> a4["GPU"]
  end
  subgraph M2["machine 2: fast links, ~500 GB/s"]
    b1["GPU"] <--> b2["GPU"] <--> b3["GPU"] <--> b4["GPU"]
  end
  M1 <-- "network, ~50 GB/s per GPU:<br/>data parallelism, once per step" --> M2
  TP["tensor parallelism:<br/>talks inside every layer"] -.-> M1
  TP -.-> M2

In code: ring_all_reduce simulates the ring step by step and counts the numbers each GPU sends; all_reduce_seconds and training_step_seconds are the two formulas; IN_MACHINE_LINK and BETWEEN_MACHINES_LINK are the illustrative link speeds.

Why it matters in practice. At scale, a cluster's network matters as much as its chips. Parallelism strategies are chosen by matching how often each kind talks to how fast each link is, and a training run that ignores the map spends its budget waiting.

Everyday picture A bookshelf of fixed width. The encyclopedia (the model's weights) must go on it, whole. Whatever space is left holds one notebook per customer being served (their KV cache), and a notebook grows with every word of the conversation. Thinner volumes, printed in a smaller number format, leave room for more notebooks.

Tiny worked example A model shaped like Llama 3 70B (80 layers, 8 key/value heads of 128 dimensions each) on one 80 GB GPU. Its KV cache costs 2 × 80 × 8 × 128 × 2 = 327,680 bytes per token at 16-bit (primer.ml.inference derives this).

This ignores activations and the serving software's own working memory, which take several more gigabytes in practice.

P, M_gpu = 70e9, 80e9
L, H_kv, d_h = 80, 8, 128
def kv_per_token(b_kv):
    # a key and a value, per layer, per KV head
    return 2 * L * H_kv * d_h * b_kv
kv_per_token(2)  # → 327680
# weight GB at 16, 8 and 4 bits
[P * bits / 8 / 1e9 for bits in (16, 8, 4)]  # → [140.0, 70.0, 35.0]
# tokens of cache beside fp8 weights, with a 16-bit cache and then an fp8 one
int((M_gpu - P * 1) // kv_per_token(2)), int((M_gpu - P * 1) // kv_per_token(1))  # → (30517, 61035)

flowchart LR
  P["parameters × bytes per weight"] --> W["weights"]
  T["tokens × bytes per token"] --> K["KV cache"]
  W --> Q{"weights + cache<br/>fit in GPU memory?"}
  K --> Q
  Q -- yes --> Y["serve: spare room means<br/>more users or longer contexts"]
  Q -- no --> N["smaller formats, fewer KV heads,<br/>shorter contexts, or more GPUs"]

In code: max_cache_tokens subtracts the weights and divides what is left by the per-token cache, using primer.ml.inference.weight_bytes and primer.ml.inference.kv_cache_bytes_per_token; LLAMA3_70B_SHAPE holds the example model's shape.

Why it matters in practice. This arithmetic comes first in every deployment: the number format sets whether a model fits, how many users share a GPU, and (because decode reads every weight per token) how fast each of them sees words appear. primer.ml.inference carries it on into batching, speculative decoding and prompt caching.

Answer each one out loud or on paper before you open it. If you can explain it, you know it.