Transformer accounting on TPU

Every later chapter divides one of three numbers by a bandwidth or a FLOP rate: parameters, FLOPs and KV-cache bytes. This chapter counts them for a decoder-only transformer, then applies the Chapter 1 roofline to attention and to mixture-of-experts layers. Most of it is hardware-neutral, so it moves quickly through what a GPU engineer already knows and spends its time on the parts where TPU shapes and numbers differ.

Notation follows the Scaling Book [SB4]: $B$ batch, $T$ query length, $S$ key/value length, $L$ layers, $D$ model width, $F$ MLP width, $N$ query heads, $K$ KV heads, $H$ head dimension, $V$ vocabulary, $E$ experts, $k$ experts per token, $G = N/K$ query heads per KV head.

Counting FLOPs

A contraction costs 2 FLOPs, one multiply and one add, per element of the product of all distinct dimensions, counting batch and contracting dimensions once [SB4]. $A[N, P] \cdot B[P, M]$ costs $2NPM$ .

The backward pass of $C = AB$ needs two products of the same size: $\partial L / \partial A = (\partial L / \partial C) B^\top$ and $\partial L / \partial B = A^\top (\partial L / \partial C)$ . Training is therefore $2 + 4 = 6$ FLOPs per multiply-accumulate in the forward pass. Since every parameter takes part in one multiply-accumulate per token, training costs

\text{FLOPs}_\text{train} \approx 6 \cdot \text{parameters} \cdot \text{tokens} \tag{4.1}

Two FLOPs per parameter per token forward, four backward. Attention's sequence-dependent term is extra and matters only at long context.

Rematerialization, recomputing activations in the backward pass instead of storing them, adds a second forward pass for the recomputed part. Saving only each layer's input brings the cost to about $8 \cdot \text{parameters} \cdot \text{tokens}$ [SB4].

Counting parameters

For a gated MLP and grouped-query attention, per layer [SB4]:

Component	Parameters per layer	Training FLOPs per layer
MLP (3 matrices)	$3DF$	$18BTDF$
Attention projections	$2DH(N + K)$	$12BTDH(N + K)$
Dot-product attention	none	$12BTSNH$
Norms	$2D$	$O(BTD)$
Vocabulary (embed + unembed)	$2DV$ in total	$6BTDV$ for the unembedding

ModelLayers Ld_model Dd_ff F (per expert)Query heads NKV heads KHead dim HVocab VExperts EActive experts kWeight bytesKV bytesContext S 32,768Chip

Total parameters8.03 B

Active per token8.03 B

MLP / attention / vocab5.64 / 1.34 / 1.05 B

Training FLOPs per token (6 × active)48.2 GFLOP

Forward FLOPs per token at S33.2 GFLOP

Attention share of forward FLOPs at S52%

KV per token131 kB

KV per sequence at S4.29 GB

Chips just to hold the weights1

16.1 GB of weights. On TPU7x Ironwood, 1 chip holds them. Each 32,768-token sequence needs 4.29 GB of KV cache, so the HBM left on that chip holds about 44 sequences.

Worked example: LLaMA-3-70B

$L = 80$ , $D = 8{,}192$ , $F = 28{,}672$ , $N = 64$ , $K = 8$ , $H = 128$ , $V = 128{,}256$ [SB6].

MLP: $3 \cdot 8192 \cdot 28672 \cdot 80 = 56.4$ B, 80% of the model.
Attention: $80 \cdot 2 \cdot 8192 \cdot 128 \cdot (64 + 8) = 12.1$ B.
Vocabulary: $2 \cdot 128256 \cdot 8192 = 2.1$ B.
Total 70.5B; the book gives 70.4B [SB6].

Training on 15T tokens takes $6 \cdot 70.5 \times 10^9 \cdot 15 \times 10^{12} = 6.3 \times 10^{24}$ FLOPs. At 40% MFU, that is 45 days on one v5p pod of 8,960 chips, the book's figure of about 44, or 9 days on one Ironwood pod of 9,216 chips derived. Computed by scripts/ch04_transformer.py.

How much of the compute is attention

For self-attention, $S = T$ , with a non-gated MLP of $F = 4D$ and $NH = D$ , the ratio of attention FLOPs to everything else is [SB4]

\frac{\text{attention FLOPs}}{\text{matmul FLOPs}} = \frac{12BT^2NH}{18BTDF + 24BTDNH} = \frac{T}{8D} \tag{4.2}

Attention overtakes the weight matmuls only when the sequence is about eight times longer than the model width.

For $D = 8{,}192$ that is $T > 65{,}536$ . For training at typical context lengths, attention is a minority of the FLOPs. It still matters, for two reasons that this formula hides:

Its intensity is different from a matmul's.
At decode time it does not shrink with batch.

Attention intensity: prefill against decode

Count the bytes attention reads. Ignoring the projections, one layer of GQA attention reads $Q$ and writes its output ( $4BTNH$ bytes in bf16), reads $K$ and $V$ ( $4BSKH$ ), and does $4BTSNH$ FLOPs [SB4, Q4]:

\text{AI}_\text{attn} = \frac{4BTSNH}{4BTNH + 4BSKH} = \frac{TSG}{TG + S} \tag{4.3}

B cancels: batching more sequences does not raise attention's intensity, because each sequence brings its own keys and values.

Prefill, $S = T$ : intensity is $TG/(G + 1)$ , which grows with the prompt. A long prompt's attention is compute-bound.
Decode, $T = 1$ : intensity is $SG/(G + S) \approx G$ . With 8 query heads per KV head that is about 8 FLOP/B, against a ridge of 313 on Ironwood. Decode attention is memory-bound at every batch size and every context length. The only levers are fewer KV bytes (GQA, quantized KV) and more bandwidth.

