HeyGen 与 Google Cloud：将 Avatar IV 部署到 TPU

HeyGen 是一个具备虚拟形象模型的 AI 视频生成平台。Avatar IV 是支撑我们口播头像视频的扩散模型栈，拥有超过 180 亿参数，并通过 Web 和 API 提供服务。 与 Google Cloud 专业 AI 基础设施性能优化团队合作后，我们将 Avatar IV 部署到一台八芯片 Trillium（v6e）主机上，并使其速度达到首个可用版本的 1.86 倍。

流式传输不会等待。Avatar IV 按视频块逐段渲染口播头像视频；如果某一块延迟，视频就会卡顿。这里的每一项优化都围绕这一时限展开。我们与 Google Cloud 团队合作，将流水线迁移到八芯片 Trillium（v6e）主机上，随后使其速度接近首个可用版本的两倍，同时每次变更都通过相同的输出质量门槛。三道障碍决定了这项工作的形态：网格中暴露在关键路径上的 all-to-all 集合通信、稀疏注意力网格中的部分块，以及 softmax 内循环中的串行依赖。本文将介绍这三项问题，以及其背后的编译器契约和质量门槛。

工作负载与迁移

Avatar IV 可将一张照片和一段音频转换为会说话、会动的人物视频。在产品背后，三个模型会在每一个视频块上依次运行：一个根据音频条件渲染动作的扩散 Transformer、第二个进行超分辨率处理的 Transformer，以及一个将潜变量转换为像素的 VAE 解码器。输出为每秒 25 帧的 720p 或 1080p 视频；视频块完成后即开始流式输出，因此后续视频块仍在渲染时便可开始播放。

Trillium 主机上的 Avatar IV 流水线。两个扩散 Transformer 和一个 VAE 解码器在每个视频块上依次运行。权重通过 FSDP 分片，序列通过 Ulysses 分片，二者都分布在同一个八芯片网格上。

与其周边生态一样，Avatar IV 最初是为 GPU 编写的。此次迁移通过 torchax 完成，后者是构建在 JAX 之上的 PyTorch 前端：生产模型代码无需修改，即可被分派至 JAX 数组并由 XLA 编译器编译。同一套模型代码现在可同时面向两个技术栈，TPU 专属工程工作则集中在硬件差异所在之处。流水线中的每种注意力变体都会分派至一个针对其形状构建的 Pallas 内核。我们还检验了完整原生 JAX 重写相对此方法能够带来什么收益。答案几乎是什么也没有：无论采用哪种方式，XLA 都会端到端编译整个流水线，因此前端成本只会在追踪（trace）时支付一次。

并行策略由算术条件决定。两个 Transformer 的 bf16 权重总计超过 36 GB，而每颗 Trillium 芯片仅有 32 GB HBM，因此权重必须通过 FSDP 分片到八颗芯片上。Trillium 的 SparseCore 是一个与主核心并行运行的协处理器，负责异步执行每层的权重聚合。在生产追踪记录中，这些聚合隐藏在计算之后：在协处理器上执行它们，会从物理上将矩阵单元从权重搬运中释放出来，因此由内存容量决定的分片不会给关键路径带来代价。Ulysses 序列并行与权重分片共同运行在同一网格上，将视频序列本身切分开来。

六个里程碑

分片策略在首个可用版本期间便已确定。之后的工作全部围绕内核与编译器展开；每项变更落地时，我们都会记录每个视频块的处理时间。下图以六个里程碑展示了这一过程：

图 1. 每个生成视频块的相对耗时，以首个可用 TPU 版本（= 1.00×）为基准进行归一化。每个里程碑均包含一组同时发布的变更。

从图中由左至右看，每个视频块的耗时降至起点的一半多一点，即实现 1.86× 加速；全程使用相同的模型和质量门槛。第一个里程碑的降幅最大，代表完整执行了已知的优化方案：以自定义注意力内核替换现成内核，固定序列并行布局，针对这一工作负载而非沿用默认值调优 XLA 标志，并让内核 tile 尺寸匹配实际形状。接下来的三道障碍，以及其后的编译器契约，则是这套常规方案无法解决的问题。

最终得到的流水线，其流式性能可与我们在 8×H100 生产环境中观察到的性能相当，同时每分钟生成视频的成本效率最高提升 25%。

