HeyGen x Google Cloud：Avatar IV 模型移植至 TPU 實現 1.86 倍加速

HeyGen 與 Google Cloud 合作，將 Avatar IV 模型移植至 Trillium (v6e) TPU。Avatar IV 是一個包含 18B+ 參數的擴散堆疊，用於生成說話頭像影片，輸出 720p 或 1080p、25 幀率的影片，分塊串流以支持邊渲染邊播放。優化後的性能達到 1.86 倍加速，相比首個可運作版本。

優化工作涉及三個主要技術瓶頸。首先是隱藏全對全集合通信。Ulysses 序列平行化在自注意力層內產生兩個全對全通信用於交換序列分片和頭分片，初始狀態下這些通信完全暴露於關鍵路徑。通過將注意力頭分組並讓每組運行獨立的全對全-注意力-全對全模式，使得傳輸能隱藏在相鄰頭組的注意力計算後面，將通信足跡從計算流中減少約 5 倍。

其次是刪除稀疏注意力掩碼。超分辨率階段的稀疏注意力核心原本以 128 的倍數分塊序列，而遮罩按幀對齊，導致約五分之一的分塊邊界不齊。HeyGen 放寬分塊大小約束至 16 的倍數（硬體在序列維度上支持的最細 bf16 分塊），使每個分塊要麼完全在掩碼內，要麼完全在掩碼外。這樣掩碼謂詞、第二遍處理和填充全部消失，單次密集注意力內迴圈即可覆蓋整個掩碼。後續重構核心體，通過調整 softmax 歸約以減少暖迴圈的暫存器流量，將該階段性能從天花板的約一半提升至三分之二，最終接近 86%。

第三是消除 softmax 序列依賴。Flash 注意力的在線 softmax 在追蹤運行最大值時產生序列依賴。團隊用預計算的上界取代運行最大值，利用 Cauchy-Schwarz 不等式證明查詢的最大邏輯值受其範數與最大鍵範數乘積的約束。通過標量預取將預計算範數傳入核心，每行能在開始時推導其上界，消除重新縮放及其序列鏈。頭選擇性檢查，98-99% 的頭符合條件，其餘回退至標準在線路徑。此外，內鍵