TPUs vs GPUs: Google's Bet on Custom Silicon

NVIDIA GPUs power virtually every AI model you've heard of. GPT-4, Claude, Llama, Stable Diffusion — all trained on NVIDIA hardware.

But Google didn't like that. So they built their own chips. TPUs — Tensor Processing Units. Custom silicon designed specifically for machine learning.

Are they better? Sometimes. Are they worse? Often. Should you care? Absolutely — because understanding when TPUs win (and when they don't) makes you a better AI engineer.

What Is a TPU?

A GPU (Graphics Processing Unit) was originally designed for rendering video game graphics. NVIDIA realized the same parallel processing that draws pixels could also multiply matrices. They added CUDA, turned the ship toward ML, and now own the market.

A TPU (Tensor Processing Unit) was designed from scratch for one thing: accelerating TensorFlow operations. No graphics heritage. No legacy baggage. Just matrix math, hardcoded in silicon.

The Key Difference: Systolic Arrays

GPUs are flexible parallel processors. They can run any CUDA code — graphics, physics simulations, crypto mining, AI.

TPUs use systolic arrays — specialized circuits that do one thing extremely fast: multiply-accumulate operations (the core of neural networks).

GPU: Flexible → Good at everything, great at ML

TPU: Specialized → Great at specific ML ops, inflexible

Think of it like this:

GPU = Swiss Army knife. Does everything.
TPU = Kitchen knife. Only cuts vegetables. But cuts them really well.

TPU Generations

Generation Year Peak Performance Memory Notes

TPU v1 2016 92 TFLOPS 8 GB HBM Inference only

TPU v2 2017 180 TFLOPS 64 GB HBM Training + inference

TPU v3 2018 420 TFLOPS 128 GB HBM 2x v2 performance

TPU v4 2021 275 TFLOPS/chip 32 GB HBM Better efficiency

TPU v5e 2023 197 TFLOPS 16 GB HBM Cost-optimized

TPU v5p 2023 459 TFLOPS 95 GB HBM Performance-optimized

TPU v6e 2024 ~600 TFLOPS TBD Latest generation

Note: TPU performance numbers are for bfloat16. GPUs usually quote FP16 or FP8 numbers, so direct comparison is tricky.

TPU v5p: The Current Flagship

Google's latest training chip:

459 TFLOPS per chip (bfloat16)

95 GB HBM memory

3.2 Tbps chip-to-chip interconnect

8960 chips in a single pod (that's 4.1 exaFLOPS)

For comparison, an H100 does ~1000 TFLOPS in FP8. But raw FLOPS don't tell the whole story.

🔧

When TPUs Win

1. Large-Scale Transformer Training

Google trained Gemini on TPUs. Not GPUs. TPUs.

Why? At massive scale, TPUs are more efficient:

GPT-4 scale training (~50,000 GPUs):

• Power: ~50 MW
• Cost: ~$1B hardware
• Network: Infiniband between nodes

Gemini scale training (~10,000 TPU v5p):

• Power: ~30 MW
• Cost: ~$500M hardware (estimated)
• Network: Built into TPU pod design

TPU pods are designed as unified systems. The network is part of the chip design, not an afterthought. At 10,000+ chip scale, this matters.

1. When You're Already in the Google Ecosystem

If you use:

TensorFlow (especially TF 2.x)

JAX (Google's ML framework)

Google Cloud Storage

BigQuery

Then TPUs integrate seamlessly. The data pipeline, the model code, the checkpointing — all designed together.

1. Specific Model Architectures

TPUs excel at:

Transformers (what doesn't these days?)

Large batch sizes (TPUs love big batches)

Specific activation functions (ReLU, GELU — optimized in silicon)

TPUs struggle with:

Dynamic shapes (variable sequence lengths)

Custom CUDA kernels (you can't write custom TPU ops easily)

Non-standard architectures (anything too exotic)

📊

When GPUs Win

1. Flexibility

Want to try the latest architecture from a research paper? It probably has PyTorch code and CUDA kernels. It probably doesn't have TPU support.

python

New paper drops with this CUDA kernel

import torch

from cuda_extension import fancy_new_attention

Works on GPU immediately

On TPU? File a feature request with Google

1. Inference Latency

TPUs are throughput-optimized. GPUs are latency-optimized.

Metric TPU v5e H100

Batch-1 latency Higher Lower

Throughput (large batch) Excellent Excellent

Cost per token (batch=1) Higher Lower

For real-time chat applications, GPUs usually win. For batch processing, it's closer.

1. Ecosystem

Tool GPU Support TPU Support

If your stack is PyTorch + Hugging Face + vLLM, TPUs are a pain. If your stack is JAX + Flax, TPUs are a joy.

1. Availability

You can buy an H100. You can't buy a TPU. TPUs are Google Cloud only.

This matters for:

On-premise deployments

Multi-cloud strategies

Edge inference

Any workload that needs to leave Google Cloud

JAX: The Secret Weapon

If you're using TPUs, you should probably be using JAX.

What Is JAX?

JAX is Google's ML framework. It looks like NumPy:

python

import jax

import jax.numpy as jnp

from jax import grad, jit, vmap

Looks like NumPy

def predict(params, x):
return jnp.dot(x, params["w"]) + params["b"]

But compiles to XLA for TPUs (and GPUs)

predict_jit = jit(predict)

Automatic differentiation

loss_grad = grad(loss_fn)

Vectorization

batch_predict = vmap(predict, in_axes=(None, 0))

JAX + XLA + TPU = Speed

XLA (Accelerated Linear Algebra) is Google's compiler. It takes your JAX code and generates optimized machine code for TPUs.

Your JAX Code

│
     ▼

XLA Compiler

│
     ▼

Optimized TPU Machine Code

This compilation step takes time on first run ("JIT compilation"), but then execution is extremely fast.

JAX vs PyTorch on TPUs

Feature JAX on TPU PyTorch on TPU

Performance Excellent Good

Ease of use Learning curve Familiar

Ecosystem Growing (Google, DeepMind) Massive

TPU-specific optimizations Native Via torch_xla

Debugging Harder Easier

Rule of thumb: If you're starting fresh on TPUs, use JAX. If you're porting existing PyTorch code, use PyTorch/XLA but expect some friction.

Programming for TPUs: What Changes

1. Static Shapes

TPUs like to know shapes at compile time:

python

GPU: Fine

def forward(x):  # x can be any shape
return model(x)

TPU: Better to pad or use fixed shapes

def forward(x):  # x is always (batch_size, seq_len)
assert x.shape == (32, 512)
    return model(x)

1. Large Batches

TPUs are designed for large batches:

python

GPU: Often optimal at batch_size=8-32

TPU: Often needs batch_size=128-512 to saturate

Use gradient accumulation to simulate large batches

@jax.jit

def train_step(params, batch):
# Process 512 samples at once on TPU
    # vs 32 on GPU with 16 gradient accumulation steps

1. bfloat16

TPUs use bfloat16 (brain floating point) instead of FP16:

Format Exponent Bits Mantissa Bits Range Precision

FP32 8 23 Large High

FP16 5 10 Small Low

bfloat16 8 7 Large Medium

bfloat16 has the same range as FP32 (good for gradients) but less precision. TPUs have native bfloat16 support; it's their preferred format.

python

JAX automatically handles this

from jax import config

config.update("jax_default_dtype", jnp.bfloat16)

🔬

Real-World Benchmarks

Training a standard Transformer (encoder-decoder, 1B parameters):

The pattern: TPUs are 20-30% faster and cheaper for pure training throughput. But the gap isn't massive, and the ecosystem tradeoffs are real.

🛡️

Should You Use TPUs?

Use TPUs If:

You're training large models (10B+ parameters)

You're comfortable with JAX or TensorFlow

You're already on Google Cloud

You need massive scale (1000+ chips)

Your model architecture is standard (Transformer variants)

Use GPUs If:

You need flexibility (research, novel architectures)

You're using PyTorch as your primary framework

You need low-latency inference

You want to run on-premise or multi-cloud

You need the latest models from Hugging Face (immediate support)

📦

The Future: TPU v6 and Beyond

Google isn't slowing down. TPU v6 is already in preview:

Sparse attention acceleration (finally!)

Better inference optimization

Larger pod sizes (16,000+ chips)

Meanwhile, NVIDIA has Blackwell (B100/B200) coming. The chip war is heating up.

The winner? Probably us — competition drives innovation and (hopefully) lower prices.

🚀

Practical Takeaways

TPUs are specialized tools — excellent for large-scale training, inflexible for research

Use JAX for TPUs — PyTorch/XLA works but JAX is the native experience

TPUs love static shapes and large batches — design your code accordingly

bfloat16 is the TPU's native language — use it for best performance

GPUs still win on flexibility and ecosystem — most teams should default to GPUs

Consider TPUs for training, GPUs for inference — a hybrid strategy works for many

🔄

What's Next?

Episode 70: The AI Chip War — TPUs vs GPUs is just one front. AMD has MI300X. Intel has Gaudi. Amazon has Trainium. Groq has the LPU. The entire semiconductor industry is pivoting to AI. We'll map the battlefield and see who might actually challenge NVIDIA's dominance.

← Previous

Ep 68: Cloud GPU Pricing

Next →

Ep 70: The AI Chip War

Next: Episode 70 — The AI Chip War

NVIDIA is worth more than $3 trillion. More than Apple. More than any company in history. Everyone wants to dethrone them.

This is part of a 98-episode series covering AI engineering from tokens to production deployment.