--- title: HeyGen x Google Cloud: Bringing Avatar IV to TPUs description: 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... canonical: https://daily.dev/posts/heygen-x-google-cloud-bringing-avatar-iv-to-tpus-qzakruxj5 twitter:card: summary_large_image twitter:site: @dailydotdev og:type: website og:site_name: daily.dev og:title: HeyGen x Google Cloud: Bringing Avatar IV to TPUs | daily.dev og:description: 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... og:url: https://daily.dev/posts/heygen-x-google-cloud-bringing-avatar-iv-to-tpus-qzakruxj5 og:image: https://api.daily.dev/og/posts/QzakrUxJ5.png og:image:alt: HeyGen x Google Cloud: Bringing Avatar IV to TPUs og:image:width: 1200 og:image:height: 630 og:locale: 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 is HeyGen's Avatar IV model on Google Cloud Trillium TPUs compared to its first working version? It runs 1.86x faster than the first working version on an eight-chip Trillium (v6e) host. The speedup came from pipelining exposed all-to-all collectives in Ulysses sequence parallelism, aligning sparse attention block sizes to frame boundaries (128 down to 16 tokens), and replacing the serial online softmax max-tracking with a precomputed Cauchy-Schwarz upper bound, all validated against strict output-hashing quality gates. _Engineers tuning diffusion models on TPUs can track hardware-specific optimization writeups like this on daily.dev._ ### How can you avoid a serial dependency in flash attention's online softmax on TPU kernels? Replace the running maximum with a precomputed upper bound derived from the Cauchy-Schwarz inequality, where a query's largest possible logit is bounded by its norm times the largest key norm. Mean-centering keys tightens the bound further. Heads whose bound is too loose (1-2% in production) fall back to the standard online softmax path within the same kernel. _Developers optimizing attention kernels can follow deep technical breakdowns like this via daily.dev._ ### Why do sparse attention kernels end up with many partial or wasted compute blocks when processing frame-based video data? Standard Pallas splash-attention kernels tile sequences in multiples of 128 to match hardware vector lanes, but video frame token spans are not multiples of 128, so roughly one in five blocks straddles a frame boundary and becomes partial. Relaxing the block-size constraint to multiples of 16 lets blocks divide the per-frame token span exactly, eliminating mask predicates, padding, and a separate merge pass for the reference-frame mask. _Teams building windowed sparse attention for video models can find kernel-tuning insights 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 - [Google splits its TPU line in two for the agentic era](https://daily.dev/posts/google-splits-its-tpu-line-in-two-for-the-agentic-era-zbjc2yyc7) · The New Stack · 0 upvotes · 0 comments - [Accelerating Spatio-Temporal Attention for Video Diffusion on TPUs](https://daily.dev/posts/accelerating-spatio-temporal-attention-for-video-diffusion-on-tpus-by8crmpum) · Google Developers · 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) ```json {"@context":"https://schema.org","@graph":[{"@type":"Organization","@id":"https://daily.dev/#organization","name":"daily.dev","url":"https://daily.dev","logo":{"@type":"ImageObject","url":"https://daily.dev/apple-touch-icon.png","width":180,"height":180},"sameAs":["https://twitter.com/dailydotdev","https://github.com/dailydotdev","https://www.linkedin.com/company/daily-dev-ltd"]},{"@type":"WebSite","@id":"https://daily.dev/#website","url":"https://daily.dev","name":"daily.dev","publisher":{"@id":"https://daily.dev/#organization"},"potentialAction":{"@type":"SearchAction","target":{"@type":"EntryPoint","urlTemplate":"https://daily.dev/search?q={search_term_string}"},"query-input":"required name=search_term_string"}}]} {"@context":"https://schema.org","@type":"TechArticle","headline":"HeyGen x Google Cloud: Bringing Avatar IV to TPUs","url":"https://daily.dev/posts/heygen-x-google-cloud-bringing-avatar-iv-to-tpus-qzakruxj5","mainEntityOfPage":{"@type":"WebPage","@id":"https://daily.dev/posts/heygen-x-google-cloud-bringing-avatar-iv-to-tpus-qzakruxj5"},"datePublished":"2026-08-13T16:00:54.684Z","dateModified":"2026-09-14T07:29:57.408Z","description":"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...","image":"https://media.daily.dev/image/upload/f_auto,q_auto/v1/posts/6819c8795a9dd09ce1ab7766133705b9?_a=AQAEuop","thumbnailUrl":"https://media.daily.dev/image/upload/f_auto,q_auto/v1/posts/6819c8795a9dd09ce1ab7766133705b9?_a=AQAEuop","isAccessibleForFree":true,"articleSection":"Google Developers","inLanguage":"en","publisher":{"@type":"Organization","name":"daily.dev","url":"https://daily.dev","logo":{"@type":"ImageObject","url":"https://daily.dev/apple-touch-icon.png","width":180,"height":180}},"author":{"@type":"Organization","name":"Google Developers","logo":"https://media.daily.dev/image/upload/t_logo,f_auto/v1/logos/ea4bbded4e7f45ccb82e121a3156535f","url":"https://daily.dev/sources/googledevs"},"commentCount":0,"discussionUrl":"https://daily.dev/posts/heygen-x-google-cloud-bringing-avatar-iv-to-tpus-qzakruxj5","interactionStatistic":[{"@type":"InteractionCounter","interactionType":{"@type":"LikeAction"},"userInteractionCount":0},{"@type":"InteractionCounter","interactionType":{"@type":"CommentAction"},"userInteractionCount":0}],"keywords":"performance","timeRequired":"PT11M"} {"@context":"https://schema.org","@type":"BreadcrumbList","itemListElement":[{"@type":"ListItem","position":1,"name":"Home","item":"https://daily.dev"},{"@type":"ListItem","position":2,"name":"Google Developers","item":"https://daily.dev/sources/googledevs"},{"@type":"ListItem","position":3,"name":"HeyGen x Google Cloud: Bringing Avatar IV to TPUs"}]} {"@context":"https://schema.org","@type":"FAQPage","@id":"https://daily.dev/posts/heygen-x-google-cloud-bringing-avatar-iv-to-tpus-qzakruxj5#faq","mainEntity":[{"@type":"Question","name":"How much faster is HeyGen's Avatar IV model on Google Cloud Trillium TPUs compared to its first working version?","acceptedAnswer":{"@type":"Answer","text":"It runs 1.86x faster than the first working version on an eight-chip Trillium (v6e) host. The speedup came from pipelining exposed all-to-all collectives in Ulysses sequence parallelism, aligning sparse attention block sizes to frame boundaries (128 down to 16 tokens), and replacing the serial online softmax max-tracking with a precomputed Cauchy-Schwarz upper bound, all validated against strict output-hashing quality gates. Engineers tuning diffusion models on TPUs can track hardware-specific optimization writeups like this on daily.dev."}},{"@type":"Question","name":"How can you avoid a serial dependency in flash attention's online softmax on TPU kernels?","acceptedAnswer":{"@type":"Answer","text":"Replace the running maximum with a precomputed upper bound derived from the Cauchy-Schwarz inequality, where a query's largest possible logit is bounded by its norm times the largest key norm. Mean-centering keys tightens the bound further. Heads whose bound is too loose (1-2% in production) fall back to the standard online softmax path within the same kernel. Developers optimizing attention kernels can follow deep technical breakdowns like this via daily.dev."}},{"@type":"Question","name":"Why do sparse attention kernels end up with many partial or wasted compute blocks when processing frame-based video data?","acceptedAnswer":{"@type":"Answer","text":"Standard Pallas splash-attention kernels tile sequences in multiples of 128 to match hardware vector lanes, but video frame token spans are not multiples of 128, so roughly one in five blocks straddles a frame boundary and becomes partial. Relaxing the block-size constraint to multiples of 16 lets blocks divide the per-frame token span exactly, eliminating mask predicates, padding, and a separate merge pass for the reference-frame mask. Teams building windowed sparse attention for video models can find kernel-tuning insights like this on daily.dev."}}]} ```