Accelerating Stable Diffusion XL Inference with JAX on Cloud TPU v5e

Hugging Face has integrated JAX support into the Diffusers library to enable high-performance, cost-efficient inference for Stable Diffusion XL (SDXL) on Cloud TPU v5e. This integration addresses the computational challenges of SDXL, whose UNet is approximately three times larger than its predecessor, by leveraging JAX's just-in-time (JIT) compilation and XLA-driven parallelism.

Technical Optimization via JAX and TPU v5e

Serving SDXL on Cloud TPU v5e achieves high efficiency through two primary software and hardware mechanisms: JIT compilation and SPMD parallelism.

JIT Compilation for Static Shapes

JAX utilizes just-in-time (JIT) compilation to trace code during the initial execution and generate optimized TPU binaries for subsequent calls. This process requires static input, intermediate, and output shapes. SDXL is highly compatible with JIT compilation because:

Constant Output Shapes: Image generation typically uses a fixed number of images and a consistent size.
Fixed-Shape Embeddings: Stable Diffusion and SDXL use fixed-shape embedding vectors (with padding) for text prompts.

While the initial compilation takes several minutes (approximately three minutes in the provided example), subsequent inference calls are significantly accelerated.

XLA Parallelism and Throughput

JAX's pmap enables single-program multiple-data (SPMD) execution, allowing workloads to be scaled across multiple XLA devices. This allows for linear scaling of image generation: for instance, a TPU with 8 chips can generate 8 images in the same time a single chip takes to create one. Cloud TPU v5e instances are available in various configurations (from 1 to 256 chips) connected by ultra-fast ICI links, allowing users to scale based on specific throughput needs.

Implementation Pipeline in JAX

Running SDXL inference with JAX involves a functional approach where model parameters are handled separately from the pipeline. Key implementation steps include:

Model Loading: Using FlaxStableDiffusionXLPipeline.from_pretrained to load the base SDXL 1.0 model.
Precision Management: Converting model parameters to bfloat16 to reduce memory usage and increase speed, while maintaining the scheduler state in float32 to prevent precision errors that lead to low-quality or black images.
Input Preparation: Using prepare_inputs to ensure prompts have consistent dimensions across invocations, which is required for JIT compilation.
Device Replication: Replicating parameters and inputs across available TPU chips (e.g., using replicate for a TPU v5e-4) and assigning unique random seeds to each chip to ensure diverse image outputs.
Execution: Calling the pipeline with jit=True to trigger the XLA compilation process.

Performance Benchmarks

Benchmarks conducted on SDXL 1.0 base with 20 steps using the Euler Discrete scheduler show that TPU v5e provides superior cost-efficiency over TPU v4.

Hardware	Batch Size	Latency	Perf/$
TPU v5e-4 (JAX)	4	2.33s	21.46
TPU v5e-4 (JAX)	8	4.99s	20.04
TPU v4-8 (JAX)	4	2.16s	9.05
TPU v4-8 (JAX)	8	4.17s	8.98

TPU v5e achieves up to 2.4x greater performance per dollar compared to TPU v4. Performance is measured by calculating throughput (batch size divided by latency per chip) and dividing that figure by the list price of the hardware.

Deployment Architecture

The current implementation uses a load-balancing server that randomly routes user requests to backend servers running on pre-allocated Cloud TPU v5e-4 instances. Each instance generates four 1024×1024 images in approximately 4 seconds (including frontend processing and communication), with the actual generation time being approximately 2.3 seconds.