Slides | TPU Deep Dive

Review deck

TPU Deep Dive

TPU lineage and roofline thinking
Anatomy of a TensorCore
ICI, OCS, slices and pods
Transformer accounting on TPU
Sharded matmuls and collectives on a torus
Choosing a mesh for training
Inference rooflines
The TPU serving stack
From Python to the MXU
Pallas, XProf and a design review
TPU and GPU architecture, side by side

Chapter 1. Day 1, first block

TPU lineage and roofline thinking

By the end you can:

explain each TPU generation, from v1 to 8t and 8i, as a response to a named bottleneck.
compute the math time, the communication time and the ridge point for any chip from its spec sheet, and state which regime a kernel is in.
derive that a bf16 matmul's arithmetic intensity is about its token batch, and give the compute-bound batch for v5e, Ironwood, H100 and TPU 8i.
show why int8 or fp8 weights with bf16 math halve the compute-bound batch, and why fp8 math doubles it again.
apply the same roofline to an interconnect link and derive when a sharded matmul is communication-bound.
name what the HBM roofline leaves out (VPU-bound ops, latency-bound transfers, achievable fraction of peak) and when each matters.

Ten generations, each a response to a bottleneck

Chip	Year	Answer to
v1	2015	inference cost: one 256×256 int8 array, DDR3-bound
v2, v3	2017–18	training: HBM, bf16, ICI 2D torus, 1,024 chips
v4	2020	scale and failures: OCS, 3D torus, SparseCore
v5e / v5p	2023	cost vs scale: 256-chip 2D vs 8,960-chip 3D
v6e	2024	256×256 MXU, 4.7× v5e
7x Ironwood	2025	inference at scale: FP8, 192 GiB, 9,216 chips
8t / 8i	2026	split: training torus vs inference SRAM + Boardfly

FLOPs outgrew bandwidth

50×

BF16 FLOP/s per chip, TPUv2 to Ironwood

10.5×

HBM bandwidth per chip over the same span

The ridge point rose from 66 to 313 FLOP/B.
Each generation needs more tokens per chip to keep the matrix units busy.
Google's retrospective: about 100× peak, 10× HBM [P-Micro26].

Two times bound every kernel

T_\text{math} = \frac{\text{FLOPs}}{C}, \qquad T_\text{comms} = \frac{\text{bytes}}{W}

Arithmetic over peak FLOP/s; bytes over the bandwidth of the link that carries them. Overlapped: the max. Not overlapped: the sum, at most 2× worse.

Compute-bound when FLOPs per byte exceeds the ridge point, $C / W_\text{hbm}$ .
On TPUs the compiler prefetches HBM into VMEM, so overlap is the default.

Ridge points

Chip	BF16	FP8	HBM BW	BF16 intensity
TPU v5e	197 TFLOP/s	—	819 GB/s	241
TPU v5p	459 TFLOP/s	459 TFLOP/s	2.77 TB/s	166
TPU v6e	918 TFLOP/s	920 TFLOP/s	1.64 TB/s	560
TPU7x Ironwood	2.31 PFLOP/s	4.61 PFLOP/s	7.38 TB/s	313
TPU 8t	996 TFLOP/sp	5.98 PFLOP/sp	6.53 TB/s	153p
TPU 8i	1.1 PFLOP/sp	8.81 PFLOP/sp	8.6 TB/s	128p
H100 SXM	989 TFLOP/s	1.98 PFLOP/s	3.35 TB/s	295
B200 (HGX)	2.25 PFLOP/s	4.5 PFLOP/s	8 TB/s	281

p provisional: from the JAX Pallas hardware table, not in Google's published spec.

A matmul's intensity is its token batch

\text{AI} = \frac{2BDF}{2BD + 2DF + 2BF} \approx B

Weight bytes dominate when B is small, and each weight byte is used once per token.

313

tokens per chip before a bf16 matmul on Ironwood is compute-bound (v5e 241, H100 295)

Tokens per replica, not sequences. Model sharding does not change it.

Worked: LLaMA-3-70B up-projection, batch 64, one Ironwood

$3.01 \times 10^{10}$ FLOPs: $T_\text{math} = 13.0\ \mu$ s
$4.75 \times 10^{8}$ bytes: $T_\text{hbm} = 64.3\ \mu$ s
Intensity 63, ridge 313
v5e: 579 µs. H100: 142 µs.

4.9×

HBM-bound. Compute-bound above 329 tokens for this shape.

Precision moves the ridge

B_\text{crit} \approx \frac{C_\text{math}}{W_\text{hbm}} \cdot \frac{\text{bytes per weight}}{2}

int8/fp8 weights, bf16 math: $B_\text{crit}$ halves. v5e 240 to 120; Ironwood example 329 to 164.
fp8 math too: peak doubles, ridge 625, $B_\text{crit}$ back to 329. Step time still halves: 64 to 32 µs.
TPU 8i fp8 ridge about 1,024 provisional. So 8i bought bandwidth and 384 MB of SRAM instead of FLOPs.

The same roofline on a link

\frac{BDF}{C} > \frac{2BF}{W_\text{link}} \iff D > \frac{2C}{W_\text{link}}

Split the contracting dimension over 2 chips only if each half's math outlasts shipping the partial sums.

8,755

v5e threshold on D

46,140

Ironwood threshold derived

What the HBM roofline leaves out

VPU work: its own roofline, critical intensity about 3 on v5p. Cheap only when fused.
Latency: about 1 µs per ICI hop. Under about 45 kB per hop on v5e, latency dominates.
Achievable peak: the book's estimate is 80–85% on H100/B200 and about 95% on TPU; the measurements are not published.

Chapter 2. Day 1, second block

Anatomy of a TensorCore

By the end you can:

