03 — The Arithmetic of Transformers (The Killer Notebook)

“Predict the number before you measure it.” — the ethos of the whole roadmap, applied at the layer where it produces the highest study signal.

This chapter is a notebook, not an essay. Every section ends with a formula, a worked example on a real model, and a prediction you should be able to make cold on a whiteboard. Build this notebook once, keep it forever. It is the artefact study partners will ask you to reproduce.

Target model for worked examples: Llama‑3‑70B (meta-llama/Meta-Llama-3-70B). All configuration numbers below are pulled from the published config.json; verify them yourself from the HF repo before you memorise anything.

hidden_size            = 8192        # d_model, aka embedding dim
num_hidden_layers      = 80          # L
num_attention_heads    = 64          # n_q_heads
num_key_value_heads    = 8           # n_kv_heads  (GQA 8:1)
head_dim               = 128         # d_head  = hidden/64
intermediate_size      = 28672       # d_mlp   (SwiGLU MLP hidden dim)
vocab_size             = 128256      # V       (untied output embedding for 70B)
rope_theta             = 500000.0    # base for RoPE
max_position_embeddings= 8192        # base; 131072 with rope_scaling on 3.1

Keep this block near your desk. Every arithmetic derivation below traces back to it.

---

1. Parameter Count from Config (derive it, don’t google it)

1.1 Per‑layer decomposition

For a modern Llama‑class decoder‑only block, ignoring the two tiny RMSNorm scale vectors (≈ 2·hidden, always negligible):

Attention block (Q, K, V, O projections):

P_attn = hidden · (n_q_heads · head_dim)    # W_Q
       + hidden · (n_kv_heads · head_dim)   # W_K
       + hidden · (n_kv_heads · head_dim)   # W_V
       + (n_q_heads · head_dim) · hidden    # W_O

For Llama‑3‑70B: d = 8192, n_q·head_dim = 8192, n_kv·head_dim = 1024.

W_Q = 8192 · 8192 =  67,108,864
W_K = 8192 · 1024 =   8,388,608
W_V = 8192 · 1024 =   8,388,608
W_O = 8192 · 8192 =  67,108,864
                     -----------
P_attn per layer  = 150,994,944  ≈ 150.99 M

MLP block (SwiGLU has three projections: gate, up, down):

P_mlp = 3 · hidden · d_mlp
      = 3 · 8192 · 28672
      = 704,643,072  ≈ 704.64 M

Per layer total: 150.99 M + 704.64 M ≈ 855.64 M.

1.2 Whole‑model total

P_layers     = 80 · 855.64 M  ≈ 68,451 M
P_embed_in   = V · hidden = 128256 · 8192 = 1,050,673,152 ≈ 1.051 B
P_embed_out  = V · hidden (untied for 70B) = 1.051 B
P_final_norm ≈ 0
                                        -----------
P_total       ≈ 68,451 + 1,051 + 1,051 = 70,553 M
                                       ≈ 70.55 B  ✓

Reported card says “70B”. You just derived it from six config numbers.

Contrast (tied embeddings): Llama‑3.2‑1B, Qwen2.5‑1.5B and SmolLM2‑1.7B tie input and output embeddings, so you subtract one copy of V · hidden. When your from‑scratch inference script’s param count is off by exactly V · hidden, this is why.

1.3 The one‑line back‑of‑envelope formula

For any decoder‑only model with SwiGLU MLP:

P ≈ L · ( 2·d·(n_q + n_kv)·d_head/n_q · d/d   +   3·d·d_mlp )
  ≈ L · ( 2·d² · (1 + n_kv/n_q)   +   3·d·d_mlp )
  +   V·d · (1 or 2, depending on tying)

If d_mlp ≈ 4d and GQA ratio is 1:g (so n_kv = n_q/g), everything collapses roughly to:

P ≈ 12 · L · d²  +  V · d · (1 or 2)

