Google's TPU Strategy: Vertical Integration and What It Tells Us About AI Compute

In 2013, Google’s engineers realized that if every user of Google Search used the company’s then-experimental voice search feature for three minutes per day, they would need to double Google’s global data center compute capacity to handle the inference load. The neural network models running voice recognition were simply too computationally expensive to run on the commodity server hardware Google had built its infrastructure on. The response was to design a custom chip.

The Tensor Processing Unit, first deployed internally in 2015 and publicly described in a 2017 paper, was not designed to be the most powerful chip. It was designed to run Google’s neural network inference as efficiently as possible, measured in operations per watt per dollar. That constraint — efficiency above raw performance — shaped every architectural decision in the TPU and differentiated it from Nvidia’s GPU approach in fundamental ways that are still relevant today.

The Architecture of Efficiency

The original TPU is, from an architectural standpoint, deeply conservative. It uses 8-bit integer arithmetic rather than 16-bit or 32-bit floating point. It has a simple, deeply pipelined design with a 65,536-unit matrix multiply unit (the MXU) that processes a 256x256 matrix multiplication per clock cycle. It has no branch prediction, no out-of-order execution, none of the complexity that makes general-purpose processors fast for arbitrary code but expensive in silicon area and power.

The bet was that neural network inference is not arbitrary code. It is a specific set of operations — matrix multiplications, followed by element-wise nonlinearities — performed repeatedly in a predictable sequence. A chip designed exclusively for these operations, without the overhead of general-purpose features, can execute them far more efficiently than a general-purpose chip.

Related reading

The original TPU delivered 92 TOPS at 40 watts and fit in a standard data center server slot. An Nvidia K80, the contemporary high-performance GPU, delivered roughly 6 TFLOPS at 300 watts. The TPU’s advantage in operations per watt for inference workloads was roughly 40:1. The paper Google published in 2017 showed that the TPU delivered a 14-16x performance advantage for neural network inference compared to contemporary CPUs and a comparable advantage over GPUs, with a 41x performance-per-watt advantage.

Training as the Second Chapter

The original TPU could only run inference, not training. Neural network training requires higher-precision arithmetic (FP32 or mixed FP16/FP32) and the ability to compute gradients through backpropagation, which the original TPU’s simplified architecture did not support.

The Cloud TPU v2, announced in 2017, added training capability by including FP16 and BF16 arithmetic support and the ability to connect multiple chips in a torus topology through high-bandwidth interconnects. The torus network — where each chip connects to four neighbors in a two-dimensional grid, forming a donut shape — allows collective communication patterns (all-reduce, for averaging gradients) to be executed efficiently across hundreds of chips without a central interconnect switch.

This torus topology, which Google has refined through five TPU generations, is one of the least-discussed but most important differentiators between the TPU system and Nvidia’s GPU cluster approach. Nvidia’s DGX cluster uses NVLink for intra-server chip communication and InfiniBand or Ethernet for inter-server communication — a hierarchical approach that creates different bandwidth tiers and requires software to be written to navigate them. The TPU pod topology (a 3D torus of up to 4,096 chips in the v4 pod) provides uniform all-to-all bandwidth, simplifying the distributed training software considerably.

The Gemini Evidence

The performance of the TPU program as infrastructure for training frontier models is best evidenced by what those chips have produced. Gemini Ultra, Google’s frontier model released in December 2023, was trained entirely on TPU v4 pods. The training run used approximately 90 exaflop-days of compute — equivalent to running a single exaflop machine for 90 days — across 4,096 TPU v4 chips per pod, with multiple pods used in parallel.

The fact that Google could orchestrate a training run of this scale on custom hardware, without relying on Nvidia at all, is a competitive capability that no other AI lab can replicate. OpenAI trains on Nvidia A100 and H100 clusters (some Microsoft-provided, some Azure). Anthropic uses a mix of Nvidia and Amazon Trainium. Meta trains on a combination of Nvidia hardware and its own MTIA (Meta Training and Inference Accelerator) chips. Only Google has the training infrastructure entirely in-house from silicon to model.

This vertical integration produces cost advantages at scale that compound over time. When Google trains a new model version, the marginal cost is the electricity plus the amortized capital cost of TPU hardware that Google owns outright. When OpenAI trains a model version, the cost is GPU cloud rental at market prices, plus Nvidia’s hardware margin, plus Microsoft Azure’s infrastructure margin. For models trained infrequently on short runs, the cloud rental approach is economically rational. For hyperscale training that runs continuously at enormous scale, owning the infrastructure becomes substantially cheaper.