draw a TensorCore and say what each unit does — scalar unit, VPU, MXU, XLU, VMEM, SMEM, vector registers and DMA engines — and which one issues instructions to the rest.
explain weight-stationary systolic dataflow and compute an MXU tile's utilization from M, K, N and the array size.
explain why a contraction or output dimension smaller than the array, or not a multiple of it, wastes MXU capacity, and estimate the waste.
explain why VMEM is a compiler-managed scratchpad rather than a cache, give its size per generation, and compute the ridge point for VMEM-resident weights.
explain (8,128) tiling and the bf16 and int8 packing, and what they imply for array shapes.
say how many devices JAX sees per chip on v4, v5e, v5p, v6e, Ironwood, 8t and 8i, and why.
say what SparseCores and 8i's Collectives Acceleration Engine are for.
map SM, warp scheduler, CUDA core, shared memory, Tensor Core and HBM onto TPU units, and name three places the mapping breaks.

A TensorCore is an SM, not a Tensor Core

132 vs 2

SMs on an H100 against TensorCores on an Ironwood chip

TensorCore = scalar unit + VPU + MXUs + XLU + VMEM.
NVIDIA Tensor Core = matrix unit inside an SM, the counterpart of the MXU.

Parts of a TensorCore

Ironwood figures, per TensorCore: two TensorCores per chip, each on its own chiplet. HBM and VMEM sizes from the JAX hardware table; VREG and VPU shapes from the Scaling Book. Earlier chips differ in counts (v5p: 4 MXUs of 128×128 per core), not in structure.

One thread of control

No warps and no hardware switching. The scalar unit issues VLIW bundles for every unit.
Latency is hidden by the schedule: the DMA for the next tile is already in flight.
XLA builds that schedule for the whole program, so every shape must be known at compile time.
The cost: a new shape is a new compile, which can take minutes for a large model.

The MXU: weight-stationary systolic array

U = \frac{M \cdot K \cdot N}{S^2 \cdot (M + K + N - 2)}

Utilization of one weight tile: M rows streamed, K × N of an S × S array filled.

256×256, full tile	M = 8	M = 256	M = 2,048
Utilization	1.5%	33%	80%

M is the token batch: decode leaves the MXU filling and draining.
K = 64 (a head dimension) caps it at 25%.

VMEM: a scratchpad, not a cache

Chip	VMEM per chip
v5e, v5p, v6e, 7x, 8t	128 MiB
8i	384 MiB
H100	30 MB SMEM + 50 MB L2

Filled by explicit DMA on a compile-time schedule.
About 7× HBM bandwidth on Ironwood (book estimate).
Weights resident in VMEM: ridge about 45 FLOP/B instead of 313.

Same matmul, HBM against VMEM

259

tokens to be compute-bound with int8 weights streamed from HBM (v5e)

tokens if the weights already sit in VMEM (at 22× HBM bandwidth)

This is the lever TPU 8i pulls with 384 MB of SRAM.

Layout is (8,128)

Vector registers are 8 sublanes × 128 lanes, so arrays are tiled 8×128.
bf16 is T(8,128)(2,1) and int8/fp8 are T(8,128)(4,1): 2 or 4 values per 32-bit sublane.
Last two dimensions should be multiples of 8 and 128. Pallas blocks must be.
Relayouts show up in profiles as copy or transpose fusions.

Devices per chip

Chip	TensorCores	JAX devices
v4, v5p	2	1 (megacore)
v5e, v6e	1	1
Ironwood	2 chiplets	2
8t / 8i	1 / 2	not documented

Ironwood's die-to-die link is 6× one ICI link: keep a chip's two chiplets on the busiest mesh axis.

Beyond the MXU

SparseCore (v4 onward): gather/scatter engines for embeddings, about 5% of die, 5–7× on embedding-heavy models. Pallas-programmable.
CAE (8i): Collectives Acceleration Engine, up to 5× lower on-chip collective latency. The JAX table still lists a SparseCore on 8i.
Trainium2: 8 NeuronCores, each a 128×128 tensor engine plus vector, scalar and GPSIMD engines, SBUF + PSUM.

Where the GPU mapping breaks

A TensorCore is not a big SM: one instruction stream, no thread switching.
A VPU lane is not a warp: contiguous tiles and masking, no per-thread addressing.
Registers are inverted: about 32 MB on an H100 against about 256 KB per TPU core. A TPU spills to VMEM instead.

Chapter 3. Day 2, first block

ICI, OCS, slices and pods

By the end you can:

describe the TPU topologies — a 2D torus on v5e and v6e, OCS-joined 4×4×4 cubes forming a 3D torus on v4, v5p and Ironwood, and Boardfly on 8i — and say which slice shapes have wraparound.
convert a Cloud-docs ICI figure (per chip, both directions, all links) into per-link one-way bandwidth and into the per-axis bandwidth used in collective cost formulas.
explain what optical circuit switches buy (reconfiguration, failure isolation, non-contiguous and twisted slices) and what they cost.
rank HBM, ICI, PCIe and DCN for Ironwood with numbers, and say which parallelism belongs on each.
compute hop counts, diameter and point-to-point transfer time on a slice, and say when a transfer is latency-bound rather than bandwidth-bound.
explain why TPU 8i abandons the torus for Boardfly, using hop latency and decode traffic.
compare TPU scale-up with NVLink domains and InfiniBand fat trees, and with Trainium's move from torus to switch.

No switch in the data path

Each chip has direct ICI links to its nearest neighbours: 4 on v5e and v6e, 6 on v4, v5p and Ironwood.
Messages to distant chips hop through intermediate chips, at about 1 µs per hop.
Per-chip bandwidth stays constant as the pod grows: about 9,000 chips at full ICI bandwidth.