The famous “12·L·d²” scaling law. Verify: 12 · 80 · 8192² = 64.4 B, off by ~10% because Llama’s d_mlp = 3.5·d (not 4·d) and the GQA discount trims attention. Close enough for studies. Own the correction.

---

2. FLOPs per Token (Forward Pass)

2.1 The “2 × params” rule and why it holds

For every parameter w participating in a matmul, one output element requires one multiply and one add — so 2 FLOPs per parameter per token.

FLOPs_forward ≈ 2 · P  per token

For Llama‑3‑70B: 2 · 70.55 B = 141.1 GFLOPs per token.

2.2 Where the 2 breaks down (be honest about it)

The 2·P rule undercounts two contributions:

1. Attention scores. Computing QK^T and attn · V costs 4 · L · n_q · d_head · S per token in the decode case (S = current context length). At long context this is no longer negligible.

2. Softmax, LayerNorm, activations, embedding lookups. Small but nonzero.

Sharper formula (from Chinchilla / Kaplan lineage):

FLOPs_forward_per_token ≈ 2·P  +  4 · L · n_q · d_head · S
                        = 2·P  +  4 · L · d · S       (since n_q · d_head = d)

At S = 128k, L = 80, d = 8192: extra term = 4·80·8192·131072 = 344 GFLOPs. That is larger than the 2·P core cost of 141 GFLOPs. This is exactly why long context is expensive and why FlashAttention exists.

2.3 Training vs inference

Training forward+backward ≈ 6 · P FLOPs/token (forward 2, backward 4 because backward computes both weight‑grad and input‑grad). Kaplan and Chinchilla scaling laws use this constant. Inference is only forward → 2 · P.

Predict this before running: on an H100 SXM (989 TFLOPS bf16 dense), the compute time to prefill one token of Llama‑3‑70B at short context is 141 GFLOPs / 989 TFLOPS ≈ 143 μs. You will find your measured prefill is much larger, dominated by launch overhead and imperfect utilization — welcome to inference.

---

3. The KV Cache — Feel the Horror

This is the formula. Memorise the derivation, not the number.

3.1 Derivation

For every token you cache, you store one K vector and one V vector per layer, per KV head:

bytes_per_token = 2 · L · n_kv_heads · head_dim · bytes_per_element

The leading 2 is because you store both K and V.

3.2 Worked: Llama‑3‑70B, per token

                                     fp16 (2B)      fp8 (1B)      int4 (0.5B)
2 · 80 · 8 · 128 · bytes  =  2·80·8·128 = 163,840 elements per token
bytes/token               =    327,680 B     163,840 B      81,920 B
                          =        320 KB        160 KB          80 KB

Yes, 320 KB per token for Llama‑3‑70B in fp16. Sit with that.

3.3 At 128k context (Llama‑3.1‑70B with YaRN)

For a single sequence of 131,072 tokens:

fp16 KV cache =    320 KB · 131,072 =   41,943,040 KB =  40.0 GiB  (per sequence)
fp8  KV cache =    160 KB · 131,072 =   20,971,520 KB =  20.0 GiB
int4 KV cache =     80 KB · 131,072 =   10,485,760 KB =  10.0 GiB

Feel the horror: at batch=8 in fp16 you need 320 GiB just for KV. That is four full H100‑80GB cards’ worth of memory before you load a single weight. This is why:

• Long context is a memory‑capacity problem long before it is a compute problem.

• fp8 KV cache is now default in vLLM/SGLang for long‑context serving (halves the pain for ~0 quality loss).

• Kimi’s Mooncake, DeepSeek’s MLA, and the entire prefix‑caching literature exist.

• MLA (Multi‑head Latent Attention) is not a curiosity — DeepSeek‑V2’s MLA cuts KV per token by ~10× vs. GQA.

3.4 GQA’s KV savings, quantified

Llama‑3‑70B has n_q = 64, n_kv = 8. If it were pure MHA (n_kv = 64), KV per token would be:

2 · 80 · 64 · 128 · 2 =  2,621,440 B  = 2.5 MiB per token