Related reading

TPU v5 and the Inference Push

The TPU v5, described in Google’s 2024 publications, represents a significant evolution beyond the inference-focused original architecture toward a balanced training and inference platform. The chip achieves 918 TFLOPS of BF16 performance, roughly comparable to the H100, but in a packaging and system integration that Google claims produces better total cost of ownership for Gemini-class models.

The more interesting development is TPU v5e — a scaled-down, inference-optimized version that reduces per-chip cost while maintaining the interconnect architecture. The torus topology scales to 256 v5e chips in a pod, providing a inference infrastructure that Google can size to match demand for Gemini Pro and Gemini Flash serving. By maintaining the same programming model across training and inference TPU variants, Google’s ML engineers can train on v5 pods and immediately deploy on v5e pods without model conversion or optimization work.

This is the key advantage of vertical integration that external analysts consistently underestimate: it is not primarily about chip performance. It is about eliminating the friction in the development cycle between training and inference deployment. When the training hardware and the inference hardware share the same software stack, the model produced by training runs on the inference hardware without adaptation. In a GPU-centric world, converting a model trained on A100s to run efficiently on H100s for inference requires non-trivial optimization work. Google’s integrated stack eliminates this.

Why Amazon and Microsoft Are Copying the Strategy

Amazon’s Trainium and Inferentia chips, and Microsoft’s Maia chip, are direct responses to the competitive threat of Google’s vertical integration. Amazon has been explicit that Trainium’s purpose is to reduce AWS’s dependence on Nvidia for its own model training infrastructure. The cost savings for Amazon’s scale — it runs Bedrock, SageMaker, and its own internal AI applications — justify enormous capital investment in custom silicon.

Microsoft’s position is more complicated. The company has committed $13 billion to OpenAI and runs Azure’s GPU cluster as core infrastructure for that relationship. Building a chip that competes with the Nvidia GPUs it is contractually obligated to provide OpenAI would create obvious conflicts. Microsoft’s Maia chip is positioned as inference infrastructure for its own Azure AI services rather than as a replacement for the Nvidia clusters OpenAI trains on.

The pattern is consistent across all three: hyperscalers who can achieve sufficient scale to amortize custom chip development are doing so, pulling value from Nvidia upward into their own stacks. This is not surprising — it is exactly what happened in networking (hyperscalers built their own switches), servers (hyperscalers built their own server designs rather than buying Dell and HP), and storage (hyperscalers built their own storage systems). AI compute is following the same trajectory.

The Developer Ecosystem Question

The TPU’s Achilles heel has historically been programmability. ML researchers train on GPUs because PyTorch, the dominant research framework, runs on GPUs natively. Running PyTorch models on TPUs requires the XLA (Accelerated Linear Algebra) compiler, which works well for models designed for it but requires changes to code written for GPU execution.

Google has invested substantially in JAX, its own ML research framework, which was designed from the start for XLA compilation and TPU execution. JAX has gained significant adoption in the ML research community — DeepMind, which Google acquired in 2014, uses JAX for its frontier research, and a meaningful portion of academic ML research has shifted to JAX since 2021. But PyTorch remains dominant in the broader ML engineering community, and the TPU’s PyTorch support, while improved, still involves more friction than native GPU execution.

Related reading

This developer ecosystem gap is the primary reason TPUs are not a Nvidia replacement for the broader market. For Google’s internal ML infrastructure, where the company can invest in JAX tooling and hire engineers who know XLA, the ecosystem gap is manageable. For an ML startup that hires engineers trained on PyTorch and wants to run on Google Cloud, the friction is real.

The long-term trajectory is toward more tooling convergence — PyTorch on XLA has improved substantially, and Google has the incentive to invest in making it better. But the current state is that TPUs are optimally productive for organizations willing to invest in the JAX/XLA ecosystem, and GPUs are optimally productive for the broader ML engineering world. Google can live with that. It does not need to win the external market to justify the TPU program — it needs its own infrastructure to be competitive with the GPU alternative, which it clearly is.

The TPU program is the clearest demonstration in the AI hardware industry of what it actually means to be vertically integrated: not just performance on a benchmark, but a self-reinforcing system where each layer reinforces the others, and where the economic benefits of integration compound over time. Every other company watching Google has reached the same conclusion and is trying to replicate it. Most will partially succeed. None will build exactly what Google has.