HeyGen x Google Cloud: Bringing Avatar IV to TPUs

Summary

HeyGen and Google Cloud ported Avatar IV, an 18B-plus-parameter diffusion stack for streamed talking-head video, to an eight-chip Trillium (v6e) host. The pipeline uses two diffusion transformers and a VAE decoder, with FSDP weight sharding and Ulysses sequence parallelism across the mesh, and runs through torchax, JAX, XLA, and shape-specific Pallas kernels. Optimization targeted exposed all-to-all collectives, frame-misaligned sparse-attention blocks, serial softmax work, layout copies, scheduler choices, and custom-kernel cost estimates; changes were measured per video chunk against output-quality gates. The resulting pipeline was 1.86× faster than the first working TPU version, comparable in performance to the 8×H100 production setup and up to 25% more cost efficient per generated-video minute, while shipped changes preserved byte-identical output or passed defined similarity review.

Takeaways

1. 1

   FSDP shards the transformers’ more-than-36-GB bf16 weights across eight Trillium chips, while SparseCore performs asynchronous weight gathers that are hidden behind computation.

2. 2

   Splitting attention heads into groups pipelined all-to-all transfers with attention, reducing the collective’s compute-stream footprint by roughly 5× without reducing wire time.

3. 3

   A two-tier output gate required byte-identical video where possible; changes affecting reduction order had to stay within a measured bf16 similarity band, and a faster residual-stream candidate was rejected for quality loss.