隐藏集合通信

Ulysses 序列并行会在每次自注意力中引入一对 all-to-all：一次将序列分片交换为头部分片，另一次再交换回来。在我们的追踪记录中，这些集合通信完全暴露在关键路径上，已占用网格双向分割带宽（bisection bandwidth）的 85–90%。网络链路没有剩余空间，这说明问题出在程序本身：由 all-to-all、注意力、all-to-all 构成的单体链条没有给 XLA 提供可用于隐藏传输的计算，因此它正确地将通信保持为同步执行。

我们与 Google Cloud 团队共同制定的修复方案，是标准的补救手段：对集合通信进行流水线化。注意力头被拆分为少量独立组，每组各自执行 all-to-all、注意力、all-to-all 模式；每一组的传输现在都可隐藏在其他组的注意力计算之后，因此 XLA 会切换为异步的 start/done 对。真正的工作在于将这一机制贯穿生产流水线和质量门槛。困难之处在于超分辨率阶段的稀疏注意力：其掩码由特定 token 顺序定义，头分组必须保持该顺序不变，从而确保每一组看到的掩码与未拆分时完全相同。在追踪记录中，集合通信在计算流上的占用大约缩小了 5 倍，而注意力本身的耗时没有变化。链路传输时间并未减少，只是被移出关键路径——而这正是时限所关心的全部。分组数量存在最佳点：拆分过多时，每次启动的开销会吞噬重叠所带来的收益。我们不止一次遇到某项集合通信变更在隔离测试中更快、但在完整流水线深度下无法保持优势的情况；在这种情况下，传输会彼此交错，也会与链路上的其他所有工作交错。只有端到端获胜的优化才算数。

图 2. 为 XLA 提供可重叠的工作。示意图，时长不按比例绘制。单体 all-to-all 没有任何可供编译器隐藏其后的工作，因此保持同步状态（上）。将注意力头拆分为独立组后，传输会形成流水线：链路仍然一次只承载一项传输，但每次中间跳转都隐藏在另一组的注意力之后，只剩首次进入和最后返回暴露在外（下）。更多分组可进一步缩短这些暴露的两端，这正是集合通信在计算流上的占用大约缩小 5 倍、而非降至零的原因。

删除掩码

流水线中规模最大的单个内核是超分辨率阶段的稀疏注意力。它处理数以万计的 token，并采用窗口化稀疏模式：每一帧关注邻近帧组成的窗口，以及一个全局参考帧。该掩码天然按帧对齐，意味着其有效区域精确贴合帧边界。通用块稀疏内核没有理由理解帧的概念。我们起初使用的 Pallas splash-attention 内核族，以 128 的倍数对序列进行分块，这与硬件向量通道相匹配；但一帧的 token 跨度并不是 128 的倍数。因此，有效块会跨越帧边界，其中大约五分之一成为部分块。未对齐的分块迫使内核承担昂贵的额外机制：凡是块跨越边界，内循环中就需要掩码谓词；同时还要对序列进行填充，以便块网格覆盖完整长度。此外，该内核将掩码的两个部分——帧窗口和参考帧——作为两次独立的注意力计算处理，写入全精度中间结果，再由一次合并过程重新读取并组合它们。

修复方案不是让带掩码的路径更快，而是删除掩码。我们将内核的块大小约束从 128 的倍数放宽至 16 的倍数；这是硬件在序列维度上支持的最细 bf16 分块：足够精细，令块大小可以精确整除每帧的 token 跨度，同时又足够大，以保持矩阵单元持续工作。现在每个块要么完全位于掩码内，要么完全位于掩码外。有效块在构造上都是完整块，参考帧的块也只是同一个在线 softmax 中额外的完整块，因此一次计算即可覆盖整个掩码。由此，掩码谓词、第二次计算和填充全部消失。

图 3. 通过构造删除掩码。说明性几何图。掩码的有效区域，即每帧关注一个帧窗口加参考帧，会贴合帧边界。通用块大小会跨越这些边界，产生的部分块迫使整个内核使用掩码逻辑和填充（左）。能够整除一帧 token 跨度的块不是完整块就是被跳过，绝不会成为部分块，因此二者都在构造上消失（右）。