Reading the bandwidth numbers

Chip	Docs, per chip both ways	Per link, one way	$W_\text{ici}$ per axis
v5e	400 GB/s / 4 links	50 GB/s (book 45)	100 GB/s
v5p, Ironwood	1,200 GB/s / 6 links	100 GB/s (book 90)	200 GB/s
8i	2,400 GB/s	not published	—

Divide by links, then by two. $W_\text{ici}$ counts the two links of one axis.

Topologies

v5e, v6e: 16×16 2D torus, 256 chips. Wraps only on 16-long axes.
v4, v5p, Ironwood: 4×4×4 cubes joined by optical switches into a 3D torus. Wraps on every axis of a whole number of cubes.

Twisted tori: offset wraparound, +70% bisection on 4×4×8.
8i Boardfly: 4-chip boards, 8-board groups, 36 groups. At most 7 hops across 1,152 chips.

Hops and wraparound

Slice	Wraparound	Diameter (hops)
v5e 16×16	both axes	16
v5e 4×4	none	6
Ironwood 4×4×4	all axes	6
Ironwood 2×2×2	none	3
Ironwood 8×16×16	all axes	20

Ring of N: farthest N/2 hops with wraparound, N − 1 without.
No wraparound: ring collectives cannot close, about 2× the communication time.

Optical circuit switches

MEMS mirrors steer light between fibres; reconfigure in milliseconds.
No packets read, no buffers, no in-network reduction.

Slices of any shape from non-adjacent cubes.
Route around failed racks: 99.98% availability on v4.
Under 5% of system cost.

Ironwood bandwidth ladder, per chip, one way

Link	GB/s	HBM ÷ link
HBM	7,380	1
ICI, all 6 links	600	12
ICI, one link	100	74
PCIe (v6e figure)	32	231
DCN	12.5	590

TP and FSDP on ICI. Data parallelism across DCN. MaxText: dcn_tensor_parallelism "never recommended".

Latency or bandwidth

T \approx \text{hops} \cdot 1\,\mu\text{s} + \frac{\text{bytes}}{\text{links} \cdot W_\text{link}}

192 µs

16.8 MB corner to corner on a v5e 4×4: bandwidth-bound

16 hops

across an 8×8×16 slice: a small decode collective is latency-bound

Against GPU and Trainium

	Scale-up domain	Hops	In-network reduction
NVL72	72 GPUs, 1.8 TB/s each	1	yes (SHARP)
Ironwood	9,216 chips, 1.2 TB/s each	20 across an 8×16×16 slice	no
8i	1,152 chips, 2.4 TB/s each	up to 7	CAE on chip
Trn3	64 or 144 chips	switched	—

GPU and Trainium converge on switched domains of 72 to 144. TPU keeps a torus for training and adds Boardfly for inference.

Pod accounting

Pod	Chips	Hosts	BF16 FLOP/s	HBM
v5e	256	32	5.0e16	4 TiB
v5p	8,960	2,240	4.1e18	840 TiB
Ironwood	9,216	2,304	2.1e19	1,728 TiB

Ironwood FP8: 42.5 EFLOP/s per pod. Google's "1.77 PB" is 9,216 × 192 GB in decimal units; the docs say 192 GiB.

Chapter 4. Day 2, second block

Transformer accounting on TPU

By the end you can:

count the FLOPs of any einsum, and derive 6 × parameters × tokens for training from 2× forward and 4× backward.
count the parameters of a dense GQA transformer and of an MoE, split into MLP, attention and vocabulary.
compute KV-cache bytes per token and per sequence, and say how many sequences fit on a chip after the weights.
derive when attention FLOPs dominate (T > 8D) and why decode attention is never compute-bound.
compute the compute-bound batch for an MoE layer and explain why it is E/k times larger than for a dense layer.
estimate training memory (10 bytes per parameter plus activations) and show that large clusters are needed for FLOPs, not memory.
name the TPU-specific consequences of model shape — padding to 256, head dimension, KV heads — and where they show up later.

Three numbers every later chapter divides

Parameters: what HBM must hold and what decode reads every step.
FLOPs: 6 × parameters × tokens to train; 2 × active parameters per token to infer.
KV-cache bytes: $2 L K H$ × bytes per token, per sequence, read every decode step.

Counting

Component	Parameters per layer	Training FLOPs per layer
Gated MLP	$3DF$	$18BTDF$
Attention projections	$2DH(N+K)$	$12BTDH(N+K)$
Dot-product attention	none	$12BTSNH$
Vocabulary	$2DV$ total	$6BTDV$

LLaMA-3-70B: MLP 56.4B, attention 12.1B, vocab 2.1B. 6.3e24 FLOPs on 15T tokens: 45 days on a v5p pod, 9 on an Ironwood pod at 40% MFU.

Worked: LLaMA-3-70B

Quantity	Value
Parameters	70.5B (MLP 80%)
Training FLOPs, 15T tokens	6.3e24
Forward FLOPs per token	about 141 GFLOP
KV per token, bf16	320 KB
Weights, bf16	141 GB: fits on one Ironwood chip (206 GB), 65 GB left for cache

Attention: small share, wrong intensity

\frac{\text{attn FLOPs}}{\text{matmul FLOPs}} = \frac{T}{8D}

Attention dominates only beyond 65,536 tokens for D = 8,192.

\text{AI}_\text{attn} = \frac{TSG}{TG + S}

Prefill: grows with T. Decode: about G ≈ 8, at any batch.

Decode attention is memory-bound at every batch size: each sequence reads its own cache.

The KV cache sets decode batch

320 KB

per token, LLaMA-3-70B in bf16

10.7 GB

per 32k-token conversation