At 128k: 320 GiB per sequence. GQA saved you 8× — turning “impossible” into “merely painful”. This is why every serious inference model since Llama‑2‑70B uses GQA or MLA.

3.5 Memory budget: predict max batch size

Deployment scenario: 1 × H100‑80GB, Llama‑3‑8B (32 layers, 8 KV heads, 128 head_dim, 4096 hidden), 8k context, fp16 weights + fp16 KV.

weights          = 2 · 8·10^9         = 16.0 GB
activations peak ≈ 2 GB                                          (rough allowance)
CUDA/PyTorch runtime overhead ≈ 2 GB
                                        -----------
usable for KV    = 80 - 16 - 2 - 2    = 60 GB

kv/token = 2·32·8·128·2  = 131,072 B ≈ 128 KB
kv per seq at 8k         = 128 · 8192 = 1.0 GB

max_batch ≈ 60 / 1.0     = ~60 sequences

Now go verify empirically with vLLM. If you land within ±20% you’ve internalised the arithmetic. If not, find the miscount (activations at prefill are the usual culprit — long prompts allocate transient buffers).

---

4. Decode Speed Prediction (Roofline for LLMs)

4.1 The formula

At batch = 1, decode is memory‑bound. Generating one token requires streaming every weight (and touching the KV cache) through HBM. So:

tokens_per_sec (batch=1) ≈  HBM_bandwidth / bytes_per_token_read
                         ≈  HBM_bandwidth / (P · bytes_per_weight)

Ignore the KV read term — it’s small compared to weights at moderate context.

4.2 Worked: Llama‑3‑8B on a 4090

weights fp16    = 16 GB    (16.06 GB precisely)
4090 HBM        = 1.008 TB/s (GDDR6X, but same math)
predicted tok/s = 1008 / 16.06 ≈ 63 tok/s

You will measure ~50–60 tok/s in llama.cpp with --flash-attn on. Predicted within 20%. Same model at Q4 (4.5 GB): 1008 / 4.5 ≈ 224 tok/s — this is why quantization is oxygen for local inference.

4.3 Worked: Llama‑3‑70B in fp16 on 2 × H100 SXM (TP=2)

Weights split, each H100 reads half:

per-GPU weight bytes = 131.4 / 2  = 65.7 GB
per-GPU HBM3         = 3.35 TB/s
predicted tok/s      = 3350 / 65.7 ≈ 51 tok/s

Real number with vLLM TP=2 is ~40–45 tok/s. Delta is TP all‑reduce overhead per layer (two AllReduces per transformer block on the critical path).

Same model on 1 × MI300X (192 GB, 5.3 TB/s HBM3e — the number that made AMD relevant to inference):

fits in one chip:  131.4 / 5300 → tok/s ≈ 40

MI300X trades peak FLOPs for capacity + bandwidth — exactly the tradeoff decode wants.

4.4 Batch‑N decode: when do you hit the ridge?

Bigger batch amortises weight reads across N sequences: bytes_moved_per_step stays the same, work_done scales linearly with N.

arithmetic_intensity(batch=N) ≈ 2 · N   (FLOPs/byte)

On H100 the ridge point is ~295 FLOPs/byte, so decode stays memory‑bound until roughly N ≈ 150. In practice you hit KV memory / interference limits far before that, which is why chunked prefill + PD disaggregation matter. But the reason batch pays for decode traces back to this single line.

4.5 Prefill is different

Prefill processes T prompt tokens in one big GEMM. Arithmetic intensity there is roughly T itself — for T = 2048, you are firmly compute‑bound. So a 70B on H100 SXM prefills at roughly:

peak_flops / (2·P) tokens/sec  =  989 TFLOPS / 141 GFLOPs/tok ≈ 7,000 tok/s (theoretical)

Real measured ~4–5k tok/s. Delta is imperfect tensor‑core utilization and attention (whose FLOPs grow as S, killing you at long prompts — hence chunked prefill).

---

5. Full Memory Model at Inference Time