第二轮优化在同一网格基础上重建了内核主体：在按帧对齐的稀疏布局上运行更快的稠密注意力内循环，并重构 softmax 归约，以减少热循环中的寄存器流量。对齐优化使该内核的性能从该注意力形状在硬件上可达到上限的大约一半，提升至接近四分之三；重建后的内核主体则达到约 86%。两轮优化合计将超分辨率阶段的耗时降低了超过 10%。

对齐消除了块与块之间的工作。但在块内部，仍存在一个串行依赖：在线 softmax。

解除 softmax 的链条

Flash 风格注意力会为每个查询行维护一个运行中的最大值；每当新的键块抬高该最大值时，它就重新缩放累加器。这种记账逻辑是流水线注意力内核最热内循环中的串行依赖。

我们与 Google Cloud 团队一起，用预先计算的上界替代了运行中最大值。根据柯西—施瓦茨不等式，一个查询可能达到的最大 logit，受其范数与最大键范数之积限制。一个由预先计算的范数组成的小型数组，通过标量预取传入内核，使每一行在开始时即可推导出其上界，不再需要在线最大值。先对键进行均值中心化可以收紧该上界，因为键共享较大的公共分量；而且这对 softmax 不变：在精确算术下输出不会变化，下文所述的质量门槛会捕捉实际舍入所造成的影响。最大值预先固定后，重新缩放及其串行链条便完全从内循环中移除。

图 4. 从 softmax 中移除链条。在线 softmax 会在键值块抬高运行中最大值时重新缩放累加器，因此每个块都要等待前一个块（上）。通过柯西—施瓦茨不等式从预先计算的范数推导出的、可证明的逐行上界，会预先固定最大值，使各块能够穿过内循环流式执行（下）。上界过松的注意力头会在同一内核中回退至在线路径。

并非每个注意力头都适合这种方式：远高于真实最大值的上界会使指数项趋向下溢。系统会按注意力头检查适用性；上界过松的头会在同一内核中回退到标准在线路径。在我们的生产数据中，98–99% 的注意力头符合条件。在任何内容发布前，我们联合审查了数学推导，并将该上界的保证扩展至窗口化稀疏掩码；在这种掩码中，每个查询行看到的键切片各不相同。串行依赖消除后，内核的最佳块几何形状也会改变，因此我们将块形状与新 softmax 共同重新调优，而非沿用旧版本的设置。

在图表中，这是首轮常规优化之后下降幅度最大的优化之一。

布局就是 ABI

这些障碍还共享一个更隐蔽的杠杆：与编译器建立显式契约。

最清晰的例子是布局。注意力的输入由一串小型操作产生：归一化、旋转位置编码、投影、头打包。编译器会为后续 all-to-all 的操作数分配特定物理布局，任何不匹配都会通过复制来修补。我们将这条操作链融合进单个 Pallas 内核，使其以集合通信所需的确切布局写出输出。内核的输出缓冲区就是集合通信的输入缓冲区。在一个案例中，投影与集合通信之间的五阶段重排链被完全消除。

图 5. 布局就是 ABI。编译器会为 all-to-all 的操作数分配物理布局，任何不匹配都会通过复制修补（上）。融合内核直接以该布局写出输出，因此内核的输出缓冲区就是集合通信的输入缓冲区，重排链随之消失（下）。

标志同样是契约。重型自注意力内核只有在采用另一种指令调度器时，才能达到调优后的速度；该调度器可更好地将 softmax 的向量计算与矩阵单元重叠。因此，每个内核都会通过名称请求使用它（XLA_TPU_FORCE_LP_LLO_SCHEDULER），而不是依赖默认设置。

另外两项契约也以同样方式带来了收益。由于 XLA 无法看到自定义内核的内部实现，它会将自定义内核计为零 FLOPs 和零字节，因此其延迟隐藏调度器会错误估价周围的一切。为我们的自研内核附加真实的成本估算

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.

The Workload, and the Port

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.

The Avatar IV pipeline on a Trillium host. Two diffusion transformers and a VAE decoder take turns on every chunk. Weights are FSDP-sharded and sequences Ulysses-sharded across the same eight-chip mesh.

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.

Six Milestones

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:

Figure 1. Relative time per generated video chunk, normalized to our first working TPU version (= 1.00×). Each milestone bundles several changes that shipped together.

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.

