$\"Gemini_Generated_Image_eshxr7eshxr7eshx\"\/$ <\/p>\n

HeyGen is an AI video technology platform that includes avatar fashions. Avatar IV, the diffusion stack behind our talking-head movies, runs on greater than 18B parameters and is on the market through Internet and API.<\/i> Working with Google Cloud’s specialised AI infrastructure efficiency optimization group, we introduced Avatar IV to an eight-chip Trillium (v6e) host and made it 1.86\u00d7 sooner than our first working model.<\/i><\/p>\n

A stream does not wait. Avatar IV renders talking-head video chunk by chunk, and if a bit is late, the video stalls. Each optimization right here was made in opposition to that deadline. Working with the Google Cloud group, we moved the pipeline onto an eight-chip Trillium (v6e) host after which made it practically twice as quick as our first working model, with each change passing the identical output-quality gates. Three partitions set the form of the work: uncovered all-to-all collectives within the mesh, partial blocks within the sparse consideration grid, and a serial dependency within the softmax interior loop. This submit covers all three, plus the compiler contracts and high quality gates behind them.<\/p>\n

The Workload, and the Port<\/b><\/h2>\n
Avatar IV turns a single picture and an audio observe right into a speaking, shifting individual. Behind the product, three fashions take activates each chunk of video: a diffusion transformer that renders movement 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 full, so playback begins whereas later chunks are nonetheless rendering.<\/p>\n<\/div>\n
\n
\n
$\"image7\"\/$ <\/p>\n
\n The Avatar IV pipeline on a Trillium host. Two diffusion transformers and a VAE decoder take activates each chunk. Weights are FSDP-sharded and sequences Ulysses-sharded throughout the identical eight-chip mesh.\n <\/p>\n<\/p><\/div><\/div>\n
\n
Avatar IV was written for GPUs, just like the ecosystem round it. The port ran by way of torchax, a PyTorch frontend on JAX: the manufacturing mannequin code runs unmodified, dispatched onto JAX arrays and compiled by the XLA compiler. The identical mannequin code now targets each stacks, and the TPU-specific engineering concentrates the place the {hardware} differs. Every of the pipeline’s consideration variants dispatches to a Pallas kernel constructed for its form. We additionally checked what a full native JAX rewrite would purchase over this method. The reply was roughly nothing: XLA compiles the entire pipeline finish to finish both method, so the frontend’s value is paid as soon as, at hint time.<\/p>\n
The parallelism was pressured by arithmetic. The 2 transformers whole greater than 36 GB of bf16 weights in opposition to 32 GB of HBM per Trillium chip, so the weights are FSDP-sharded throughout the eight chips. Trillium’s SparseCore, a co-processor that runs alongside the principle core, carries every layer’s weight gathers asynchronously. In manufacturing traces, these gathers conceal behind compute: working them on the co-processor bodily frees the matrix unit from weight motion, so a sharding pressured by reminiscence arithmetic prices nothing on the important path. Ulysses sequence parallelism combines with the burden sharding over the identical mesh, splitting the video sequence itself.<\/p>\n
Six Milestones<\/b><\/h2>\n
The sharding technique was finalized throughout the first working model. All the things that adopted was kernel and compiler work, and we tracked time per chunk as every change landed. The chart beneath reveals that file as six milestones:<\/p>\n<\/div>\n
\n
\n
$\"image2$ <\/p>\n
\n Determine 1. Relative time per generated video chunk, normalized to our first working TPU model (= 1.00\u00d7). Every milestone bundles a number of adjustments that shipped collectively.\n <\/p>\n<\/p><\/div><\/div>\n
\n
Transferring from left to proper throughout the graph, time per chunk falls to only over half its beginning worth\u2014a 1.86\u00d7 speedup\u2014with the identical mannequin and high quality gates all through. The primary milestone is the most important drop, representing the recognized playbook executed in full: customized consideration kernels instead of inventory ones, the sequence-parallel format locked in, the XLA flag set tuned in opposition to this workload reasonably than left at defaults, and kernel tile sizes matched to its shapes. The three partitions that comply with, and the compiler contracts after them, are what that playbook couldn’t attain.<\/p>\n
The result’s a pipeline that streams at efficiency similar to what we see from our 8\u00d7H100 manufacturing setup whereas being as much as 25% extra value environment friendly per minute of generated video.<\/p>\n
Hiding the Collective<\/b><\/h2>\n
Ulysses sequence parallelism places a pair of all-to-alls inside each self-attention: one to commerce sequence shards for head shards, and one to commerce again. On our traces these collectives sat absolutely uncovered, already shifting bytes at 85\u201390% of the mesh’s bisection bandwidth. The wire had no headroom left, which pointed on the program itself: a monolithic chain of all-to-all, then consideration, then all-to-all provides XLA nothing to schedule the switch behind, so it accurately retains it synchronous.<\/p>\n
The repair, labored out with the Google Cloud group, was the usual treatment: pipeline the collective. The eye heads break up into just a few unbiased teams, every working its personal all-to-all, consideration, all-to-all sample, and every group’s switch now has sibling consideration to cover behind, so XLA switches to asynchronous begin\/achieved pairs. The work was carrying that by way of the manufacturing pipeline and the standard gates. The laborious case was the super-resolution stage’s sparse consideration: its masks is outlined over a selected token ordering, and the pinnacle grouping has to depart that ordering untouched, so that each group nonetheless sees precisely the masks it might have seen unsplit. Within the traces, the collective’s footprint on the compute stream collapsed roughly 5\u00d7, with consideration time itself unchanged. The wire time did not shrink. It moved off the important path, which is all of the deadline cares about. The group rely has a candy spot: over-split, and the per-launch overhead eats what the overlap buys. Greater than as soon as, a collective change that measured sooner in isolation failed to carry at full pipeline depth, the place the transfers interleave with one another and with every thing else on the wire. Nothing right here counted till it gained finish to finish.<\/p>\n<\/div>\n
\n
\n
$\"image4$ <\/p>\n
\n Determine 2. Giving XLA one thing to overlap with. Schematic, durations to not scale. A monolithic all-to-all leaves the compiler nothing to cover it behind, so it stays synchronous (high). Cut up the heads into unbiased teams and the transfers type a pipeline: the wire nonetheless carries one switch at a time, however every center hop hides behind one other group’s consideration, leaving solely the primary entry and final return uncovered (backside). Extra teams shrink these uncovered ends additional, which is how the collective’s footprint on the compute stream collapses roughly 5\u00d7, to not zero.\n <\/p>\n<\/p><\/div><\/div>\n
\n
Deleting the Masks<\/b><\/h2>\n
The biggest single kernel within the pipeline is the super-resolution stage’s sparse consideration. It runs over tens of hundreds of tokens with a windowed sparsity sample: every body attends to a window of close by frames plus one international reference body. That masks is frame-aligned by nature, which means its reside areas snap precisely to border boundaries. A general-purpose block-sparse kernel has no purpose to learn about frames. The Pallas splash-attention household we began from tiles the sequence in multiples of 128, matching the {hardware}’s vector lanes, and a body’s token span is not a a number of of 128. So reside blocks straddle body boundaries, and roughly one in 5 comes out partial. The misaligned tiling forces the costly equipment round it: masks predicates within the interior loop wherever a block straddles a boundary, and a padded sequence so the block grid tiles the total size. On high of that, the kernel dealt with the 2 elements of the masks, the body window and the reference body, as two separate consideration passes, writing full-precision intermediates {that a} merge go re-read to mix them.<\/p>\n
The repair was to not make the masked path sooner. It was to delete the masks. We relaxed the kernel’s block-size constraint from multiples of 128 all the way down to multiples of 16, the best bf16 tiling the {hardware} helps alongside the sequence dimension: high-quality sufficient that block sizes can divide the per-frame token span precisely, and nonetheless giant sufficient to maintain the matrix unit fed. Now each block is both fully contained in the masks or fully exterior it. Stay blocks are full by development, and the reference body’s blocks are merely extra full blocks in the identical on-line softmax, so a single go covers the entire masks. With that, the masks predicates, the second go, and the padding all disappear.<\/p>\n<\/div>\n
\n
\n
$\"image8\"\/$ <\/p>\n
\n Determine 3. The masks, deleted by development. Illustrative geometry. The masks’s reside area (every body attending to a window of frames plus a reference) snaps to border boundaries. Generic block sizes straddle these boundaries, and the partial blocks they create power masks logic and padding on the entire kernel (left). Blocks that divide the body’s token span are full or skipped, by no means partial, so each disappear by development (proper).\n <\/p>\n<\/p><\/div><\/div>\n
\n
A second spherical rebuilt the kernel physique across the similar grid: the sooner dense-attention interior loop working over the frame-aligned sparse format, with softmax reductions restructured to chop register site visitors within the scorching loop. The alignment spherical lifted the kernel from about half of the ceiling this consideration form can attain on the {hardware} to just about three-quarters of it, and the rebuilt physique closed to about 86%. Collectively the 2 rounds took greater than ten p.c off the super-resolution stage.<\/p>\n
Alignment eliminated the work between blocks. Inside them, one serial dependency remained: the web softmax.<\/p>\n
Unchaining the Softmax<\/b><\/h2>\n
Flash-style consideration carries a working most per question row, rescaling its accumulator at any time when a brand new block of keys raises it. That bookkeeping is a serial dependency within the hottest interior loop of the pipeline’s consideration kernels.<\/p>\n
Along with the Google Cloud group, we changed the working max with a precomputed higher sure. By the Cauchy\u2013Schwarz inequality, a question’s largest doable logit is bounded by its norm occasions the most important key norm. A tiny array of precomputed norms, fed to the kernel by way of scalar prefetch, lets every row derive its sure because it begins, and the web max is now not wanted. Imply-centering the keys first tightens the sure, since keys share a big frequent element, and it is softmax-invariant: the outputs are unchanged in actual arithmetic, and the standard gates described beneath catch what rounding does in observe. With the max mounted up entrance, the rescaling and its serial chain drop out of the interior loop fully.<\/p>\n<\/div>\n
\n
\n
$\"image3\"\/$ <\/p>\n
\n Determine 4. Eradicating the chain from the softmax. On-line softmax rescales its accumulator at any time when a key-value block raises the working max, so each block waits on the final (high). A provable per-row sure, derived from precomputed norms through the Cauchy\u2013Schwarz inequality, fixes the max up entrance so blocks stream by way of the interior loop (backside). Heads the place the sure is simply too unfastened fall again to the web path inside the identical kernel.\n <\/p>\n<\/p><\/div><\/div>\n
\n
Not each head is an effective match: a sure far above the true most pushes the exponentials towards underflow. Eligibility is checked per consideration head, and heads whose sure is simply too unfastened fall again to the usual on-line path inside the identical kernel. On our manufacturing knowledge, 98\u201399% of heads qualify. Earlier than something shipped, a joint audit of the maths prolonged the sure’s ensures to our windowed sparse masks, the place every question row sees a distinct slice of the keys. As soon as the serial dependency is gone, the kernel’s optimum block geometry strikes, so we re-tuned block shapes along with the brand new softmax reasonably than inheriting them from the outdated one.<\/p>\n
On the chart, this is likely one of the steepest drops after the opening playbook.<\/p>\n
Format Is the ABI<\/b><\/h2>\n
The partitions shared a quieter lever: express contracts with the compiler.<\/p>\n
The clearest instance is format. Consideration’s inputs are produced by a sequence of small operations: normalization, rotary embeddings, projections, head packing. The compiler assigns a selected bodily format to the operand of the all-to-all that follows, and any mismatch will get patched with copies. We fused that chain right into a single Pallas kernel that writes its output in precisely the format the collective needs. The kernel’s output buffer is the collective’s enter buffer. In a single case, a five-stage repack chain between the projections and the collective disappeared outright.<\/p>\n<\/div>\n
\n
\n
$\"image6\"\/$ <\/p>\n
\n Determine 5. Format is the ABI. The compiler assigns the all-to-all a bodily format for its operand, and any mismatch is patched with copies (high). The fused kernel writes its output immediately in that format, so the kernel’s output buffer is the collective’s enter buffer and the repack chain disappears (backside).\n <\/p>\n<\/p><\/div><\/div>\n
\n
Flags are contracts too. The heavy self-attention kernels solely hit their tuned velocity below an alternate instruction scheduler that higher overlaps the softmax’s vector work with the matrix unit. So every kernel requests it by title (`XLA_TPU_FORCE_LP_LLO_SCHEDULER<\/code>) reasonably than trusting the defaults.<\/p>\n`
`Two extra contracts paid off the identical method. XLA costs a customized kernel at zero FLOPs and nil bytes, since it might’t see inside, so its latency-hiding scheduler misprices every thing round it. Attaching trustworthy value estimates to our in-house kernels purchased again time with no kernel or graph change in any respect: the scheduler merely repriced what it might conceal.<\/p>\n`
`And since the unit of cargo on TPU is a compiled program, our executables are launch artifacts: cached, versioned, and promoted between environments like mannequin weights, gated on zero recompiles and bit-exact output.<\/p>\n`
Proving the Pixels Did not Change<\/b><\/h2>\nEach optimization right here handed a two-tier output gate earlier than it shipped.<\/p>\n Tier one is byte-identical. The delivered video should hash equal to baseline, body for body. Re-tilings and scheduling adjustments that depart each discount order intact are held to this bar, they usually meet it. One change you would possibly count on to indicate up in pixels met the identical bar: we moved consideration’s remaining matrix product to bf16 on the matrix unit, and the delivered video nonetheless hashed similar at each resolutions. No matter numeric distinction it launched vanished when the frames have been quantized to output pixels, earlier than the video was ever encoded.<\/p>\n Tier two covers adjustments that alter the compiled program’s discount order. These are held to the slim similarity band that bf16 reassociation itself produces, a band we measured independently. Something beneath that band goes to the mannequin’s house owners for a blind, frame-by-frame overview earlier than it ships. And every lever was measured one by one in opposition to a twin baseline: two runs of the identical construct that should hash similar earlier than a delta is trusted. One candidate moved the diffusion transformers’ residual stream all the way down to bf16 and measured sooner finish to finish, however its output fell beneath the band. The speedup was popping out of high quality, so we discarded it and shipped the slower model.<\/p>\n Acknowledgements<\/b><\/h2>\nOur due to the Google Cloud group for his or her assist all through this work. The group is Google Cloud’s engineering group for optimizing large-scale AI workloads on its accelerators, working with clients from early proof of idea by way of full manufacturing.<\/p>\n \nHeyGen Workforce:<\/b> Alireza Dolatabadi, Rui Zhang, Onee Yekeh, Rong Yan, Charly Hong<\/li>\n Google Cloud AI Workforce:<\/b> Rishabh Manoj, Sagar Chapara, Prisha Jain, Hitesh Yadav, Shamik Ray<\/li>\n Google Cloud Subject & Buyer Engineering Groups:<\/b> Travis Martin, Jennifer Liang<\/li>\n<\/ul>\nThis work was a collaboration between Google Cloud and HeyGen. We might prefer to thank everybody concerned for his or her assist all through this undertaking.<\/p>\n<\/div><\/div>\n\n","protected":false},"excerpt":{"rendered":" HeyGen is an AI video technology platform that includes avatar fashions. Avatar IV, the diffusion stack behind our talking-head movies, runs on greater than 18B parameters and is on the market through Internet and API. Working with Google Cloud’s specialised AI infrastructure efficiency optimization group, we introduced Avatar IV to an eight-chip Trillium (v6e) host […]<\/p>\n","protected":false},"author":2,"featured_media":17985,"comment_status":"open","ping_status":"open","sticky":false,"template":"","format":"standard","meta":{"footnotes":""},"categories":[56],"tags":[5950,1458,234,81,10260,7308],"class_list":["post-17983","post","type-post","status-publish","format-standard","has-post-thumbnail","hentry","category-software","tag-avatar","tag-bringing","tag-cloud","tag-google","tag-heygen","tag-tpus"],"_links":{"self":[{"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/posts\/17983","targetHints":{"allow":["GET"]}}],"collection":[{"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/posts"}],"about":[{"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/types\/post"}],"author":[{"embeddable":true,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/users\/2"}],"replies":[{"embeddable":true,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=%2Fwp%2Fv2%2Fcomments&post=17983"}],"version-history":[{"count":1,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/posts\/17983\/revisions"}],"predecessor-version":[{"id":17984,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/posts\/17983\/revisions\/17984"}],"wp:featuredmedia":[{"embeddable":true,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/media\/17985"}],"wp:attachment":[{"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=%2Fwp%2Fv2%2Fmedia&parent=17983"}],"wp:term":[{"taxonomy":"category","embeddable":true,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=%2Fwp%2Fv2%2Fcategories&post=17983"},{"taxonomy":"post_tag","embeddable":true,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=%2Fwp%2Fv2%2Ftags&post=17983"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}