GPU memory used  ≈  W                              # weights
                 +  2·L·n_kv·d_head·bytes · Σ seq_len_i   # KV, summed over active batch
                 +  activations_peak               # transient forward buffers, bigger at prefill
                 +  runtime_overhead               # CUDA context, PyTorch caching allocator, ~1-3 GB

Common failure modes to name from memory:

• OOM at prefill only: activation buffers for long prompts (attention S² intermediates when not using FA). Fix: FlashAttention, chunked prefill.

• OOM mid‑stream: KV growth from long‑running sessions. Fix: paged KV, preemption/eviction.

• Fragmentation OOM: contiguous KV allocation on variable lengths. Fix: PagedAttention (Phase 4).

• Slow steady degradation: PyTorch caching allocator holding freed blocks. torch.cuda.empty_cache() diagnostically.

---

6. Napkin Ledger — Numbers to Have Memorised

Pin these to memory; they are the study vocabulary.

Quantity	Value	Comment
H100 SXM peak bf16 dense	989 TFLOPS	fp8 is 2×; sparsity claims are marketing
H100 SXM HBM3 bandwidth	3.35 TB/s	ridge ≈ 295 FLOPs/byte
H200 HBM3e bandwidth	4.8 TB/s	same compute as H100, more BW → decode gain
A100 80 GB HBM2e	2.0 TB/s, 312 TFLOPS bf16	old workhorse
MI300X	5.3 TB/s, 192 GB, ~1.3 PFLOPS bf16	capacity king
RTX 4090	1.008 TB/s, 165 TFLOPS bf16	local‑LLM ridge ≈ 165
RTX 3090	0.936 TB/s, 71 TFLOPS bf16	kernel dev workhorse
NVLink 4 (Hopper)	900 GB/s per GPU	TP‑friendly
PCIe Gen5 x16	~63 GB/s each way	why TP hates PCIe
DDR5 dual channel	~60–100 GB/s	CPU inference ceiling
Fastest CPU inference	~5–15 tok/s @ 8B Q4	consistent with above

Also memorise: fp16 = 2 B, fp8 = 1 B, int4 = 0.5 B, bf16 = 2 B. And 1 GB ≈ 10⁹ B ≈ 2³⁰ B (the industry sloppiness is worth ~7%, be aware but don’t be pedantic in casual conversation).

---

7. Notebook Layout (build this file, keep it forever)

Turn this chapter into a Jupyter notebook with these cells:

Cell 1 : ModelConfig dataclass (matches your Phase 1 inference script)
Cell 2 : def param_count(cfg) -> int           # verify against real HF checkpoint
Cell 3 : def flops_per_token(cfg, seq_len)     # returns forward FLOPs
Cell 4 : def kv_cache_bytes(cfg, seq_len, bs, dtype)  # returns bytes
Cell 5 : def max_batch(cfg, gpu_mem_gb, seq_len, weight_bits, kv_bits)
Cell 6 : def predicted_decode_tps(cfg, hbm_bw_gbs, weight_bits)
Cell 7 : Roofline plot: ridge = peak_flops/bw ; scatter your kernels on it
Cell 8 : Table: {model} × {precision} × {gpu} → predicted vs. measured tok/s

Every time you touch a new model, run the notebook and add a row. Every time your prediction is >20% off, write the post‑mortem in the notebook (not on a sticky note). Over 12 months this becomes the single most valuable file in your repo.

---

Exit Test for This Chapter

You should, on a whiteboard, in under 15 minutes, given only the six numbers of a model’s config plus a GPU spec:

1. Derive the parameter count within ±5%.

2. Compute KV cache per token in fp16 exactly.

3. State whether a target (model, batch, context) fits on a given GPU with fp16 weights and fp16 KV, and if not, what the smallest viable quantization is.

4. Predict batch‑1 decode tok/s within ±25%.

5. Explain why prefill is compute‑bound and decode is memory‑bound in one sentence, with the arithmetic to back it.

If you can do all five, you have earned the right to open a CUDA file in Phase 2.