Hiding the Collective

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.

Figure 2. Giving XLA something to overlap with. Schematic, durations not to scale. A monolithic all-to-all leaves the compiler nothing to hide it behind, so it stays synchronous (top). Split the heads into independent groups and the transfers form a pipeline: the wire still carries one transfer at a time, but each middle hop hides behind another group's attention, leaving only the first entry and last return exposed (bottom). More groups shrink those exposed ends further, which is how the collective's footprint on the compute stream collapses roughly 5×, not to zero.

Deleting the Mask

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.

Figure 3. The mask, deleted by construction. Illustrative geometry. The mask's live region (each frame attending to a window of frames plus a reference) snaps to frame boundaries. Generic block sizes straddle those boundaries, and the partial blocks they create force mask logic and padding on the whole kernel (left). Blocks that divide the frame's token span are full or skipped, never partial, so both disappear by construction (right).

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.

Unchaining the 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.

Figure 4. Removing the chain from the softmax. Online softmax rescales its accumulator whenever a key-value block raises the running max, so every block waits on the last (top). A provable per-row bound, derived from precomputed norms via the Cauchy–Schwarz inequality, fixes the max up front so blocks stream through the inner loop (bottom). Heads where the bound is too loose fall back to the online path inside the same kernel.

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.

Layout Is the ABI

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 copies. We fused that chain into a single Pallas kernel that writes its output in exactly the layout the collective wants. The kernel's output buffer is the collective's input buffer. In one case, a five-stage repack chain between the projections and the collective disappeared outright.

Figure 5. Layout is the ABI. The compiler assigns the all-to-all a physical layout for its operand, and any mismatch is patched with copies (top). The fused kernel writes its output directly in that layout, so the kernel's output buffer is the collective's input buffer and the repack chain disappears (bottom).

Flags are contracts too. The heavy self-attention kernels only hit their tuned speed under an alternative instruction scheduler that better overlaps the softmax's vector work with the matrix unit. So each kernel requests it by name (XLA_TPU_FORCE_LP_LLO_SCHEDULER) rather than trusting the defaults.

Two more contracts paid off the same way. XLA prices a custom kernel at zero FLOPs and zero bytes, since it can't see inside, so its latency-hiding scheduler misprices everything around it. Attaching honest cost estimates to our in-house kernels bought back time with no kernel or graph change at all: the scheduler simply repriced what it could hide.

And because the unit of shipment on TPU is a compiled program, our executables are release artifacts: cached, versioned, and promoted between environments like model weights, gated on zero recompiles and bit-exact output.

Proving the Pixels Didn't Change

Every optimization here passed a two-tier output gate before it shipped.

Tier one is byte-identical. The delivered video must hash equal to baseline, frame for frame. Re-tilings and scheduling changes that leave every reduction order intact are held to this bar, and they meet it. One change you might expect to show up in pixels met the same bar: we moved attention's final matrix product to bf16 on the matrix unit, and the delivered video still hashed identical at both resolutions. Whatever numeric difference it introduced vanished when the frames were quantized to output pixels, before the video was ever encoded.

Tier two covers changes that alter the compiled program's reduction order. Those are held to the narrow similarity band that bf16 reassociation itself produces, a band we measured independently. Anything below that band goes to the model's owners for a blind, frame-by-frame review before it ships. And each lever was measured one at a time against a dual baseline: two runs of the same build that must hash identical before a delta is trusted. One candidate moved the diffusion transformers' residual stream down to bf16 and measured faster end to end, but its output fell below the band. The speedup was coming out of quality, so we discarded it and shipped the slower version.

Acknowledgements

Our thanks to the Google Cloud team for their help throughout this work. The team is Google Cloud's engineering group for optimizing large-scale AI workloads on its accelerators, working with customers from early proof of concept through full production.

• HeyGen Team: Alireza Dolatabadi, Rui Zhang, Onee Yekeh, Rong Yan, Charly Hong
• Google Cloud AI Team: Rishabh Manoj, Sagar Chapara, Prisha Jain, Hitesh Yadav, Shamik Ray
• Google Cloud Field & Customer Engineering Teams: Travis Martin, Jennifer Liang

This work was a collaboration between Google Cloud and HeyGen. We'd like to thank everyone involved for their support throughout this project.