HeyGen 将其 180 亿+ 参数的 Avatar IV 视频生成模型通过 torchax 和 XLA 移植至 Google Cloud 的 Trillium（v6e）TPU，并在八芯片网格上利用 FSDP 和 Ulysses 序列并行。为实现实时流处理的 1.86 倍加速，工程团队对暴露的全对全集体通信进行了流水线化处理，将稀疏注意力块大小对齐以消除掩码填充，并采用预计算的柯西 - 施瓦茨上界绕过 softmax 的串行依赖。这些定制的 Pallas 内核与编译器优化仅在通过严格的两级质量门禁后部署，以确保像素输出在字节层面完全一致或在数学上等价。