8 Ironwood chips: 141 GB of weights, about 1.5 TB free.
About 140 conversations at 32k fit.
The matmuls want 313 tokens per chip, 2,500 across 8 chips.
Capacity, not FLOPs, limits the batch.

MoE: memory follows E, compute follows k

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

Each expert sees only its share of tokens, so the batch must be E/k times larger.

	E/k	Compute-bound batch
DeepSeek-V3 shape, v5e, int8 weights	32	3,840
Same, Ironwood, fp8 weights, bf16 math	32	about 5,000
Same, Ironwood, fp8 math too	32	about 10,000

Training memory is not the constraint

Adam: 10 bytes per parameter. 700 GB for 70B.
LLaMA-3-70B at a 4M-token batch: 11.2 TB with activations. Fits on 54 Ironwood chips.
On 54 chips it would take years. Over a pod it is 1.3 GB per chip.

Shape on TPU

Pad D, F, H and V to multiples of 256 (v6e, Ironwood). LLaMA-3 already is.
Head dim 128 fills half a 256-deep MXU per head; 64 fills a quarter.
K = 8 KV heads: split the cache by heads at most 8 ways, then by batch.
Ragged MoE batches fight static shapes: capacity padding, ragged_dot, megablox, SparseCore.

Chapter 5. Day 3, first block

Sharded matmuls and collectives on a torus

By the end you can:

read and write sharding notation, such as A[I_X, J_Y] and the unreduced marker {U_X}, and the equivalent JAX PartitionSpec. I can say what each device holds for a given mesh and spec.
classify any sharded matmul into the four cases and name the collective each needs.
derive the bidirectional-ring AllGather time V / W_ici and explain why it does not depend on the number of chips. I can extend it to several axes, to AllReduce (2×) and to AllToAll (¼).
say when a collective is latency-bound, and compute the threshold on v5e and Ironwood.
explain collective matmul (decomposing an AllGather into collective-permutes overlapped with partial matmuls) and compare it with async tensor parallelism on GPUs.
relate the book's formulas to NCCL's algbw and busbw, and compare how quickly TPU and GPU collectives reach peak bandwidth.

Sharding notation

Notation	PartitionSpec	Meaning
$A[I_X, J_Y]$	`P('X', 'Y')`	I over X, J over Y
$A[I_{XY}, J]$	`P(('X','Y'), None)`	I over X and Y together
$A[I, J]\{U_X\}$	`unreduced={'X'}`	partial sums still to add over X

One mesh axis per array dimension at most; unmentioned axes replicate.

Four cases for A[I, J] · B[J, K]

Case	Sharding	Collective
1	J not sharded on either	none
2	J sharded on one operand	AllGather it
3	J sharded on both, same axis	AllReduce, or ReduceScatter at half cost
4	I and K sharded on the same axis	AllGather one first

Tensor parallelism's second matmul is case 3. FSDP's weights are case 2.

Ring collective costs

T_\text{AG} = T_\text{RS} = \frac{V}{W_\text{ici}}, \quad T_\text{AR} = \frac{2V}{W_\text{ici}}, \quad T_\text{A2A} = \frac{V}{4W_\text{ici}}

V per device; independent of ring size. Divide by the number of axes used. Floor: N/2 µs of hop latency.

1 GB per device, Ironwood	AG	AR	A2A	AG over 3 axes
Time	5 ms	10 ms	1.25 ms	1.7 ms

Latency-bound below one microsecond of link

45 kB

per hop on v5e

100 kB

per hop on Ironwood

Below this, time is hops × 1 µs. Decode collectives live here.
Without wraparound the ring is open: a 4-long v5e axis took 560–680 µs against a 373 µs formula.

Collective matmul

Split the AllGather into N − 1 collective-permutes; multiply the shard you have while the next one arrives.
XLA: windowed einsum, above xla_jf_spmd_threshold_for_windowed_einsum_mib.
Book, shard_map by hand: 311 µs → 244 µs, against 224 µs with no communication.
GPU equivalents: PyTorch async-TP (SymmetricMemory), TE userbuffers, Flux. Same algorithm (Wang et al.), more machinery, because NCCL competes for SMs.

ICI roofline: Ironwood is about 5× stricter

Chip	$C/W_\text{ici}$ per axis
v5p (book)	2,550
Ironwood (docs)	11,535

Moving weights (AllGather W) is compute-bound when B > $C/W_\text{ici}$ .
Moving activations beats moving weights at large B only when D > $2C/W_\text{ici}$ : 23,000 on Ironwood.

Against NCCL

Same ring math: busbw factors $(n-1)/n$ and $2(n-1)/n$ .
H100 node AllReduce: about 370 GB/s only at about 10 GB messages; a 58 MB tensor-parallel AllReduce gets about 150 GB/s.
TPU v5e AllGather: about 95% of peak at about 10 MB.
AllGather of bf16[1024, 16384] over 8 devices: H100 node 65 µs; one Ironwood axis 168 µs; three axes about 56 µs.

Chapter 6. Day 3, second block

Choosing a mesh for training

By the end you can:

derive the condition under which each training scheme is compute-bound — data parallelism, FSDP, tensor parallelism, FSDP plus TP, and cross-slice data parallelism over DCN.
compute the optimal FSDP degree, X_opt, and pick a mesh for a given model, batch and slice.
redo the Scaling Book's LLaMA-3-70B plan on a v5p pod, and explain why the same batch is communication-bound on a full Ironwood pod.
explain why pipeline parallelism is central to GPU recipes and secondary on TPUs, citing LLaMA-3's actual GPU configuration.
map these choices to MaxText configuration — ici_* against dcn_* axes and logical axis rules.

One number decides the mesh

\alpha = \frac{C}{W_\text{ici}}

FLOPs a chip must do per byte moved on one torus axis.

