HeyGen x Google Cloud: Bringing Avatar IV to TPUs

HeyGen is an AI video generation platform featuring avatar models. Avatar IV, the diffusion stack behind our talking-head videos, runs on more than 18B parameters and is available via Web and API. Working with Google Cloud's specialized AI infrastructure performance optimization team, we brought Avatar IV to an eight-chip Trillium (v6e) host and made it 1.86× faster than our first working version. A stream doesn't wait. Avatar IV renders talking-head video chunk by chunk, and if a chunk is late, the video stalls. Every optimization here was made against that deadline. Working with the Google Cloud team, we moved the pipeline onto an eight-chip Trillium (v6e) host and then made it nearly twice as fast as our first working version, with every change passing the same output-quality gates. Three walls set the shape of the work: exposed all-to-all collectives in the mesh, partial blocks in the sparse attention grid, and a serial dependency in the softmax inner loop. This post covers all three, plus the compiler contracts and quality gates behind them. Avatar IV turns a single photo and an audio track into a talking, moving person. Behind the product, three models take turns on every chunk of video: a diffusion transformer that renders motion conditioned on the audio, a second transformer that super-resolves it, and a VAE decoder that turns latents into pixels. The output is 720p or 1080p video at 25 frames per second, streamed as chunks complete, so playback starts while later chunks are still rendering. Avatar IV was written for GPUs, like the ecosystem around it. The port ran through torchax, a PyTorch frontend on JAX: the production model code runs unmodified, dispatched onto JAX arrays and compiled by the XLA compiler. The same model code now targets both stacks, and the TPU-specific engineering concentrates where the hardware differs. Each of the pipeline's attention variants dispatches to a Pallas kernel built for its shape. We also checked what a full native JAX rewrite would buy over this approach. The answer was roughly nothing: XLA compiles the whole pipeline end to end either way, so the frontend's cost is paid once, at trace time. The parallelism was forced by arithmetic. The two transformers total more than 36 GB of bf16 weights against 32 GB of HBM per Trillium chip, so the weights are FSDP-sharded across the eight chips. Trillium's SparseCore, a co-processor that runs alongside the main core, carries each layer's weight gathers asynchronously. In production traces, those gathers hide behind compute: running them on the co-processor physically frees the matrix unit from weight movement, so a sharding forced by memory arithmetic costs nothing on the critical path. Ulysses sequence parallelism combines with the weight sharding over the same mesh, splitting the video sequence itself. The sharding strategy was finalized during the first working version. Everything that followed was kernel and compiler work, and we tracked time per chunk as each change landed. The chart below shows that record as six milestones: Moving from left to right across the graph, time per chunk falls to just over half its starting value-a 1.86× speedup-with the same model and quality gates throughout. The first milestone is the largest drop, representing the known playbook executed in full: custom attention kernels in place of stock ones, the sequence-parallel layout locked in, the XLA flag set tuned against this workload rather than left at defaults, and kernel tile sizes matched to its shapes. The three walls that follow, and the compiler contracts after them, are what that playbook could not reach. The result is a pipeline that streams at performance comparable to what we see from our 8×H100 production setup while being up to 25% more cost efficient per minute of generated video. Ulysses sequence parallelism puts a pair of all-to-alls inside every self-attention: one to trade sequence shards for head shards, and one to trade back. On our traces those collectives sat fully exposed, already moving bytes at 85-90% of the mesh's bisection bandwidth. The wire had no headroom left, which pointed at the program itself: a monolithic chain of all-to-all, then attention, then all-to-all gives XLA nothing to schedule the transfer behind, so it correctly keeps it synchronous. The fix, worked out with the Google Cloud team, was the standard remedy: pipeline the collective. The attention heads split into a few independent groups, each running its own all-to-all, attention, all-to-all pattern, and each group's transfer now has sibling attention to hide behind, so XLA switches to asynchronous start/done pairs. The work was carrying that through the production pipeline and the quality gates. The hard case was the super-resolution stage's sparse attention: its mask is defined over a specific token ordering, and the head grouping has to leave that ordering untouched, so that every group still sees exactly the mask it would have seen unsplit. In the traces, the collective's footprint on the compute stream collapsed roughly 5×, with attention time itself unchanged. The wire time didn't shrink. It moved off the critical path, which is all the deadline cares about. The group count has a sweet spot: over-split, and the per-launch overhead eats what the overlap buys. More than once, a collective change that measured faster in isolation failed to hold at full pipeline depth, where the transfers interleave with each other and with everything else on the wire. Nothing here counted until it won end to end. The largest single kernel in the pipeline is the super-resolution stage's sparse attention. It runs over tens of thousands of tokens with a windowed sparsity pattern: each frame attends to a window of nearby frames plus one global reference frame. That mask is frame-aligned by nature, meaning its live regions snap exactly to frame boundaries. A general-purpose block-sparse kernel has no reason to know about frames. The Pallas splash-attention family we started from tiles the sequence in multiples of 128, matching the hardware's vector lanes, and a frame's token span isn't a multiple of 128. So live blocks straddle frame boundaries, and roughly one in five comes out partial. The misaligned tiling forces the expensive machinery around it: mask predicates in the inner loop wherever a block straddles a boundary, and a padded sequence so the block grid tiles the full length. On top of that, the kernel handled the two parts of the mask, the frame window and the reference frame, as two separate attention passes, writing full-precision intermediates that a merge pass re-read to combine them. The fix was not to make the masked path faster. It was to delete the mask. We relaxed the kernel's block-size constraint from multiples of 128 down to multiples of 16, the finest bf16 tiling the hardware supports along the sequence dimension: fine enough that block sizes can divide the per-frame token span exactly, and still large enough to keep the matrix unit fed. Now every block is either entirely inside the mask or entirely outside it. Live blocks are full by construction, and the reference frame's blocks are simply more full blocks in the same online softmax, so a single pass covers the whole mask. With that, the mask predicates, the second pass, and the padding all disappear. A second round rebuilt the kernel body around the same grid: the faster dense-attention inner loop running over the frame-aligned sparse layout, with softmax reductions restructured to cut register traffic in the hot loop. The alignment round lifted the kernel from about half of the ceiling this attention shape can reach on the hardware to nearly three-quarters of it, and the rebuilt body closed to about 86%. Together the two rounds took more than ten percent off the super-resolution stage. Alignment removed the work between blocks. Inside them, one serial dependency remained: the online softmax. Flash-style attention carries a running maximum per query row, rescaling its accumulator whenever a new block of keys raises it. That bookkeeping is a serial dependency in the hottest inner loop of the pipeline's attention kernels. Together with the Google Cloud team, we replaced the running max with a precomputed upper bound. By the Cauchy-Schwarz inequality, a query's largest possible logit is bounded by its norm times the largest key norm. A tiny array of precomputed norms, fed to the kernel through scalar prefetch, lets each row derive its bound as it starts, and the online max is no longer needed. Mean-centering the keys first tightens the bound, since keys share a large common component, and it's softmax-invariant: the outputs are unchanged in exact arithmetic, and the quality gates described below catch what rounding does in practice. With the max fixed up front, the rescaling and its serial chain drop out of the inner loop entirely. Not every head is a good fit: a bound far above the true maximum pushes the exponentials toward underflow. Eligibility is checked per attention head, and heads whose bound is too loose fall back to the standard online path inside the same kernel. On our production data, 98-99% of heads qualify. Before anything shipped, a joint audit of the math extended the bound's guarantees to our windowed sparse masks, where each query row sees a different slice of the keys. Once the serial dependency is gone, the kernel's optimal block geometry moves, so we re-tuned block shapes together with the new softmax rather than inheriting them from the old one. On the chart, this is one of the steepest drops after the opening playbook. The walls shared a quieter lever: explicit contracts with the compiler. The clearest example is layout. Attention's inputs are produced by a chain of small operations: normalization, rotary embeddings, projections, head packing. The compiler assigns a specific physical layout to the operand of the all-to-all that follows, and any mismatch gets patched with