--- title: "HeyGen x Google Cloud: Bringing Avatar IV to TPUs" url: https://daily.dev/posts/heygen-x-google-cloud-bringing-avatar-iv-to-tpus-qzakruxj5 source_url: https://developers.googleblog.com/heygen-x-google-cloud-bringing-avatar-iv-to-tpus type: article source: "Google Developers" published: 2026-08-13T16:00:54.684Z updated: 2026-08-13T16:01:25.873Z tags: ["performance"] reading_time: 11 upvotes: 0 comments: 0 language: en --- > ## Documentation Index > Fetch the complete documentation index at: https://daily.dev/llms.txt > Use this file to discover all available pages before exploring further. # HeyGen x Google Cloud: Bringing Avatar IV to TPUs **[Google Developers](https://daily.dev/sources/googledevs)** · 11 min read · 0 upvotes · 0 comments ## Summary HeyGen and Google Cloud's AI infrastructure team ported Avatar IV, an 18B+ parameter talking-head video diffusion pipeline, from GPUs to an eight-chip Trillium (v6e) TPU host using torchax (a PyTorch-on-JAX frontend), achieving a 1.86x speedup over the first working version. The engineering tackled three bottlenecks: exposed all-to-all collectives in Ulysses sequence parallelism (fixed by splitting attention heads into pipelined groups so XLA could overlap transfers), partial/misaligned blocks in sparse attention (fixed by shrinking block-size granularity from 128 to 16 tokens to eliminate mask predicates and padding), and a serial dependency in online softmax (replaced with a precomputed Cauchy-Schwarz upper bound covering 98-99% of attention heads). Additional wins came from explicit compiler contracts: fused kernels matching collective layouts, forcing an alternative instruction scheduler via an XLA flag, and attaching honest cost estimates to custom kernels. Every change passed a two-tier output-quality gate (byte-identical hashing or a measured similarity band), and the result matches 8xH100 GPU performance while being up to 25% more cost-efficient per minute of video. ## Full article daily.dev links to this article rather than hosting it. Read it at the original source: ## Questions this post answers ### How much faster did HeyGen's Avatar IV video model get after optimizing it for Google Cloud Trillium TPUs? Avatar IV, HeyGen's 18B+ parameter talking-head video diffusion pipeline, became 1.86x faster on an eight-chip Trillium (v6e) TPU host compared to its first working version, while passing the same output-quality gates throughout. The result streams at performance comparable to an 8xH100 GPU setup while being up to 25% more cost efficient per minute of generated video. _Engineers tuning diffusion models across TPU and GPU stacks can follow real-world performance breakdowns like this on daily.dev._ ### How do you fix exposed all-to-all collectives that block overlap in Ulysses sequence parallelism on TPUs? Split the attention heads into independent groups, each running its own all-to-all, attention, all-to-all sequence, so every group's collective transfer has sibling attention work to hide behind. This lets the XLA compiler switch from synchronous to asynchronous start/done transfer pairs, which cut the collective's footprint on the compute stream roughly 5x in HeyGen's Avatar IV pipeline without changing attention time itself. _Anyone debugging exposed collectives in multi-chip attention pipelines can track techniques like this via daily.dev._ ### How can you eliminate the serial dependency from online softmax's running max in flash attention? Replace the running max with a precomputed upper bound derived from the Cauchy-Schwarz inequality: a query's largest possible logit is bounded by its norm times the largest key norm. Precomputing these norms via scalar prefetch removes the serial rescaling chain from the inner loop; on HeyGen's production data 98-99% of attention heads qualified, with the rest falling back to the standard online softmax path. _Developers optimizing flash-attention kernels can find deep technical breakdowns like this on daily.dev._ ## Similar posts on daily.dev - [Our eighth generation TPUs: two chips for the agentic era](https://daily.dev/posts/our-eighth-generation-tpus-two-chips-for-the-agentic-era-0d0ehgeeq) · Hacker News · 0 upvotes · 0 comments - [How Lightricks trains video diffusion models at scale with JAX on TPU](https://daily.dev/posts/how-lightricks-trains-video-diffusion-models-at-scale-with-jax-on-tpu-cikcoqltz) · Google Cloud · 0 upvotes · 0 comments --- Tags: [#performance](https://daily.dev/tags/performance) [View this post on daily.dev](https://daily.dev/posts/heygen-x-google-cloud-bringing-avatar-iv-to-tpus-qzakruxj5)