This is the structure from the "per-example weights" problem in Chapter 1: every sequence's KV cache is used by exactly one query. The MLP, by contrast, shares one weight matrix across the whole batch and does become compute-bound at large batch. Chapter 7 builds the decode step time from these two parts.

The KV cache

Each layer stores one key and one value vector per KV head per token:

\text{KV bytes per token} = 2 \cdot L \cdot K \cdot H \cdot \text{bytes per value} \tag{4.4}

Two (K and V) times layers times KV heads times head dimension.

For LLaMA-3-70B in bf16: $2 \cdot 80 \cdot 8 \cdot 128 \cdot 2 = 327{,}680$ bytes, about 320 KB per token. One 32k-token conversation holds 10.7 GB of cache, and one 128k-token conversation holds 43 GB. Without GQA, with $K = N = 64$ , those figures would be 8× larger [SB4].

Worked example: How many long conversations fit

LLaMA-3-70B in bf16 on 8 Ironwood chips. The weights take 141 GB, 17.6 GB per chip, leaving about 189 GB per chip of the 206 GB. At 10.7 GB per 32k-token sequence, each chip holds the cache for about 17 such sequences, or about 140 across the 8 chips derived. That is a batch of 140 at 32k context, below the 313 tokens per chip an Ironwood matmul needs to be compute-bound, before any working memory is reserved.

KV capacity, not compute, limits decode batch at long context. Ironwood's 192 GiB of HBM and 8i's 288 GB exist for this reason [D-7x] [G-8dd].

Mixture of experts

An MoE layer replaces one MLP with $E$ expert MLPs, of which $k$ run per token. Parameters scale with $E$ ; FLOPs per token scale with $k$ [SB4]. The Scaling Book's chapter 7 model, made into an MoE with $E = 16$ and $k = 2$ , grows from 18.4B to 212B parameters while its active parameters per token grow only to 31B [SB7, Q5].

The roofline penalty follows directly. A batch of $B$ tokens spreads over $E$ experts, so each expert's weight matrix sees about $Bk/E$ tokens. For each expert matmul to be compute-bound, $Bk/E$ must exceed the dense threshold:

B > \frac{E}{k} \cdot \frac{C_\text{math}}{W_\text{hbm}} \cdot \frac{\text{bytes per weight}}{2} \tag{4.5}

An MoE needs E/k times the batch of a dense model to keep its matrix units busy, because each expert sees only its share of the tokens.

With DeepSeek-V3's sparsity ( $E = 256$ , $k = 8$ , so $E/k = 32$ ) and int8 weights on v5e, $B > 3{,}840$ tokens [SB4, Q8]. On Ironwood with fp8 weights and bf16 math, $B > 5{,}000$ . With fp8 math as well, the ridge doubles and $B > 10{,}000$ derived. MoE decode is memory-bound at any realistic batch. MoE also adds two AllToAll collectives per layer, to route tokens to experts and back, which Chapter 5 costs.

Training memory

Mixed-precision Adam keeps bf16 parameters plus fp32 first and second moments: 10 bytes per parameter, 700 GB for a 70B model [SB5]. Activations add more. The Scaling Book's LLaMA-3-70B example, a 4M-token batch with two checkpoints per layer, needs about 10.5 TB of activations, 11.2 TB in total [SB6]. That fits on about 110 v5p chips or 54 Ironwood chips derived. Training on that few chips would take about 10 years.

Large training clusters exist for FLOPs, not memory. Spread over a full pod, the memory is about 1.3 GB per chip. Chapters 5 and 6 are about the communication cost of spreading the FLOPs, not about fitting the model.

TPU-specific consequences of shape

Pad to the array. D, F, the head dimension and the vocabulary should be multiples of 256 on v6e and Ironwood. LLaMA-3's are: $D = 8{,}192$ , $F = 28{,}672 = 112 \cdot 256$ and $V = 128{,}256 = 501 \cdot 256$ .
Head dimension sets attention's MXU efficiency. $H = 128$ fills half of a 256-deep array per head; $H = 64$ fills a quarter (Chapter 2). Attention kernels compensate by packing heads.
KV heads limit KV sharding. With $K = 8$ , the KV cache can be split by heads at most 8 ways. Beyond that it must be split by batch or by sequence, at the cost of extra collectives (Chapter 7).
Static shapes. Sequence length and batch are compile-time constants. Serving engines pad them into buckets (Chapter 8).

Worked problems from the book

Worked example: Counting a 4096-wide model

The Scaling Book's Chapter 4, problem 1 [SB4]: $D = 4{,}096$ , $F = 4D$ , $V = 32{,}000$ , $L = 64$ , multi-head attention with $NH = D$ , int8 KV cache.

Parameters: $L(3DF + 4D^2 + 2D) + 2DV = 17.4$ B. The book states "about 16B"; the exact sum of its own formula is 17.4B.
Attention fraction of parameters: $4D^2 / (4D^2 + 12D^2) = 1/4$ .
KV per token: $2 \cdot 64 \cdot 4096$ bytes $= 512$ KiB in int8.

Worked example: DeepSeek-V3's training efficiency

The Scaling Book's Chapter 4, problem 7 [SB4]. DeepSeek-V3 reported 2.79M H800-hours, 14.8T tokens and 37B active parameters, trained in FP8 at 1.513 PFLOP/s dense per GPU.

Needed: $6 \cdot 37 \times 10^9 \cdot 14.8 \times 10^{12} = 3.3 \times 10^{24}$ FLOPs.
Available: $1.513 \times 10^{15} \cdot 2.79 \times 10^6 \cdot 3600 = 1.52 \times 10^{25}$ .
Utilization: about 21.6%; the book rounds to 21.7%.