A GPU achieves generality through programmability: thousands of scalar cores (CUDA cores)
that can execute any program, augmented by Tensor Cores that execute one specific operation
(matrix multiply) faster. The CUDA runtime decides what to run on each core at execution time.
A TPU achieves efficiency through specialization: a systolic array that performs matrix
multiply through data movement alone — no instruction dispatch, no SIMT scheduling,
no warp overhead. The XLA compiler decides everything ahead of time. The result is
dramatically higher energy efficiency and utilization for matrix multiply, and dramatically
lower flexibility for everything else. Understanding this trade-off is the key to
understanding where TPUs beat GPUs, where they cannot, and why the world still buys NVIDIA.
Part 1 — The Systolic Array: Data Flows, Computation Happens
The most elegant compute architecture for matrix multiply ever built
A systolic array is a grid of simple processing elements (PEs). Data flows through the grid
in a regular, synchronized wave — like blood pulsing through arteries (hence "systolic").
No complex control logic. No instruction cache. No branch predictor. Each PE does exactly
one thing: receive A from the left, receive B from the top, compute A×B and add to its
running accumulator, pass A to the right and B downward. That's it.
When data finishes flowing through all PEs, you have computed the matrix product.
4×4 systolic array — computing C = A × B, step by step
The systolic array computes C = A × B entirely through data movement.
No program runs on the PEs — they are hardwired to do MAC (multiply-accumulate) and
pass values to neighbors. This eliminates the entire GPU software layer:
no CUDA thread scheduler, no warp dispatcher, no register file management.
In exchange, the array can only compute matrix multiply — nothing else.
The TPU adds a separate Vector Processing Unit (VPU) for operations like softmax,
layer normalization, and activation functions that the systolic array cannot handle.
Why the systolic array achieves near-100% utilization on matrix multiply
A GPU Tensor Core achieves 85–95% utilization for GEMM only when cuBLAS optimally tiles
the matrix into shared memory and schedules warps to hide HBM latency. This requires
engineering expertise and GPU-generation-specific tuning (the CUDA library moat from Lesson 6).
A systolic array achieves near-100% PE utilization automatically for any matrix multiply
that fits in the array's tile size — because there is no scheduling overhead.
Every PE is doing a MAC every clock cycle, as long as data is flowing.
The only waste is pipeline fill/drain at the edges of each tile (minor, ~0.1% for large matrices).
Part 2 — GPU vs. TPU: A Side-by-Side Architectural Comparison
Two philosophies for the same problem
Dimension
NVIDIA H100 GPU
Google TPU v4
Compute paradigm
SIMT — 132 SMs, each with 128 CUDA cores + 4 Tensor Cores. Warp-level parallelism. Any compute pattern supported.
Systolic array (MXU) — 128×128 grid of PEs for matrix multiply. VPU for vector ops. Only matrix multiply is first-class.
Utilization for GEMM
85–95% with cuBLAS (requires expert tuning). Naive CUDA: 10–20%.
~97–99% automatically for aligned matrix sizes. No tuning required. XLA compiler handles alignment.
Utilization for irregular ops
High — CUDA cores handle any pattern. Custom kernels allow full flexibility.
Low — irregular memory access, sparse ops, custom activation functions poorly supported. XLA fusion helps but has limits.
XLA (Accelerated Linear Algebra) compiler. Developer writes JAX/TensorFlow; XLA generates optimized kernel code. No manual kernel writing.
Compile model
JIT (just-in-time) — kernels compiled at runtime. Fast first-run, flexible for dynamic shapes.
AOT (ahead-of-time) — entire computation graph compiled before execution. Longer compile time (minutes for large models), but execution is fully optimized and static.
Interconnect
NVLink 4.0 (900 GB/s per GPU bidirectional). NVSwitch for 72-GPU pools.
ICI (Inter-Chip Interconnect) — 3D torus topology. TPU v4: 600 GB/s per chip ICI bandwidth. No NVSwitch needed — torus routing.
TPUs are not a GPU replacement. They are a dedicated matrix multiply engine with a
software stack designed to compile entire model graphs to hardware-specific code.
The question is not "which is faster" — TPUs win on GEMM efficiency; GPUs win on flexibility.
The real question: "can your workload be expressed entirely as a static, compiled computation
graph of matrix multiplies?" If yes, TPU is more efficient. If your model has dynamic shapes,
custom ops, or non-standard architectures, XLA compilation either fails or degrades performance.
Google's own inference for Gemini and Search runs on TPU because these are stable, known
workloads compiled once and run billions of times. Research models (novel architectures,
experimental kernels) still run on GPU first.
Part 3 — XLA: The Compiler That Makes TPUs Usable
Why you write JAX, not systolic array assembly
The systolic array does matrix multiply through data movement. But a Transformer is not
just matrix multiply — it includes softmax (a normalization operation), ReLU/SiLU/GeLU
(activation functions), layer normalization, and residual additions. And above the
hardware, you write Python using JAX or TensorFlow — not systolic array instructions.
XLA (eXtended Linear Algebra) bridges this gap.
XLA compilation pipeline — from Python to TPU instructions
JAX / TensorFlow model code (Python)
Developer writes: jnp.dot(Q, K.T) / jnp.sqrt(d_k)
↓
HLO (High-Level Operations) Graph
Hardware-independent IR. Each Python op → one or more HLO ops. Graph captures computation structure.
↓
XLA Optimizer — Op Fusion + Tiling + Layout
Op fusion: softmax(QKᵀ) → fused into one HBM pass (like FlashAttention). Tiling: matrices cut into 128×128 tiles for MXU. Layout: choose row-major vs. col-major per op.
↓
LLO (Low-Level Operations) — TPU-specific
Maps MXU tiles to systolic array. Maps VPU ops (softmax, norm, activation) to vector unit. Schedules DMA prefetches so data arrives before MXU needs it.
↓
TPU Program Binary (static, compiled)
Runs with no runtime overhead. All HBM prefetches pre-scheduled. No kernel launch overhead per op. Executes deterministically in fixed time.
XLA's key optimization: operator fusion. In PyTorch on GPU, each layer operation
(GEMM → LayerNorm → GeLU → another GEMM) triggers a separate CUDA kernel launch with a
separate HBM read/write for intermediate results. XLA identifies chains of operations and
fuses them into one kernel that reads inputs once and writes outputs once.
For a Transformer, XLA might fuse an entire FFN block (linear → activation → linear) into
one HBM access, achieving the FlashAttention-like bandwidth reduction automatically.
This is XLA's primary advantage over CUDA's JIT kernel model for static workloads.
JAX — the native programming model for TPUs
JAX is Google's ML framework designed around XLA from the ground up. It looks like NumPy
but every operation is JIT-compiled to XLA, which targets CPU/GPU/TPU.
# JAX: NumPy API + XLA compilation + functional transforms import jax.numpy as jnp from jax import jit, grad, vmap, pmap
# --- The four key JAX transforms ---
# 1. jit: compile this function to XLA (runs on CPU/GPU/TPU) @jit defattention(Q, K, V, mask=None):
d_k = Q.shape[-1]
scores = jnp.matmul(Q, K.T) / jnp.sqrt(d_k) # → XLA matmul → MXU if mask is not None:
scores = scores + mask * -1e9
attn = jax.nn.softmax(scores, axis=-1) # → XLA → VPU return jnp.matmul(attn, V) # → XLA matmul → MXU
# 2. grad: automatic differentiation — works through any jitted function
loss_and_grads = grad(loss_fn, has_aux=True)(params)
# 3. vmap: vectorize a function over a batch dimension (removes explicit batch loops)
batched_attention = vmap(attention, in_axes=(0, 0, 0))
# 4. pmap: data-parallel over multiple TPU chips / GPUs # Each chip receives a shard of the batch and runs the same function
parallel_train_step = pmap(train_step, axis_name='batch')
# The key difference from PyTorch: # - JAX functions are pure (no in-place mutation, no global state) # - XLA traces the computation graph on first call, compiles to TPU binary # - Subsequent calls execute the compiled binary with no Python overhead # - Dynamic shapes (different seq_len per call) cause recompilation — major pain point
The JAX footgun: dynamic shapes cause recompilation
XLA compiles programs for a specific input shape. If your training loop sees sequences of
length 512 then 256 then 1024, XLA recompiles three times — each compilation takes
30 seconds to 5 minutes for large models. This is the biggest pain point of the JAX/XLA
ecosystem for research code: every new input shape triggers recompilation.
Solutions: padding to fixed lengths (wasteful of compute for short sequences),
bucketing (group similar-length sequences together), or
jax.jit with abstract shapes (experimental, TPU v4+).
PyTorch's JIT compilation and dynamic shapes are more forgiving — which is why
most ML research still starts in PyTorch even at Google, then migrates to JAX/TPU
once the architecture is stable.
Part 4 — TPU Generations: From v1 to v5p
Each generation doubled performance and expanded the interconnect scale
TPU v1
2016 (Gmail)
92 TOPS
INT8 only. Inference-only. 8 GB HBM. First deployed for Gmail Smart Reply and Search ranking.
TPU v2
2017
45 TFLOPS
BF16 training supported. 16 GB HBM. First Cloud TPU pods (64 chips). AlphaGo trained here.
TPU v3
2018
420 TOPS
Liquid-cooled. 32 GB HBM/chip. 1,024-chip pods. First large-scale BERT/T5 training.
v5p: 32 GB HBM, 8,960-chip pod, 4× ICI vs. v4. v5e: inference-optimized. Gemini 1.0 trained on v5p.
The TPU v1 origin story — why Google built it
In 2013, Google's engineers projected that if every user ran a 5-minute/day voice search with
Google's then-state-of-the-art neural network (using CPUs), Google would need to double its
worldwide data center capacity. The economics of GPU training were also alarming given
that Google's primary inference hardware was CPUs (GPUs were power-hungry for 24/7 inference).
Jeff Dean commissioned a custom inference accelerator that became TPU v1 in 2014, deployed
in 2016. The key design decision: optimize for INT8 inference of matrix multiply, minimize
everything else. The result: 30× better performance/watt vs. server CPUs for neural network
inference. This single decision gave Google a 1–2 year AI infrastructure lead in 2016–2018.
Part 5 — TPU Pod Topology: The 3D Torus That Replaced NVSwitch
Why Google doesn't need NVSwitch — and why that matters
NVIDIA's NVSwitch (Lesson 6) is a dedicated switching chip that connects 72 GPUs in an
all-to-all topology. For more than 72 GPUs, you need InfiniBand networking — a slower,
packet-switched interconnect. This creates a bandwidth cliff between intra-node (NVLink,
900 GB/s) and inter-node (InfiniBand, ~800 Gbps = 100 GB/s) communication.
TPU pods use a different topology: a 3D torus. Each chip is connected to
six neighbors (±x, ±y, ±z) via ICI (Inter-Chip Interconnect). Messages route through
the torus to reach any chip. There is no NVSwitch equivalent needed — the torus handles
all-to-all communication through multi-hop routing.
TPU v4 pod — 4×4×4 cube excerpt of the 8×8×8×16 full pod
The 3D torus topology is the TPU pod's key structural advantage for training at 1,000+ chip scale.
NVIDIA's architecture requires NVSwitch for 72-GPU all-to-all connectivity, then InfiniBand
for inter-node communication. The bandwidth drop from NVLink (900 GB/s) to InfiniBand (100 GB/s)
is 9×. TPU's ICI provides consistent 600 GB/s per direction across all 4,096 chips in the pod —
no bandwidth cliff. This is why Gemini Ultra and PaLM 540B training runs
(requiring 4,000+ chips) ran natively on TPU pods.
Part 6 — Investment Implications of the TPU
Google's internal AI cost advantage — and who profits from it
~$75B
Google 2025 capex
Significant portion is TPU manufacturing, data center, and networking.
1.7×
TPU energy efficiency vs. H100
275 TFLOPS BF16 at 170W vs. H100's 3,958 TFLOPS at 700W — per-TFLOP watt advantage for Google.
TSMC
TPU manufacturer
TPU v4/v5 on 7nm/5nm TSMC. Google's internal silicon spend is TSMC revenue — a second large HPC customer alongside NVIDIA.
~12%
NVIDIA revenue from Google (est.)
Google buys H100/B200 for GCP customer instances even while using TPU internally. The two coexist.
Three investment implications of the TPU for a long-term TSMC holder
Google is a significant TSMC HPC customer. TPU v5p is on TSMC 5nm with CoWoS-S packaging. As Google scales Gemini compute requirements, TPU volume grows — directly benefiting TSMC HPC node revenue. Google's AI capex growing from $30B (2023) to $75B+ (2025) includes a large TSMC wafer component. TSMC benefits from Google's AI scaling regardless of whether Google or NVIDIA wins.
TPU reduces NVIDIA's Google-specific TAM, not total AI compute TAM.
Google's TPU compute replaces what would otherwise be H100/B200 orders from Google.
But Google still buys NVIDIA GPUs for GCP customer instances (because customers write CUDA code).
The net effect: Google's own AI R&D (Gemini training, Search, YouTube recommendations) does not flow to NVIDIA; GCP customer revenue does. This is a bounded reduction to NVIDIA's TAM, not an existential threat.
XLA's cross-hardware ambition is the longer-term story.
XLA increasingly targets NVIDIA GPUs (via CUDA backend), AMD GPUs (via ROCm), and future hardware.
JAX's hardware-agnostic compilation is the same vector Triton represents (Lesson 6) — if
XLA can generate near-optimal kernels for any hardware automatically, the CUDA software moat
erodes without anyone needing to write a competing library. Google releasing JAX open-source
and contributing heavily to cross-hardware XLA is not altruism — it is strategic: lower the
value of CUDA lock-in, make hardware more fungible, reduce Google's dependency on NVIDIA
when building XLA-compiled models.
The systolic array resurgence: Cerebras, Groq, and Amazon Trainium
The TPU's systolic array design influenced an entire generation of AI accelerators.
Cerebras WSE-3 (the 900,000-core wafer-scale engine) uses a mesh of processing elements
connected by on-wafer interconnect — the same data-flow-through-hardware principle.
Groq's LPU uses a "Tensor Streaming Processor" with deterministic, static scheduling
(no cache hierarchy, no speculation) — XLA-like ahead-of-time compilation.
Amazon Trainium uses a custom NeuronCore matrix engine with the Neuron SDK (XLA-based compiler).
All three companies took Google's core insight — dedicated matrix multiply hardware +
compiler-managed data movement + ahead-of-time compilation — and applied it.
The common target: the inference era, where workloads are stable enough for AOT compilation
and the efficiency premium over GPU is maximized.
None of these are publicly investable as of 2025 (Cerebras IPO pending, Groq and Trainium private/captive). The public proxy remains TSMC and SK Hynix — which supply the HBM and wafer fab regardless of which accelerator architecture wins.
Technical quiz
1. In a systolic array computing C = A × B, each Processing Element (PE) does five things per clock cycle. Which of the following correctly lists all five?
Fetch A from registers, fetch B from shared memory, compute A×B, store the product in L1 cache, advance the program counter
Receive A from the left neighbor, receive B from the upper neighbor, compute A×B and add to local accumulator, pass A to the right neighbor, pass B to the lower neighbor
Read A and B from HBM, compute A×B+C, write result to shared memory, signal completion to warp scheduler, increment address counter
Decode instruction, read operands from register file, compute A×B, write to accumulator register, fetch next instruction
2. XLA's operator fusion benefit most closely parallels which technique in the CUDA/NVIDIA ecosystem?
cuBLAS's optimal GEMM tiling that maximizes Tensor Core utilization for a given matrix shape
FlashAttention's tiled computation that avoids materializing the full attention score matrix in HBM by computing attention in SRAM tiles — reducing HBM reads/writes for intermediate results
NCCL's AllReduce ring algorithm that reduces inter-GPU communication traffic by overlapping gradient communication with backward-pass computation
NVIDIA Tensor Core's MMA instruction that computes a 4×4 matrix multiply-accumulate in one clock cycle
3. TPU pods use a 3D torus interconnect instead of NVIDIA's NVSwitch architecture. What specific advantage does this provide at 4,000+ chip scale?
The 3D torus eliminates latency entirely — all chips communicate in exactly one hop regardless of distance
The 3D torus achieves higher peak bandwidth than NVSwitch for any number of chips by using bidirectional links
The torus provides consistent ICI bandwidth (600 GB/s/dir) across all 4,096 chips without requiring a centralized switching chip — avoiding the bandwidth cliff that occurs when NVIDIA systems scale beyond 72 GPUs (NVLink) and must use InfiniBand (9× lower bandwidth)
The 3D torus enables custom routing protocols that prioritize gradient AllReduce communication over other traffic types, improving training efficiency
4. JAX's XLA compilation fails or causes repeated recompilation when working with dynamic shapes. Which real-world serving scenario does this make most problematic for TPU deployment?
Training a model where each epoch has a different number of total tokens — XLA must recompile for every epoch
Inference serving where user prompts have variable lengths — each unique sequence length triggers XLA recompilation, causing compilation overhead (30s–5min) that makes real-time serving with dynamic input shapes very difficult without aggressive padding or bucketing strategies
Distributed training where different chips receive shards of different sizes due to uneven data distribution
Fine-tuning a pre-trained model on a new task, since the fine-tuning graph shape differs from the pre-training graph
5. From a TSMC long-term investment perspective, Google scaling its TPU compute (and thus TSMC wafer orders) while also buying NVIDIA GPUs for GCP customers is best characterized as:
A risk — if Google's TPU program succeeds, TSMC loses the Google wafer revenue since Google would manufacture TPUs in-house
Neutral — Google's TPU revenue to TSMC is offset exactly by the NVIDIA GPU revenue TSMC loses when Google switches to TPU internally
Additive — Google scaling TPU production generates direct TSMC HPC node revenue (TPU on TSMC 5nm), while Google still ordering NVIDIA GPUs for GCP generates additional TSMC revenue through NVIDIA's continued orders. Both flows increase with AI scaling, making TSMC a beneficiary of Google's AI spending regardless of GPU vs. TPU allocation
A modest risk to TSMC's HPC utilization since Google's TPU investments reduce Google's long-term NVIDIA orders, and NVIDIA is TSMC's largest HPC customer
JAX quickstart documentation — Run jax.jit, jax.grad, and jax.vmap in Colab (free TPU access). Build intuition for what XLA compilation feels like vs. PyTorch's eager mode.
Lesson 10: Scaling laws — the Chinchilla result, compute-optimal training, and why understanding the relationship between model size, data, and compute shapes every major AI lab's chip purchasing decision
Ask me anything. Good follow-ups for a tech worker:
"Run an actual JAX attention implementation on Google Colab's free TPU and compare it to PyTorch's eager mode" ·
"Why does XLA pad matrices to multiples of 128 on TPU v4, and what happens to utilization if your matrix is 130×130?" ·
"How does pmap in JAX implement model parallelism differently from PyTorch's DDP?" ·
For an investor:
"If Google's Gemini 2 training requires 10× the compute of Gemini 1, and 70% runs on TPU v5p, how does that flow through to TSMC's HPC segment revenue?" ·
"How does Cerebras's wafer-scale approach differ from TPU's systolic array, and what does that imply about their respective addressable markets?" ·
"What is the right way to model TSMC's exposure to Google AI capex vs. NVIDIA AI capex?"