Chip	α
v5p	2,550 (book), 2,295 (docs)
Ironwood	11,535
TPU 8t	about 2,490 (provisional)
H100, across nodes	2,475

Conditions for compute-bound training

Scheme	Condition	v5p	Ironwood
DP / FSDP	$B/X > \alpha/M_X$	850 (3 axes)	3,845
TP	$Y < M_Y F / \alpha$	11-way	2.5-way
FSDP + TP	$B/N > \alpha^2/(M_X M_Y F)$	113	2,320
DP over DCN	$B/\text{slice} > C / W_\text{dcn}$	73,440	185,000

For LLaMA-3-70B ( $F$ = 28,672), tokens per chip or per slice.

LLaMA-3-70B, 4M-token batch

v5p pod, 8,960 chips

468 tokens per chip, threshold 113.
$X_\text{opt}$ = 1,619: about 2,048 FSDP × 4 TP.
44 days at 40% MFU.

Ironwood pod, 9,216 chips

455 tokens per chip, threshold 2,320.
Communication-bound by about 2×.
Compute-bound up to about 1,800 chips; about 7,200 at a 16M batch.

Why Ironwood is different

5×

v5p's BF16 FLOPs per chip

1×

v5p's ICI bandwidth per link

Larger batches, fewer chips per slice plus DCN data parallelism, pipeline stages, overlap.
TPU 8t doubles ICI: α back to about 2,490 with the provisional BF16 figure.

GPU recipe for comparison

H100 α: 2,200 in node, 2,475 across nodes. Similar thresholds; different topology.
NVLink stops at 8 GPUs, so TP stays in the node and PP + DP carry the scale.
LLaMA-3 405B: 16,384 H100s, TP 8 × PP 16 × DP 128, 16M tokens, 41% MFU.
DeepSeek-V3: 2,048 H800s, EP 64 × PP 16 × DP 2.

MaxText encodes the chapter

ici_fsdp_parallelism: -1: FSDP fills the slice.
dcn_data_parallelism: -1: data parallelism across slices.
dcn_tensor_parallelism: "never recommended".
logical_axis_rules: named dimensions (embed, mlp, heads) to mesh axes.

Worked: LLaMA-3-405B on 8 v5p pods (derived)

Quantity	Value
Parameters / FLOPs	406B / 3.8e25 (15.6T tokens)
Time at 40% MFU	about 33 days
Per pod, 16M batch	2.1M tokens (DCN needs 73,440)
Per chip	234 tokens (FSDP+TP needs 61)
Mesh	about 1,000 FSDP × 8 TP per pod, 8 DP across pods

Chapter 7. Day 4, first block

Inference rooflines

By the end you can:

state the decode step-time lower bound and the general step-time formula, and compute both for a model on a slice.
compute the compute-bound decode batch for each weight and math precision on v5e, Ironwood, 8i and H100, and explain why KV growth makes it unreachable at long context.
explain why prefill is compute-bound and decode is not, and what that implies for disaggregated serving and the prefill-to-decode server ratio.
choose a decode sharding — model parallelism up to F/(Bβ), the KV cache split by heads then by batch, never FSDP — and say when its collectives are latency-bound.
explain int8, fp8 and fp4 weights, activations and KV on TPU, say which generations have native fp8, and describe how Qwix applies post-training quantization.
explain speculative decoding as spending idle FLOPs in a memory-bound step.

Two workloads in one server

Prefill

Whole prompt in one pass; batch is T tokens.
Compute-bound. Sets time to first token.
8k tokens of 70B: 0.91 s on 16 v5e, 157 ms on 8 Ironwood.

Decode

One token per sequence per step.
Reads every weight and every cache.
Memory-bound. Sets time per token.

Decode step time

T_\text{step} = \frac{B \cdot \text{KV}}{W_\text{hbm}} + \max\left(\frac{2B \cdot \text{params}}{C}, \frac{\text{param bytes}}{W_\text{hbm}}\right)

Attention always reads every cache; the MLP is weights or FLOPs, whichever is longer.

B_\text{crit} \approx \frac{C_\text{math}}{W_\text{hbm}} \cdot \frac{\text{bytes per weight}}{2}

	bf16	fp8 weights, bf16 math	fp8 throughout
Ironwood	313	156	313
8i (provisional)	128	64	512

LLaMA-3-70B, 8 Ironwood chips, 32k context

Weights, KV	Batch	Step	KV / weights / FLOPs	tok/s/chip
bf16	64	14.0 ms	11.6 / 2.4 / 0.5	570
bf16	128	25.7 ms	23.3 / 2.4 / 1.0	623
fp8	64	7.0 ms	5.8 / 1.2 / 0.5	1,140

KV-bound at every batch. Fewer bytes doubles throughput; more FLOPs does nothing.

Sharding decode

No FSDP: weights over ICI every step, an order-of-magnitude regression.
No in-server DP, no sequence parallelism: use replicas.
Model parallelism, far: ICI binds only when $Y > F/(B\beta_\text{ici})$ ; $\beta_\text{ici}$ ≈ 8 on v5e, 37 on Ironwood.
KV cache: by heads (at most K ways), then by batch with two AllToAlls per layer.
Collectives are latency-bound: 8-way below 360 kB (v5e), 800 kB (Ironwood).

Disaggregated serving

Prefill servers: compute, little memory, batch 1.
Decode servers: bandwidth and HBM for many caches, heavy model parallelism.
Cost: shipping KV caches between them.

3 : 1

prefill to decode servers for 8k prompts and 512-token answers (book example)

Quantization and other levers

