EngineersofAI LogoCoursesBlogResearch LabAI LettersThe LabCode BankInteractive 3DKodrEarnestJobs
Skip to main content
Interactive 3D/TPU Architecture and Systolic Arrays
Workload: transformer training
TPUs shine: systolic array maps perfectly to dense matrix ops in attention and FFN layers.
Hardware Comparison (Single Chip)
HardwareBF16 TFLOPSMem BW (GB/s)TDP (W)$/hrRel. Perf
TPU v40 PFLOPS1,200170$3.22
100%
A1000 PFLOPS2,000400$3
72%
H1001 PFLOPS3,350700$8
140%
MXU Systolic Array (128×128, showing 8×8)
Each cell is a Processing Element (PE). Data flows right (matrix A) and down (matrix B). Every PE multiplies and accumulates in parallel - no memory traffic after initial load.
TPU v4 Architecture
Matrix Multiply Unit (MXU)
128×128 systolic array, BF16
Vector Unit (VPU)
activation functions, elementwise ops
HBM2e Memory
32 GB, 1.2 TB/s bandwidth
High-Speed ICI Mesh
4 chips/host, 100 Gbps inter-chip
TPU Pod (v4)
4096 chips, 1.1 EFLOPS BF16
TPU Controls
Workload
Systolic Array: data flows through a grid of Processing Elements, each doing a multiply-accumulate. No random memory access - data arrives predictably like a heartbeat.

Why TPUs win at transformers: attention and FFN are pure dense matrix multiplications - exactly what MXUs are optimized for.

Why GPUs win at irregular workloads: GPUs have general-purpose SMs that can handle branching, sparse ops, and custom CUDA kernels.

TPU Architecture and Systolic Arrays - Interactive Visualization

Google's Tensor Processing Units (TPUs) are application-specific integrated circuits designed from the ground up for dense matrix multiplications - the core operation in transformer training. The key hardware innovation is the Matrix Multiply Unit (MXU): a 128×128 systolic array where data flows through Processing Elements like a heartbeat, each PE doing a multiply-accumulate without touching memory. This eliminates memory bandwidth bottlenecks for matmul-heavy workloads. TPU v4 Pods scale to 4096 chips connected via a high-speed ICI mesh, delivering 1.1 exaFLOPS of BF16 compute - used to train Gemini and PaLM 2. GPUs retain an advantage for irregular, sparse, or custom CUDA workloads where the systolic array's rigid data flow is a poor fit.

  • Systolic array: data flows right (matrix A) and down (matrix B) - no random memory access after initial load
  • TPU v4 vs A100: transformers favor TPU (1.0× vs 0.72×); CNNs favor A100 (1.0× vs 0.70×)
  • Irregular workloads: TPU achieves only 0.25× relative performance - sparse ops misalign with systolic flow
  • TPU Pod: 4096 chips on ICI mesh = 1.1 EFLOPS BF16, used for Gemini and PaLM 2 training

Part of the EngineersOfAI Interactive 3D - free interactive visualizations covering every major concept in machine learning and AI engineering. Hover any element for a plain-English explanation. No code required.