TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings
2023/04/04 by Norman P. Jouppi, Jouppi, Norman P., George Thomas Kurian +27 · 1 voice · 59 citations
Computer Science · Engineering · #Advanced Optical Network Technologies #Interconnection Networks and Systems #Photonic and Optical Devices #cs.AI #cs.AR #cs.LG #cs.PF
paper · pdf · doi:10.48550/arxiv.2304.01433
openalex publication_date 2023/04/04 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28
Abstract
In response to innovations in machine learning (ML) models, production workloads changed radically and rapidly. TPU v4 is the fifth Google domain specific architecture (DSA) and its third supercomputer for such ML models. Optical circuit switches (OCSes) dynamically reconfigure its interconnect topology to improve scale, availability, utilization, modularity, deployment, security, power, and performance; users can pick a twisted 3D torus topology if desired. Much cheaper, lower power, and faster than Infiniband, OCSes and underlying optical components are <5% of system cost and <3% of system power. Each TPU v4 includes SparseCores, dataflow processors that accelerate models that rely on embeddings by 5x-7x yet use only 5% of die area and power. Deployed since 2020, TPU v4 outperforms TPU v3 by 2.1x and improves performance/Watt by 2.7x. The TPU v4 supercomputer is 4x larger at 4096 chips and thus ~10x faster overall, which along with OCS flexibility helps large language models. For similar sized systems, it is ~4.3x-4.5x faster than the Graphcore IPU Bow and is 1.2x-1.7x faster and uses 1.3x-1.9x less power than the Nvidia A100. TPU v4s inside the energy-optimized warehouse scale computers of Google Cloud use ~3x less energy and produce ~20x less CO2e than contemporary DSAs in a typical on-premise data center.
Cited by
- Photonic convolutional neural network with pre-trained in situ training
- MoX: Efficient MoE Routing on Direct-Connect Topologies
- ADEPT: Architecture-Driven Energy-Efficient CNN Fine-Tuning on PIM Accelerators
- Enabling Spatially Fine-Grained DVFS in Neural Processing Units for Energy-Efficient LLM Serving
- Auto-Scaling Heterogeneous Neural Processing Units for Energy and Cost-Efficient LLM Serving
- Forgetting in Language Models: Capacity, Optimization, and Self-Generated Replay
- Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation Is Wasteful
- Life-Cycle Emissions of AI Hardware: A Cradle-To-Grave Approach and Generational Trends
- Application-Driven Architecture Exploration for Cross-Layer Heterogeneous Systems
- Beyond Prefill-Decode Disaggregation: Dissecting LLM Inference for Heterogeneous Platforms via Dynamic Operator Scheduling
- Focus: A Streaming Concentration Architecture for Efficient Vision-Language Models
- DCO: Dynamic Cache Orchestration for LLM Accelerators through Predictive Management
- Performance Benchmarking of Tensor Trains for accelerated Quantum-Inspired Homogenization on TPU, GPU and CPU architectures
- EventQueues: Autodifferentiable spike event queues for brain simulation on AI accelerators
- RoMe: Row Granularity Access Memory System for Large Language Models
- Efficient Kernel Mapping and Comprehensive System Evaluation of LLM Acceleration on a CGLA
- PointCNN++: Performant Convolution on Native Points
- Cohet: A CXL-Driven Coherent Heterogeneous Computing Framework with Hardware-Calibrated Full-System Simulation
- CAMformer: Associative Memory is All You Need
- Towards a future space-based, highly scalable AI infrastructure system design
- NX-CGRA: A Programmable Hardware Accelerator for Core Transformer Algorithms on Edge Devices
- AccelOpt: A Self-Improving LLM Agentic System for AI Accelerator Kernel Optimization
- Green AI: A systematic review and meta-analysis of its definitions, lifecycle models, hardware and measurement attempts
- PerfDojo: Automated ML Library Generation for Heterogeneous Architectures
- Detecting Anomalies in Machine Learning Infrastructure via Hardware Telemetry
- NeuronMM: High-Performance Matrix Multiplication for LLM Inference on AWS Trainium
- Persona Generators: Generating Diverse Synthetic Personas for Arbitrary Contexts
- HalleluBERT: Let every token that has meaning bear its weight
- SindBERT, the Sailor: Charting the Seas of Turkish NLP
- AsyncHZP: Hierarchical ZeRO Parallelism with Asynchronous Scheduling for Scalable LLM Training
- Enabling Reconfiguration-Communication Overlap for Collective Communication in Optical Networks
- Mixed-Precision Quantization for Language Models: Techniques and Prospects
- From Loop Nests to Silicon: Mapping AI Workloads onto AMD NPUs with MLIR-AIR
- xLLM Technical Report
- MX+: Pushing the Limits of Microscaling Formats for Efficient Large Language Model Serving
- CCL-Bench 1.0: A Trace-Based Benchmark for LLM Infrastructure
- When Light Bends to the Collective Will: A Theory and Vision for Adaptive Photonic Scale-up Domains
- SPAD: Specialized Prefill and Decode Hardware for Disaggregated LLM Inference
- From Tokens to Layers: Redefining Stall-Free Scheduling for MoE Serving with Layered Prefill
- Toward Co-adapting Machine Learning Job Shape and Cluster Topology
- Short-circuiting Rings for Low-Latency AllReduce
- Are We Scaling the Right Thing? A System Perspective on Test-Time Scaling
- High Clockrate Free-space Optical In-Memory Computing
- Design-Space Exploration of Distributed Neural Networks in Low-Power Wearable Nodes
- PCCL: Photonic circuit-switched collective communication for distributed ML
- StreamTensor: Make Tensors Stream in Dataflow Accelerators for LLMs
- CompAir: Synergizing Complementary PIMs and In-Transit NoC Computation for Efficient LLM Acceleration
- Toward Lifelong-Sustainable Electronic-Photonic AI Systems via Extreme Efficiency, Reconfigurability, and Robustness
- Accelerating Frontier MoE Training with 3D Integrated Optics
- Electricity Demand and Grid Impacts of AI Data Centers: Challenges and Prospects
- Dato: A Task-Based Programming Model for Dataflow Accelerators
- AI Compute Architecture and Evolution Trends
- When Routers, Switches and Interconnects Compute: A processing-in-interconnect Paradigm for Scalable Neuromorphic AI
- TOAST: Fast and scalable auto-partitioning based on principled static analysis
- Towards Lock Modularization for Heterogeneous Environments
- ReGate: Enabling Power Gating in Neural Processing Units
- Sustainable AI Training via Hardware-Software Co-Design on NVIDIA, AMD, and Emerging GPU Architectures
- RailX: A Flexible, Scalable, and Low-Cost Network Architecture for Hyper-Scale LLM Training Systems
- Tensor Processing Unit [wikipedia]
Discussions
Related