Weights to 8 bits: halves the weight read and capacity; B_crit halves if math stays bf16.
KV to fp8: halves the dominant long-context term.
Native fp8 math: Ironwood and TPU 8; v5p and v6e list fp8 at the BF16 rate.
Qwix: intercepts dot_general, einsum and ragged_dot; regex rules per module.
Speculative decoding: k tokens scored for the cost of one in a memory-bound step.
MoE: E/k × the batch, latency-bound AllToAlls.

Chapter 8. Day 4, second block

The TPU serving stack

By the end you can:

describe how TPU serving got from JetStream through PyTorch/XLA and tpu-inference to the October 2026 split, and say which path to use for a new PyTorch or JAX port.
explain what torchax does to a torch.nn.Linear call, why tpu-inference needs the model as a pure function that returns its KV cache, and where the code ends up as HLO.
explain how continuous batching survives a compiler that wants static shapes — padded shape buckets compiled ahead, a flattened token batch, and Ragged Paged Attention's three sub-kernels.
read the tpu-inference KV-cache layout and explain its packing dimension from (8,128) tiling.
map tpu-inference's mesh axes to the decode sharding rules of Chapter 7.
size a deployment from the Scaling Book's LLaMA-3-70B serving analysis, and solve the book's 405B problem under a latency limit.

Three changes in 18 months

When	Stack	Status
2024–25	JetStream	archived 2026-02-01
to 2025	vLLM on PyTorch/XLA	deprecated; repo archived 2026-11-15
Oct 2025	tpu-inference: JAX models + PyTorch via torchax	3.6× / 2.1× over PyTorch/XLA
Oct 2026	vLLM JAX for JAX models; vllm-torchtpu for PyTorch	current

New PyTorch port: vllm-torchtpu. Existing torchax ports keep working.

Three paths, one scheduler

As of 2026-10-23. tpu-inference is being renamed vLLM JAX and focuses on JAX models; its torchax path still exists. New PyTorch-native TPU serving is in vllm-torchtpu. SGLang has the same split (sglang-jax, sglang-torchtpu).

torchax

torchax.Tensor is a torch.Tensor holding a jax.Array.
__torch_dispatch__ maps each ATen op to jnp / lax; jax.jit traces it to StableHLO.
The model must be a pure function: model_fn(weights, kv_caches, inputs) -> (kv_caches, ...).
The "sandwich": PyTorch model code inside; JAX sampler, KV cache and RPA kernel around it.

Continuous batching with static shapes

Flatten every query token in the step into one array; metadata carries per-sequence structure.
Pad token and sequence counts to buckets; compile them all before serving.
Ragged Paged Attention handles true lengths inside the kernel: fused KV write, three sub-kernels (decode, prefill, mixed).
86% bandwidth in decode, 73% FLOPs in prefill (Llama-3-8B, Ironwood).

KV cache layout

(\text{blocks}, \text{block size}, \lceil 2K/p \rceil, p, \text{head dim}), \qquad p = \frac{32}{\text{bits}}

K and V interleaved, packed 2 (bf16) or 4 (8-bit) per 32-bit sublane: the (8,128)(2,1) and (4,1) tilings.

Mesh: ('data', 'model') by default; heads, KV heads, MLP and experts on model. Also expert, attn_dp, dcp, pcp.

Case study: LLaMA-3-70B on v5e

int8 KV: 164 kB per token. Batch 32 at 8k: 113 GB, fits a 4×2.
4×2: 17 ms per step, about 235 tokens/s per chip. 4×4: 8.6 ms, same per chip.
bf16 on 4×8 at batch 64: HBM 18 µs, ICI 12 µs, math 5 µs. HBM-bound.
Bigger slices fit disproportionately more KV: about 1.6× tokens/s per chip from doubling.

LLaMA-3-405B under 15 ms per token (derived)

Slice	Weight-read floor	Limit
36 v5e	13.8 ms	about 17 sequences
64 v5e	7.7 ms	ICI caps batch near 100: about 12 ms, 130 tok/s/chip
256 v5e	1.9 ms	per-chip throughput levels off at about 149
8 Ironwood, fp8	6.9 ms	—

Answer: replicas of 64 chips.

Chapter 9. Day 5, first block

From Python to the MXU

By the end you can:

trace a function from Python through jaxpr, StableHLO, HLO and XLA's passes to LLO and TPU instructions, and dump each stage.
say what I control on TPU — shardings, constraints, layouts, remat, donation, flags and Pallas — and what I do not.
explain Shardy against GSPMD, compare JAX's three sharding modes (auto, explicit, manual shard_map), and say when to drop to each.
compare XLA whole-program compilation with CUDA Graphs and torch.compile with Inductor.
explain TorchTPU in detail — the PrivateUse1 backend, deferred StableHLO ops, its three eager modes, torch.compile without Inductor, the compile cache, tpu_dist and the SPMD handshake, and Pallas through jax_op.
contrast TorchTPU with JAX, torchax and PyTorch/XLA, and the TPU stack with AWS's Neuron stack.

Python to MXU

jaxpr: JAX's traced program, global shapes.
StableHLO: platform-independent MLIR, with sdy.sharding annotations.
HLO and XLA passes: partitioning, fusion, layout, HBM↔VMEM placement, scheduling.
LLO: programs DMAs, systolic pushes and units directly. Closed, in libtpu.
Machine code in instruction memory.

Pallas enters at stage 3 as a custom call, compiled by Mosaic.

One sharded MLP, compiled (real output, CPU, 2×2 mesh)

gelu(x @ w1) @ w2 with w1 P(None,'model'), w2 P('model',None): per-device shapes and one all-reduce

