In Part 1 we watched the Word Factory answer a question. The Typewriter wrote the answer one token at a time, and every single token needed a full trip through the machine.
Today we climb inside and ride along with one token on that trip.
We'll count every bit of math it causes, room by room. By the end you'll be able to look at any language model and guess how much work one word costs, with a rule small enough to fit on a sticky note.
And I'll tell you about a bug this counting helped me catch in Part 1.
How do we count "work"?
Inside a language model, almost everything is one move repeated billions of times: multiply two numbers, then add the result to a running total.
Engineers count these tiny math steps as FLOPs (floating-point operations). One multiply is 1 FLOP. One add is 1 FLOP.
Here's a tiny example. A word is turned into a list of numbers (just 2 here, to keep it small), and we push it through a table of numbers with 2 rows and 3 columns:
word: [2, 3]
table: [ 1, 0, 4 ]
[ 5, 2, 1 ]
result: [ 2×1 + 3×5, 2×0 + 3×2, 2×4 + 3×1 ] = [17, 6, 11]
Look at what happened: every number in the table got used exactly once, multiplied by something and added to a total. The table has 6 numbers, so that's 6 multiply-and-adds.
Chips do a multiply-and-add as one combined move, and everyone counts it as 2 FLOPs. So:
work for one word through a table = 2 × (rows × columns) FLOPs
This table cost 2 × 6 = 12 FLOPs. (Strictly there were only 3 additions here, not 6, but counting 2 per table number keeps the bookkeeping simple, and for big tables the difference is tiny.)
Hold on to that rule. It's going to do almost all the work today.
A map of the inside
The Reader and the Typewriter from Part 1 are the same machine, so there's only one building to explore. Our model has the shape of GPT-2 small, a real model from 2019. Here's the building:
┌────────────────────────────────────────┐
│ SCOREBOARD: score all 50,257 words │ → next word
└────────────────────────────────────────┘
▲
┌────────────────────┬───────────────────┐
floor 12 │ look-back room │ thinking room │
└────────────────────┴───────────────────┘
▲
... 10 more floors ...
▲
┌────────────────────┬───────────────────┐
floor 1 │ look-back room │ thinking room │
└────────────────────┴───────────────────┘
▲
┌────────────────────────────────────────┐
│ DICTIONARY: word → 768 numbers │
└────────────────────────────────────────┘
▲
"sky"
Our token rides the elevator from the bottom to the top. Let's follow it.
Stop 1: The dictionary
First, the token needs a backpack of numbers. The model has a dictionary with one row per word piece: 50,257 rows, each holding 768 numbers. The token number for "sky" says which row to grab.
That's it. No multiplying, just finding the right row. Like looking up a word in an actual dictionary.
Work: 0 FLOPs. (Practically free.)
From now on, our token is those 768 numbers. Every room it visits changes them a little.
Stop 2: The look-back room (attention)
This is where the token looks back at the words that came before it. Engineers call it attention. It happens in four steps.
a) Make three cards. The token's 768 numbers are pushed through one big table (768 rows × 2,304 columns) to make three cards:
- a Question card: "what am I looking for?"
- a Label card: "what am I about?"
- a Contents card: "what can I share?"
b) Write in the notebook. The Label and Contents cards go into the notebook (the KV cache from Part 1). Every earlier token already left its own two cards there.
c) Read the notebook. The token holds up its Question card and compares it with every Label in the notebook, to see which earlier words matter most right now. Then it blends together their Contents cards: lots of the important ones, a little of the rest.
d) Tidy up. The blend goes through one more table (768 × 768) to become 768 numbers again, which get added to the token's backpack.
Here's the work, using our rule:
make 3 cards: 2 × 768 × 2,304 = 3.54 million FLOPs
tidy-up table: 2 × 768 × 768 = 1.18 million FLOPs
read the notebook: 4 × 768 × (number of notes in the notebook)
That last line is different. It doesn't depend on a fixed table. It depends on how long the conversation is so far. With 128 notes it's 0.39 million FLOPs. With 8,192 notes it's 25 million. Keep an eye on it.
Stop 3: The thinking room (MLP)
Next, the token thinks on its own for a bit. Engineers call this the MLP (multi-layer perceptron, a fancy name for "two tables with a squish in the middle").
a) Expand. The 768 numbers go through a table that makes them 4 times bigger: 3,072 numbers. More room to think.
b) Squish. Each of the 3,072 numbers goes through a little curve that shrinks small and negative values. (This is what lets the model learn things more interesting than straight lines.)
c) Shrink back. Another table squeezes the 3,072 numbers back down to 768, and they're added to the backpack.
expand: 2 × 768 × 3,072 = 4.72 million FLOPs
shrink: 2 × 3,072 × 768 = 4.72 million FLOPs
Notice something? The thinking room's two tables are twice as big as the look-back room's tables. So in every floor, thinking is about two-thirds of the table work.
Then the elevator goes up, and the token does it all again. 12 floors in total.
Stop 4: The scoreboard
At the top, the token's 768 numbers are compared with every word in the dictionary to give each one a score. The highest score becomes the next word in the answer.
That's one more table: 768 × 50,257.
scoreboard: 2 × 768 × 50,257 = 77.2 million FLOPs
That's a big one. Remember it.
Let's count it for real
Counting on paper is nice, but let's measure. I wrote forward_pass.py: the same GPT-2-small-shaped building in about 120 lines of NumPy. Like in Part 1, its numbers are random, so its words are gibberish, but it does exactly the same amount of math as a trained model of this shape.
The trick is one little helper. Every time the token goes through a table, it adds up the FLOPs and starts a stopwatch:
def mm(part, x, W):
"""Push the token's numbers x through table W, and keep score."""
t = time.perf_counter()
y = x @ W # the actual math
secs[part] += time.perf_counter() - t # how long it took
flops[part] += 2 * W.shape[0] * W.shape[1] # 2 x rows x columns
return y
And here's one floor of the building. You don't need to follow every line, just spot the two rooms:
for p in self.layers: # 12 floors
# 👀 Look-back room
h = norm(x)
q, k, v = np.split(mm("make Q, K, V", h, p["qkv"]), 3)
scores = compare(q, notebook_labels) # read the notebook...
mix = blend(scores, notebook_contents) # ...and blend
x = x + mm("attention output", mix, p["out"])
# 🧠 Thinking room
h = norm(x)
up = gelu(mm("expand 768 -> 3072", h, p["up"]))
x = x + mm("shrink 3072 -> 768", up, p["down"])
(The real file uses the same steps with a bit more bookkeeping. The notebook is pre-filled with notes from earlier tokens, so we can test different conversation lengths.)
I ran one token at a time through the building, 30 times, repeated the whole thing 3 times, and kept the middle run. Plain CPU, 2 cores, no GPU.
The results
Here's one token with 128 notes in the notebook (a short chat):
notebook holds 128 tokens | one token = 252.2 million FLOPs, 35.3 ms
part MFLOPs work % ms time %
1. look up the token 0.00 0.0% 0.01 0.0%
2. attention: make Q, K, V 42.47 16.8% 4.85 13.7%
3. attention: read the notebook 4.72 1.9% 1.95 5.5%
4. attention: output 14.16 5.6% 1.91 5.4%
5. thinking: expand 768 -> 3072 56.62 22.4% 6.26 17.7%
6. thinking: shrink 3072 -> 768 56.62 22.4% 6.39 18.1%
7. score every word in the dictionary 77.19 30.6% 8.29 23.5%
other (norms, softmax, adds) 0.45 0.2% 5.68 16.1%
("MFLOPs" means millions of FLOPs, added up over all 12 floors.)
And here's how it changes as the conversation gets longer:
notes in notebook 128 1,024 8,192
FLOPs per token 252.2 M 285.6 M 553.0 M
time per token 35.3 ms 44.7 ms 105.0 ms
notebook share of the work
1.9% 13.2% 54.6%
Five things we just discovered
1. It's (almost) all tables
Add up the table steps (2, 4, 5, 6 and 7): 247 of the 252 million FLOPs, 98%. Reading the notebook was 1.9%. Everything else (the squish curve, tidying, adding) was 0.2%.
So when people say "AI is just multiplying matrices", they're nearly right. The tables are the model. These tables are what engineers call parameters (or weights): the numbers the model learned during training.
2. The sticky-note rule
Since every table number gets used once per token, here's a shortcut. Count the numbers in all the tables:
12 floors × 7,077,888 numbers per floor = 84,934,656
scoreboard: 768 × 50,257 = 38,597,376
total = 123,532,032 (about 124 million)
Double it: 247.1 million FLOPs. The measured total was 252.2 million. The small difference is the notebook reading plus the tiny "other" steps.
That gives us the rule for the sticky note:
FLOPs per token ≈ 2 × parameters + 4 × floors × width × notes in the notebook
For a model with 8 billion parameters, that's about 2 × 8 = 16 billion FLOPs for every single word it writes (a little less, really, because the dictionary lookup is free), before you even count the notebook.
3. Long chats change where the work goes
With 128 notes, reading the notebook was a rounding error (1.9%). With 8,192 notes, it was more than half of all the work (54.6%) and two-thirds of the time (66%).
The tables don't change size during a chat. The notebook does. Every new word adds another note that every future word has to read. That's why really long conversations and long documents get slow and expensive, and why so much engineering goes into making the notebook smaller and faster to read.
4. The scoreboard is huge in small models
For our little model, scoring 50,257 words was 30.6% of all the work: the single biggest step! That's because the dictionary is big and the rest of the model is small.
In big models, it flips. I fed the shape of a popular 8-billion-parameter open model (Llama 3.1 8B) into a small calculator, flop_counter.py, that uses the same rule:
Llama 3.1 8B shape
notebook 128: 15.08 GFLOPs | attention tables 17.8% read the notebook 0.4% thinking tables 74.8% score the dictionary 7.0%
notebook 8,192: 19.30 GFLOPs | attention tables 13.9% read the notebook 22.2% thinking tables 58.4% score the dictionary 5.4%
notebook 32,768: 32.19 GFLOPs | attention tables 8.3% read the notebook 53.4% thinking tables 35.0% score the dictionary 3.3%
(GFLOPs means billions of FLOPs.) The scoreboard drops to about 7%, and the thinking room takes three-quarters of the work. And look at the last line: with a 32,768-token conversation, the notebook is again more than half the work, even for a big model.
5. Counting work is not the same as counting time
Look at the "other" row again: 0.2% of the work, but 16.1% of the time. Those little steps do almost no math, but each one is a separate small job, and every job has a setup cost. It's like a delivery driver whose packages weigh nothing but who still has to park at every house. Real inference engines glue small steps together (engineers call it kernel fusion) for a similar reason.
And here's a bigger surprise. Our CPU did 252 million FLOPs in 35 ms. That's about 7 billion FLOPs per second. Is that the CPU's top speed?
Let's check with bus_test.py. It pushes tokens through 36 tables of 768 × 3,072 numbers (340 MB, the same size as all 12 floors), first one token per trip, then many tokens per trip:
for batch in [1, 16, 64, 256]:
x = rng.standard_normal((batch, 768)).astype(np.float32)
...
for W in tables: # one trip through every table
y = x @ W
tables: 340 MB
1 tokens per trip: 17.4 ms per trip, 17.36 ms per token, 9.8 GFLOP/s
16 tokens per trip: 82.7 ms per trip, 5.17 ms per token, 32.9 GFLOP/s
64 tokens per trip: 106.9 ms per trip, 1.67 ms per token, 101.7 GFLOP/s
256 tokens per trip: 299.7 ms per trip, 1.17 ms per token, 145.1 GFLOP/s
Same tables. Same CPU. 15 times more math per second when 256 tokens share the trip.
It's the school bus from Part 1 again. With one token per trip, the CPU spends almost all its time hauling the 340 MB of tables from memory, and very little time doing math. With 256 tokens, it hauls the tables once and does 256 tokens' worth of math with them.
So the FLOP count tells you how much math a word needs. It does not tell you how long it takes. For typing one word at a time, the clock is mostly set by how many bytes have to be moved, not by how much math there is. That's the whole story of Part 3.
How counting FLOPs caught a bug in Part 1
When I first ran this post's code, the tables took more than 6 times longer inside the model than they did in the bus test. Same tables, same CPU. With FLOPs counted, the gap was impossible to miss.
The cause was one innocent-looking line:
scores = q @ K / np.sqrt(64) # np.sqrt(64) is a 64-bit NumPy number...
In recent versions of NumPy, dividing our 32-bit numbers by a 64-bit NumPy number quietly turns the result into 64-bit numbers. Every step after that worked with 64-bit numbers too, which meant converting every table on every token. Twice the bytes, plus extra copying.
The fix is tiny:
scores = q @ K / 64 ** 0.5 # a plain Python number keeps everything 32-bit
The same line was in Part 1's tiny engine. With it fixed, typing went from 107 ms to about 27 ms per token, and reading went from 6.5 ms to about 4.4 ms per token. So reading is about 6 times cheaper per token than typing, not 16 times like I first wrote. The lessons in Part 1 all still hold, but I've corrected the numbers there and added a note.
The code for this post now has a guard so it can't happen silently again:
assert x.dtype == np.float32, "something switched us to 64-bit math"
Lesson: if your measured speed is far from what you'd expect for the work, don't trust the measurement. Go find the reason.
Cheat sheet
| Stop | Real name | Work for one token | Grows with the chat? |
|---|---|---|---|
| Dictionary | Embedding lookup | ~0 | No |
| Make 3 cards | Q, K, V projection | 2 × width × 3 × width | No |
| Read the notebook | Attention over the KV cache | 4 × width × notes | Yes |
| Tidy up | Attention output projection | 2 × width × width | No |
| Thinking room | MLP | 2 × width × 4 × width × 2 tables | No |
| Scoreboard | LM head (unembedding) | 2 × width × dictionary size | No |
| Rule | What it says |
|---|---|
| Sticky note | FLOPs per token ≈ 2 × parameters + 4 × floors × width × notes |
| Small models | The scoreboard can be the biggest single step |
| Big models | The thinking room is most of the work |
| Long chats | Reading the notebook takes over |
| One token at a time | Time is set by moving bytes, not by doing math |
What's next
We just saw that the same CPU does 15 times more math per second when tokens share the trip. Part 3: Prefill vs Decode: Two Phases, Two Bottlenecks digs into exactly that: why reading your question is limited by math, why typing the answer is limited by memory, and what that means for every chatbot server.
Every number in this post comes from actually running these three scripts (a few snippets above are shortened for reading; the files are the exact code I ran): forward_pass.py, bus_test.py and flop_counter.py. The tiny model used random weights on a 2-core CPU, so it shows the pattern, not the speed of any real GPU.
Top comments (0)