HeyGen x Google Cloud: Bringing Avatar IV to TPUs

HeyGen x Google Cloud：将 Avatar IV 带到 TPUs 上 - Google Developers Blog

Google Tag Manager (noscript)

结束 Google Tag Manager (noscript)

HTML

HeyGen x Google Cloud：将 Avatar IV 带到 TPUs 上

2026年8月13日

分享

• Facebook

• Twitter

• LinkedIn

• 邮件

HeyGen 是一个提供虚拟人像模型的人工智能视频生成平台。Avatar IV 是我们生成 talking-head 视频的扩散模型堆栈，拥有超过 180 亿个参数，可通过 Web 和 API 调用。通过与 Google Cloud 的专用 AI 基础设施性能优化团队合作，我们将 Avatar IV 部署到八芯片 Trillium (v6e) 主机上，使其运行速度比我们最初的版本快了 1.86 倍。

视频流不会等待。Avatar IV 逐块渲染 talking-head 视频，如果某一块渲染延迟，视频就会卡顿。所有优化都围绕着这个截止时间展开。与 Google Cloud 团队合作，我们将整个流水线迁移到八芯片 Trillium (v6e) 主机上，并使其速度几乎达到最初版本的两倍，所有修改都通过了相同的输出质量检测。三个技术难点定义了工作的方向：在网格中暴露了所有对所有（all-to-all）的集体通信操作、稀疏注意力网格中的部分块处理，以及 softmax 内部循环中的串行依赖。本文将涵盖这三个难点，以及它们背后的编译器合约和质量检测机制。

工作负载与移植

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

在 Trillium 主机上的 Avatar IV 流水线。两个扩散变换器和一个 VAE 解码器轮流处理每个片段。权重通过 FSDP 分片，序列通过 Ulysses 在相同八芯片网格中分片。

Avatar IV 最初是为 GPU 编写的，其生态系统也围绕 GPU 构建。移植工作通过 torchax 完成，这是一个基于 JAX 的 PyTorch 前端：生产模型代码无需修改，直接部署到 JAX 数组并由 XLA 编译器进行编译。现在，相同的模型代码可以同时针对两种架构进行优化，而 TPU 特有的工程实现则集中在硬件差异部分。流水线中每种注意力机制都会分发到为其形状定制的 Pallas 内核上。我们还评估了完全重写为原生 JAX 的潜在收益。结果发现几乎没有优势：无论哪种方式，XLA 都能端到端地编译整个流水线，因此前端成本只需在追踪时支付一次。

并行性是由算术强制实现的。两个变换器总共拥有超过 36 GB 的 bf16 权重，而每个 Trillium 芯片只有 32 GB 的 HBM，因此权重必须在八芯片间通过 FSDP 进行分片。Trillium 的 SparseCore 是一个与主核并行运行的协处理器，它异步地为每一层收集权重。在生产追踪中，这些收集操作隐藏在计算之后：在协处理器上运行它们可以物理上释放矩阵单元，使其免于权重移动的干扰，因此由内存算术强制的分片在关键路径上没有成本。Ulysses 序列并行性与同一网格上的权重分片相结合，将视频序列本身进行分割。

分片策略在首个可用版本中最终确定。此后所有工作均围绕内核和编译器展开，我们跟踪每个更改落地后每块数据的耗时。下图显示了这一记录的六个里程碑：

图1. 每个生成视频块的相对耗时，以我们首个可用TPU版本（= 1.00×）为基准。每个里程碑包含若干同时发布的更改。

从图的左侧向右侧移动时，每个块的耗时降至初始值的一半以上（实现1.86倍加速），整个过程使用相同的模型和质量门控。第一个里程碑是最大的降幅，代表完整执行已知方案：使用自定义注意力内核替代标准内核，锁定序列并行布局，针对该工作负载调整XLA标志而非使用默认值，以及将内核分块大小与形状匹配。随后的三个瓶颈及之后的编译器合约是该方案无法触及的领域。

最终实现的流水线在生成视频性能上可与我们的8×H100生产环境相当，同时每分钟视频生成的性价比提升最高达25%。

隐藏集体通信

Ulysses序列并行性在每个自注意力操作中嵌入了两个all-to-all通信：一个用于将序列分片交换为头分片，另一个用于反向交换。在我们的性能追踪中，这些集体通信完全暴露，已占用网格双工带宽的85-90%。链路已无任何余量，这指向程序本身：一个单体的all-to-all→注意力→all-to-all链式结构使XLA无法在后台调度传输，因此正确保持同步。

与Google Cloud团队协作的解决方案是标准方案：将集体通信流水化。将注意力头拆分为若干独立组，每组运行自己的all-to-all→注意力→all-to-all模式，每个组的传输现在可以隐藏在其他组的注意力操作后，使XLA切换为异步启动/完成对。这项工作需要将该方案贯穿生产流水线和质量门控。最具挑战的是超分辨率阶段的稀疏注意力：其掩码定义在特定的token顺序上，头分组必须保持该顺序不变，确保每个组仍能看到与未拆分时完全相同的掩码。在追踪中，集体通信在计算流中的占用面积缩小约5倍，注意力时间本身保持不变。链路时间并未减少，但它已移出关键路径，而这正是截止日期所关心的全部内容。组数量存在最佳值：过度拆分会导致每次启动的开销抵消重叠带来的收益。不止一次出现单独测试更快的集体通信更改，在完整流水线深度下因传输相互交织而未能保持性能。在端到端通过之前，这里的所有优化都不计数。

