Google’s Cloud TPU v5e is making a strong case for itself as the budget-friendly workhorse for generative AI inference. At less than half the cost of its predecessor, TPU v4, the v5e is purpose-built to bring large-scale training and deployment within reach for organizations that couldn’t previously justify the expense.
The team behind the Diffusers JAX integration has built a demo that runs Stable Diffusion XL on several TPU v5e-4 instances, each packing 4 TPU chips. The result: four 1024×1024 images served in about 4 seconds, including format conversions, communications overhead, and frontend processing. Strip away that overhead and the actual generation time drops to roughly 2.3 seconds.
The speed comes from two JAX features working in tandem. First, just-in-time compilation traces code on the first run and generates highly optimized TPU binaries that get reused on subsequent calls. The catch is that all input, intermediate, and output shapes must be static — change a shape and you trigger a costly recompilation. Image generation fits this constraint neatly: fixed image counts and dimensions mean constant output shapes, and Stable Diffusion’s text embeddings are fixed-shape vectors by design.
Second, JAX’s pmap function expresses single-program multiple-data programs, letting workloads scale across multiple XLA devices without compromising performance. Run SDXL on a TPU with 8 chips and you get 8 images in the same time it takes 1 chip to produce a single image. The v5e instances scale from 1-chip shapes all the way to 256-chip pods with ultra-fast ICI links between chips.
One implementation detail stands out for anyone attempting this themselves: model parameters load in 32-bit precision by default, but converting them to bfloat16 saves memory and accelerates computation. However, the scheduler state must stay in float32 — precision errors otherwise accumulate and produce degraded or even black images. It’s a small but critical caveat that separates a working deployment from a frustrating debugging session, and it’s the kind of practical knowledge that rarely makes it into the headline benchmarks.
