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.
Table of contents
The Workload, and the PortSix MilestonesHiding the CollectiveDeleting the MaskUnchaining the SoftmaxLayout Is the ABIProving the Pixels Didn't ChangeAcknowledgementsQuestions 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.