图2. 为XLA提供可重叠的操作。示意图，持续时间未按比例。单体all-to-all使编译器无处隐藏，因此保持同步（顶部）。将头拆分为独立组后，传输形成流水线：链路仍一次只传输一个数据包，但每个中间跳转都可隐藏在另一组的注意力后，仅暴露第一个入口和最后一个返回（底部）。更多组会进一步缩小这些暴露的端点，这就是集体通信在计算流中的占用面积缩小约5倍的原因（而非完全消失）。

删除掩码

流水线中最大的单个内核是超分辨率阶段的稀疏注意力。它以窗口稀疏模式处理数万个标记：每个帧会关注附近帧的窗口以及一个全局参考帧。该掩码本质上是帧对齐的，意味着其活动区域会精确对齐帧边界。通用的块稀疏内核没有理由了解帧的概念。我们最初采用的Pallas splash-attention家族将序列分块为128的倍数，与硬件的向量通道匹配，而帧的标记跨度并非128的倍数。因此活动块会跨越帧边界，约五分之一的块会部分填充。这种错位的分块方式迫使内核增加额外开销：在内层循环中任何跨越边界的块都需要掩码谓词判断，同时需要填充序列以使块网格覆盖完整长度。此外，内核将掩码的两个部分（帧窗口和参考帧）作为两次独立的注意力传递处理，生成需要合并传递重新读取的全精度中间结果。

解决方案不是让掩码路径运行得更快，而是直接删除掩码。我们将内核的块大小限制从128的倍数放宽到16的倍数——这是硬件在序列维度支持的最细粒度bf16分块：精度足够使块大小能整除每帧的标记跨度，同时仍能保持矩阵单元的持续运算。现在每个块要么完全在掩码内，要么完全在掩码外。活动块通过构造是完整的，参考帧的块只是同一在线softmax中更多的完整块，因此单次传递即可覆盖整个掩码。这样，掩码谓词、第二次传递和填充都消失了。

图3. 通过构造删除掩码。示意图。掩码的活动区域（每个帧关注帧窗口和参考帧）对齐帧边界。通用块大小会跨越这些边界，其产生的部分块迫使整个内核需要掩码逻辑和填充（左）。能整除帧标记跨度的块要么完整要么跳过，从不会部分填充，因此这两个问题通过构造自然消失（右）。

第二轮重构围绕相同网格重新构建内核主体：更快的密集注意力内层循环处理帧对齐的稀疏布局，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 和零字节，因为它无法查看内部结构，因此其隐藏延迟的调度器对周围所有内容都进行了错误定价。将诚实的成本估算附加到我们内部的内核上，无需对内核或图进行任何更改，就直接买回了时间：调度器只需重新评估它能够隐藏的部分。

由于 TPU 的发货单位是编译后的程序，我们的可执行文件属于发布制品：被缓存、版本化，并在不同环境之间传递（如模型权重），且依赖零次重新编译和精确到比特的输出。

证明像素未发生改变

此处的每个优化在发布前都通过了双重输出验证。

第一层验证要求字节完全一致。交付的视频必须逐帧与基线版本的哈希值相等。所有保持每帧约简顺序不变的重新分块和调度更改都需满足这一标准，且均符合要求。一个你可能预期会在像素上体现的更改同样通过了该标准：我们将注意力机制的最终矩阵乘积转移到矩阵单元的 bf16 精度上，但交付的视频在两种分辨率下哈希值仍完全一致。该更改引入的数值差异在帧被量化为输出像素时就已经消失，甚至在视频编码之前就已消失。

第二层验证覆盖改变编译程序约简顺序的更改。这些更改需满足 bf16 重关联本身产生的狭窄相似性范围，该范围我们独立进行了测量。任何低于该范围的更改在发布前都需要由模型所有者进行盲测逐帧审查。每个更改都单独与双基线进行对比测试：同一构建的两次运行必须哈希值一致，才能信任其差异。其中一个候选方案将扩散变压器的残差流降级到 bf16 并实现了端到端加速，但其输出低于相似性范围。该加速是以质量为代价实现的，因此我们弃用了该方案，选择了更慢但质量更高的版本。

致谢

感谢 Google Cloud 团队在整个工作过程中提供的帮助。该团队是 Google Cloud 用于在其加速器上优化大规模 AI 工作负载的工程团队，与客户从早期概念验证到完整生产阶段均有合作。

• HeyGen 团队：Alireza Dolatabadi, Rui Zhang, Onee Yekeh, Rong Yan, Charly Hong

• Google Cloud AI 团队：Rishabh Manoj, Sagar Chapara, Prisha Jain, Hitesh Yadav, Shamik Ray

• Google Cloud 现场与客户工程团队：Travis Martin, Jennifer Liang

这项工作是 Google Cloud 与 HeyGen 的合作成果。我们想感谢所有参与人员在整个项目中提供的支持。

posted in:

• AI

• Cloud

• 案例研究

• 学习中心

上一篇

下一篇

相关文章

列表

AI

案例研究

社区

使用深度学习和 Keras 解码宇宙信号

2026 年 8 月 27 日

操作指南

在 TPU 上运行 Ray，第 2 部分：Ray AI 库

2026 年 7 月 24 日

云服务

公告

最佳实践

云 TPU 上长上下文多模态嵌入推理的企业级精度

2026 年 8 月 26 日

导航点 /