ENTRY %main.0_spmd (param: bf16[4,512], param.1: bf16[512,1024], param.2: bf16[1024,512]) -> bf16[4,512] {
  %param = bf16[4,512]{1,0} parameter(0), sharding={devices=[2,1,2]<=[4] last_tile_dim_replicate}, metadata={op_name="x"}
  %param.1 = bf16[512,1024]{1,0} parameter(1), sharding={devices=[1,2,2]<=[2,2]T(1,0) last_tile_dim_replicate}, metadata={op_name="w1"}
  %param.2 = bf16[1024,512]{1,0} parameter(2), sharding={devices=[2,1,2]<=[2,2]T(1,0) last_tile_dim_replicate}, metadata={op_name="w2"}
  %dot = f32[4,1024]{1,0} dot(%wrapped_convert, %wrapped_convert.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(mlp)/dot_general" stack_frame_id=2}
  %dot.1 = f32[4,512]{1,0} dot(%convert_convert_fusion, %wrapped_convert.2), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(mlp)/dot_general" stack_frame_id=3}
  %all-reduce.1 = f32[4,512]{1,0} all-reduce(%dot.1), channel_id=1, replica_groups=[2,2]<=[4], use_global_device_ids=true, to_apply=%add.clone_promoted, metadata={op_name="jit(mlp)/dot_general" stack_frame_id=3}
  ROOT %wrapped_convert.3 = bf16[4,512]{1,0} fusion(%all-reduce.1), kind=kLoop, calls=%wrapped_convert_computation.3, metadata={op_name="jit(mlp)/dot_general" stack_frame_id=3}

What you control

You control

Shardings and with_sharding_constraint
Layouts (AUTO)
Remat policy, buffer donation
XLA flags via LIBTPU_INIT_ARGS
Pallas kernels

You do not

TPU code generation
VLIW scheduling
Fusion beyond flags
VMEM inside fusions

Three sharding modes

Mode	View	Collectives
Auto	global	compiler inserts; steer with constraints ("tickling", about 60% of the work)
Explicit	global; `float32[8@X,512@Y]`	compiler inserts; ambiguous ops must give `out_sharding`
Manual, `jax.shard_map`	per device	you write psum, all_gather, ppermute

Shardy has been the default propagator since JAX 0.7.0.

Not CUDA Graphs

CUDA Graph: replay compiled kernels, less launch overhead, no cross-kernel optimization.
Inductor: fuses per captured graph.
XLA: one program per step. Fusion, every layout, VMEM placement and collective timing chosen together.
Price: static shapes. A new shape means a compile, not a slower kernel.

TorchTPU

PrivateUse1 backend, torch.device("tpu"), plain tensors.
Each ATen op becomes a deferred StableHLO builder; views and in-place ops are emulated.
Modes: strict eager (per op), debug, fused eager (per region), torch.compile (FX graph to XLA, no Inductor).
3-tier compile cache; PjRT and libtpu; async dispatch.
tpu_dist: DDP, FSDP2, DTensor; collectives are StableHLO, no NCCL.
SPMD handshake: reconciles rank-dependent PyTorch code with XLA's SPMD collectives.
Pallas via torch.tpu.pallas.jax_op. Qwen3-MoE-235B on 8,192 v7 chips at 27.1% MFU.

Four ways PyTorch reaches a TPU

	Tensor	Capture	To XLA
PyTorch/XLA (deprecated)	lazy `xla`	until `mark_step`	HLO
torchax	subclass over `jax.Array`	`jax.jit`	StableHLO via jaxpr
TorchTPU eager	plain, `PrivateUse1`	per op or region	StableHLO
TorchTPU `torch.compile`	plain	Dynamo FX	StableHLO per graph

AWS did the same: TorchNeuron (PrivateUse1), compiled by neuronx-cc instead of XLA.

Chapter 10. Day 5, second block

Pallas, XProf and a design review

By the end you can:

write and explain a Pallas TPU kernel — grid, BlockSpec index maps, the (8,128) rule, memory spaces, implicit double buffering, dimension_semantics and scratch accumulators — and say how it reaches the hardware through Mosaic.
choose a Pallas tile size from arithmetic intensity and the VMEM budget.
say when a Pallas kernel is worth writing on TPU and when XLA already does the job.
capture a TPU profile from JAX and from tpu-inference, and read the Trace Viewer, including async -start/-done pairs.
decode an HLO op line — shape, layout, tiling, memory space, fusion kind, called computation and replica groups.
find a communication-overlap failure and a sharding mistake in a profile, and use the op profile and roofline tools to choose what to fix.
run a design review — model, chip, SLO, mesh, quantization, batch, expected step time — and state the bottleneck in one sentence.

Pallas TPU in one kernel

pl.pallas_call(
    matmul_kernel,
    grid=(m // bm, n // bn, k // bk),          # sequential, reduction axis last
    in_specs=[pl.BlockSpec((bm, bk), lambda i, j, kk: (i, kk)),
              pl.BlockSpec((bk, bn), lambda i, j, kk: (kk, j))],
    out_specs=pl.BlockSpec((bm, bn), lambda i, j, kk: (i, j)),
    scratch_shapes=[pltpu.VMEM((bm, bn), jnp.float32)],     # fp32 accumulator
    compiler_params=pltpu.CompilerParams(
        dimension_semantics=("parallel", "parallel", "arbitrary")))

BlockSpecs drive double-buffered HBM→VMEM DMAs. Last two block dims: multiples of 8 and 128.
No thread blocks: the grid is a loop on one TensorCore.

Tile size is arithmetic intensity

Tile	FLOP/B	VMEM	Ironwood (ridge 313)
256²	128	1.5 MiB	memory-bound
512²	256	4 MiB	memory-bound
1024²	512	12 MiB	compute-bound
2048²	1024	40 MiB	compute-bound

Write Pallas for access patterns XLA cannot express: paged attention (scalar prefetch), grouped MoE, quantized matmuls, collective matmul (remote DMA).

Capture and read a trace

with jax.profiler.trace(dir): ... then xprof --logdir dir.
tpu-inference: PHASED_PROFILING_DIR, USE_JAX_PROFILER_SERVER.
Read the XLA Ops row of TPU:0. Layers repeat.
Synchronous collective: a long bar alone, the core idle.
Overlapped: all-gather-start … compute … all-gather-done.

Reading an op

%fusion.3 = bf16[32,32,4096]{2,1,0:T(8,128)(2,1)S(1)} fusion(...), kind=kCustom, calls=%all-reduce-scatter.3

{2,1,0}: minor-to-major. T(8,128)(2,1): tiles, bf16 pairs packed.
S(1): VMEM. No S: HBM.
calls=%all-reduce-scatter: a ReduceScatter. One matmul per fusion at most.
replica_groups={{0,16,...,112},...}: 8 per group, 8-way model parallel.

Tools

Tool	Answers
Op Profile	which ops cost the time
Roofline Model	which resource each op is bound by (HBM, VMEM lines)
Memory Viewer	peak HBM and what is live
Graph Viewer	HLO graph and shardings
Megascale Viewer	DCN traffic across slices

Book's check: 95.6 ms predicted, 96 ms measured.

Design review: LLaMA-3-70B on Ironwood, 30 ms/token at 32k

Replica	Batch	Step	tok/s/chip
4 chips, bf16	63	27.7 ms	569
4 chips, fp8 W + KV	140	27.9 ms	1,257
8 chips, fp8 W + KV	294	27.9 ms	1,316

KV-read-bound. Choose 4-chip fp8 replicas; prefill on separate 4-chip servers (337 ms TTFT), about 3 per decode replica.

Chapter 11. Synthesis

TPU and GPU architecture, side by side

By the end you can:

state the design bet behind GPUs, TPUs and Trainium, and what each gives up.
compare current chips quantitatively — compute per precision, HBM capacity and bandwidth, ridge point and on-chip memory — for Ironwood, TPU 8t and 8i, H100, B200, GB200, Rubin, Trainium2 and Trainium3.
compare scale-up and scale-out designs, and derive their consequences for AllGather and AllToAll and for training thresholds.
compare the programming models — kernels, compilation, PyTorch paths and partitioning.
say, for a given workload, which design has the structural advantage, and why.
name the 2025–26 convergence trends and the evidence for each.

Three bets

	Cores per chip	Latency hiding	On-chip memory	Compilation
GPU (H100)	132 SMs	hardware warps	30 MB SMEM + 50 MB L2 cache	eager kernels, graphs
TPU (Ironwood)	2 TensorCores	compiler schedule	128 MiB VMEM, no cache	whole-program XLA
Trainium2	8 NeuronCores	concurrent engines	224 MiB SBUF + PSUM	neuronx-cc

Compute and memory

Chip	BF16	FP8	HBM	HBM BW	BF16 intensity
TPU7x Ironwood	2.31 PFLOP/s	4.61 PFLOP/s	192 GiB	7.38 TB/s	313
TPU 8t	996 TFLOP/sp	5.98 PFLOP/sp	216 GiB	6.53 TB/s	153p
TPU 8i	1.1 PFLOP/sp	8.81 PFLOP/sp	288 GiB	8.6 TB/s	128p
H100 SXM	989 TFLOP/s	1.98 PFLOP/s	75 GiB	3.35 TB/s	295
GB200 (NVL72)	2.5 PFLOP/s	5 PFLOP/s	173 GiB	8 TB/s	313
Rubin (VR NVL72)	4 PFLOP/s	17.5 PFLOP/sp	268 GiB	19.2 TB/s	208
Trainium3	671 TFLOP/s	2.52 PFLOP/s	144 GiB	4.9 TB/s	137

p provisional: from the JAX Pallas hardware table, not in Google's published spec.

Flagship ridge points cluster near 300; inference designs push it down.

Switch against torus

Switch (NVLink, NeuronSwitch)

One hop to every peer in 8, 72 or 144.
In-switch reduction (SHARP).
Cheap all-to-all and small collectives.
Beyond the domain: scale-out at about 50 GB/s per GPU (H100).

Torus (ICI)

6 neighbours, multiple hops.
No in-network reduction.
Per-chip bandwidth constant to about 9,000 chips.
Bisection grows as $N^{2/3}$ .

1 GB collectives, idealized

	8	64	512 chips
AllToAll, Ironwood	1.25 ms	0.078 ms	0.020 ms
AllToAll, H100	0.28 ms	0.31 ms	0.039 ms
AllToAll, GB200 NVL72	0.14 ms	0.017 ms	0.039 ms
AllGather, Ironwood	5.0 ms	1.7 ms	1.7 ms
AllGather, GB200 NVL72	1.0 ms	1.1 ms	1.1 ms

The switch wins inside its domain; the torus wins past it. Practical GPU reach is lower than these figures.

Where each wins

Workload	Advantage
Large dense pretraining	TPU pod (v5p-class α); GPU close
MoE training	NVL72 inside 72; torus at multi-rack scale
Low-latency decode	NVL72 one hop; TPU 8i at most 7 hops
Long-context decode	most HBM bandwidth and capacity per dollar
Recommendation (embeddings)	TPU SparseCore
Changing ops, dynamic shapes	GPU

Converging, not the same

GPU Tensor Cores grow (warp → SM → 2 SMs) and gain Tensor Memory: toward the MXU and VMEM.
Trainium moved from torus to switch, and toward NVLink Fusion.
TPU split: 8t torus for training, 8i Boardfly and 384 MB SRAM for inference.
PyTorch moved to native backends on TPU and Trainium.
Unchanged: 100+ SMs and eager execution against 2 cores and whole-program compilation.

1 / 0Arrow keys to move, F for full screen.