---
title: "HeyGen x Google Cloud: Bringing Avatar IV to TPUs"
description: "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."
---

# HeyGen x Google Cloud: Bringing Avatar IV to TPUs

[Google](https://yomu.fyi/company/google) · Aug 13, 2026

**Type:** Benchmark

## 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.

## Context

Avatar IV renders talking-head video chunk by chunk, so a late chunk stalls playback. Its two transformers contain more than 36 GB of bf16 weights, exceeding the 32 GB of HBM on each Trillium chip, while Ulysses sequence parallelism introduces all-to-all communication inside attention.

## Approach / What changed

The pipeline was moved to an eight-chip Trillium host through torchax, JAX, and XLA, with FSDP weight sharding, Ulysses sequence parallelism, and shape-specific Pallas kernels. Optimizations pipelined collectives, removed unnecessary sparse-attention masking, overlapped softmax work, eliminated layout repacks, tuned compiler scheduling, and added kernel cost estimates and output-quality gates.

## Takeaways

- 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.
- 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.
- 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.

**Tags:** [Google Cloud](https://yomu.fyi/topic/gcp), [JAX](https://yomu.fyi/topic/jax), [Performance](https://yomu.fyi/topic/performance), [TPU](https://yomu.fyi/topic/tpu)

- Source: [Google](https://developers.googleblog.com/heygen-x-google-cloud-bringing-avatar-iv-to-tpus)
- Source URL: https://developers.googleblog.com/heygen-x-google-cloud-bringing-avatar-iv-to-tpus
- Ingested by Yomu: 2026-09-09T14:27:42.657Z

[Read original post](https://developers.googleblog.com/heygen-x-google-cloud-bringing-avatar-iv-to-tpus)