TPU and GPU architecture, side by side

This chapter puts the previous ten side by side. Each section takes one axis — compute, memory, network, software and workload fit — and lines up NVIDIA, Google and AWS on it, with numbers. Earlier chapters' callouts made these comparisons locally; here they form one argument.

Three design bets

GPU: many independent cores, hardware-managed latency hiding. An H100 has 132 SMs, each with 4 Tensor Cores, a large register file and up to 64 resident warps. A hardware L2 cache sits behind them. Independent kernels run on different SMs, and the warp scheduler hides memory latency by switching threads [N-wp] [N-tune, §1.4]. The Scaling Book's summary: GPUs "just work" on new tasks more often and lean far less on the compiler. The price is that their performance is harder to reason about and harder to bring to the roofline [SB12].

TPU: a few large cores, compiler-managed everything. An Ironwood chip has two TensorCores, each a single instruction stream driving two 256×256 systolic arrays. 64 MiB of VMEM per core is filled by DMAs scheduled at compile time, and XLA compiles the whole step. The Scaling Book's summary: TPUs "can often get closer to peak roofline performance with less work", and cost less to build. The price is the compiler: static shapes, whole-program compilation, and poor tolerance for irregular access patterns.

Trainium: in between. A Trainium2 chip has 8 NeuronCores. Each has a 128×128 systolic tensor engine plus separate vector, scalar and general-purpose SIMD engines running concurrently, a 28 MiB SBUF scratchpad in 128 partitions, and a dedicated 2 MiB accumulation buffer, PSUM [A-trn2-nki]. It is systolic and scratchpad-based like a TPU, and multi-core like a GPU, though with 8 cores instead of 132.

Google's own framing, in the PyTorch blog, is the same: "a TPU TensorCore is fundamentally a single-threaded execution processor … On the contrary, a GPU includes many identical compute cores" [T-pt].

Compute and memory, per chip

Sources are in scripts/specs.py. NVIDIA per-GPU figures are dense, halved from sparse, and per-GPU values for rack products are derived from rack totals [N-gb200] [N-rubin]. Trainium3's FP8 is MXFP8 [A-trn3]. TPU 8 BF16 and FP8 figures are provisional [J-hw]. Rubin's figures are NVIDIA's product-page claims; CES coverage quoted higher memory bandwidth.

What the table says:

• Ridge points cluster around 300 for this generation's flagships in BF16: Ironwood 313, GB200 312, H100 295, B200 281. The batch-size intuitions of Chapters 1 and 7 transfer directly between them.
• Inference-oriented designs push the ridge down. TPU 8i (128, provisional), Trainium3 (137) and Rubin (208) spend more on bandwidth relative to BF16 compute, and put their FLOPs into low precision.
• HBM capacity per chip: Rubin and 8i at 288 GB, Ironwood at 192 GiB, B200 and GB200 at about 180–186 GB, Trainium3 at 144 GiB.

Sources: [N-tune, §1.4], [SB12], [J-hw], [A-trn2].

Network: switch against torus

Sources: [SB12], [N-gb200], [SB2], [G-8dd], [A-trn2-sys], [A-trn3-sys].

Beyond the scale-up domain, GPUs and Trainium use an InfiniBand or Ethernet fat tree. An 8-GPU H100 node has 400 GB/s of egress, 50 GB/s per GPU. A GB200 NVL72 rack has 3.6 TB/s [SB12]. TPUs use DCN between pods: 12.5 GB/s per Ironwood chip, and up to 4× that on 8t.

What each design makes cheap.

• A switch makes all-to-all traffic cheap, and makes small, latency-bound collectives fast: one hop, plus SHARP for reductions. An AllToAll inside an NVLink node costs about half what it does on a TPU ring [SB12].
• A torus makes scale cheap. The pod adds links in proportion to chips, so per-chip bandwidth stays constant up to about 9,000 chips with no switch silicon. Ring collectives along an axis get full bandwidth at 10 MB messages.

Practical reach differs from these formulas. The Scaling Book reports that an 8×H100 AllReduce peaks around 370 GB/s of 450 only near 10 GB messages, and that multi-node AllToAll may not reach 50% of theory [SB12]. A v5e AllGather reaches about 95% of peak at about 10 MB [SB3].

Training rooflines, side by side

The same derivation as Chapter 6, with each platform's bandwidth [SB12]:

GPU recipes follow from the 8-GPU boundary: tensor parallelism inside a node, then pipeline, expert and data parallelism across nodes. LLaMA-3 used TP 8 × PP 16 × DP 128; DeepSeek-V3 used EP 64 × PP 16 × DP 2 (Chapter 6). TPU recipes follow from the pod: FSDP and TP inside one ICI domain, data parallelism across pods. The Ironwood row is the outlier, discussed in Chapter 6.

Software and programming model

Sources: [N-graphs], [T-pt], [A-torch], [A-nki], Chapters 9–10.

Where each design has the structural advantage

Price changes these comparisons, and it is not in any spec sheet. In February 2025, the Scaling Book's Google Cloud prices gave v5e 5.8 × 10¹⁷ FLOPs per dollar, against 3.3 × 10¹⁷ for an H100 [SB8].

Convergence, 2025–26

• GPUs are adding TPU features.
  • Blackwell added 256 KB of Tensor Memory per SM, because accumulators no longer fit in registers [SB12].
  • Tensor Cores have grown each generation: fed by one warp on Ampere, a whole SM on Hopper and two SMs on Blackwell [SB12].
  • Both move toward larger matrix units with a dedicated accumulator store fed by bulk copies (TMA), closer to an MXU fed from VMEM than to warps in registers.
• Trainium moved toward the GPU. It replaced Trn2's torus with Trn3's NeuronSwitch [A-trn3-sys], and AWS has previewed NVLink Fusion for Trainium4 [A-trn4].
• TPU split by workload. 8t keeps the training torus; 8i adopts a low-diameter network and a large SRAM [G-8dd].
• PyTorch moved off torch-xla on both TPU and Trainium, to native PrivateUse1 backends (Chapter 9).

The evidence against full convergence: GPUs still have over 100 independent SMs, a hardware cache and eager execution. TPUs still have two cores per chip, a compiler-managed scratchpad and whole-program compilation. The design bets of the first section have not changed; the boundaries between them have moved.

Worked problems from the book