CoCoDiff: Optimizing Collective Communications for Distributed Diffusion Transformer Inference Under Ulysses Sequence Parallelism

分散式 Diffusion Transformer (DiT) 在生成高解析度影像時，Ulysses 序列平行的集合通訊往往佔據超過 80% 的推理時間。最新推出的 CoCoDiff 推理引擎透過重構通訊架構與計算排程，在 Aurora 超級電腦上實現了平均 3.6 倍、最高 8.4 倍的推理加速。

解析 Ulysses 序列平行與硬體通訊瓶頸

DiT 模型如 Stable Diffusion 3 (SD3)、FLUX 與 Sora 已成為現代生成式 AI 的核心。隨著影像解析度提升，模型參數量與序列長度呈二次方成長，跨多個 GPU 的分散式推理成為必備基礎設施。目前業界最主流的擴展方法是 Ulysses sequence parallelism（Ulysses 序列平行：將輸入序列分片至多個 GPU 處理），藉此打破單一硬體的記憶體限制。

然而，Ulysses 架構在每一次注意力層計算前後，都必須執行全對全（all-to-all）的集合通訊來重新分配 Q/K/V（查詢、鍵、值張量）。團隊在 Aurora 超級電腦上測試 20B 參數的 Qwen-Image 模型發現，通訊開銷極其高昂。在解析度 $768\times 768$ 時，通訊佔整體推理時間的 64%；當解析度提升至 $2304\times 2304$ 時，通訊佔比更突破 80%。

這項通訊瓶頸難以單純透過隱藏延遲來解決。注意力機制的計算必須等待通訊完全結束才能開始，且 Diffusion 模型是由數十個循序的降噪步驟組成，形成嚴格的前向依賴。更棘手的是，以 Aurora 節點為例，GPU 內部的 tile（運算區塊）頻寬高達 185 GB/s，但跨 GPU 的互連頻寬僅有 15 GB/s，兩者相差逾 12 倍，導致傳統的單一集體通訊效率極低。

TAPA 拓撲感知拆解：將集體通訊拆為硬體對齊兩階段

為了解決硬體架構頻寬極度不對稱的問題，CoCoDiff 首先引入了 TAPA (Tile-Aware Parallel All-to-all) 機制。傳統的集合通訊函式庫在處理多階層拓撲時，往往未考慮底層運算資料的語意，導致跨 GPU 與 GPU 內部的通訊同時觸發，塞滿低頻寬通道。

TAPA 選擇將每一次的 all-to-all 通訊強制拆分為兩個明確的階段。第一階段（Phase 1）限定在單一 GPU 內部進行，讓同一張 GPU 內的兩個 tile 利用 185 GB/s 的高頻寬完成資料交換。在此階段結束後，每個 tile 已經擁有分配給該 GPU 的完整資料備份。

第二階段（Phase 2）才啟動跨 GPU 的資料交換。TAPA 將此階段設計為兩個平行的 6 節點環狀通訊，完全對齊 Aurora 內部的 Xe Link 互連拓撲。這種拆解策略避免了跨環路的多節點跳轉路由，最大化利用硬體的聚合頻寬，為後續的非同步計算創造了操作空間。

V-First 計算排程：利用 Q/K 運算時間隱藏 V 張量通訊

成功將通訊拆解後，CoCoDiff 的第二個創新是 V-First 排程。在標準 Transformer 架構中，Q、K、V 張量皆透過線性投影產生，因此通訊通常會等到三者皆計算完畢後才一併啟動。但研究團隊觀察到現代 DiT 架構存在明顯的運算不對稱性。

除了基本的線性投影，模型還會額外對 Q 和 K 執行 RMSNorm（均方根正規化，維持數值穩定）以及 RoPE（旋轉位置編碼，處理相對位置）。相反地，V 只需要參與後續的加權求和，不需要經過正規化或位置編碼。這種差異使得 V 的計算能大幅提前完成。

V-First 排程便利用這段時間差，一旦 V 的投影計算完畢，立刻在背景啟動 V 的 Phase 1 通訊。在 Aurora 上的實測顯示，Q 與 K 的額外處理大約需要 1.0 毫秒，而 V 的第一階段通訊僅需 0.3 毫秒。透過這套排程，V 的初期通訊延遲被完全隱藏在 Q 與 K 的計算時間背後，打破了通訊與計算無法重疊的限制。

V-Major 選擇性通訊：過濾降噪冗餘張量減少跨節點流量

為解決跨 GPU（Phase 2）低頻寬通道的壅塞，CoCoDiff 導入了 V-Major 選擇性通訊。Diffusion 模型的影像生成需要經歷數十次反覆疊代的降噪步驟，相鄰時間步產生的中間張量具備高度的時間冗餘性，亦即數值變化極小。

V-Major 首度將這種冗餘性應用於分散式通訊減量。在每一個降噪步驟中，系統會將最新計算出的 V 向量與前一步驟的 V 向量進行 L1 距離比較。只有當向量變化超過特定閾值時，該組對應的 Q、K、V 投影才會被標記為活躍（active）並投入跨 GPU 通訊；變化微小的部分則直接讀取快取資料。

針對不同降噪階段的特性，CoCoDiff 設計了動態快取比例。在早期需要全域結構更新的步驟，系統會傳送較多張量；到了後期微調細節的階段，則大幅提高快取命中率。為了防止近似誤差跨層累積，系統設定每 10 個步驟進行一次完整的全量通訊，確保生成的影像品質不會因為快取而崩壞。

96 晶片實測：最高 8.4 倍加速且維持高保真影像品質

研究團隊將 CoCoDiff 整合至開源套件 xDiT 與 Diffusers 中，並在 Aurora 平台上擴展至最高 8 個節點（包含 96 個 Intel GPU tiles） 進行驗證。測試涵蓋了 FLUX.1-dev、Qwen-Image、SD3.5 等多款主流大模型，並跨越三種高解析度設定。

實測結果顯示，結合三大機制的 CoCoDiff 在各項配置下均展現顯著效能。相較於原始的 Flat Ulysses 基準，CoCoDiff 提供平均 3.6 倍的端到端加速。在處理 SD3.5 模型、$2304\times 2304$ 解析度與 4 節點的極端負載下，加速比飆升至 8.4 倍。其中，V-Major 的通訊減量在多節點擴展時成為突破效能天花板的關鍵。

在生成品質評估方面，團隊使用美國阿貢國家實驗室的 X 光微電腦斷層掃描影像進行修復測試。數據證實，採用 CoCoDiff 的輸出在 PSNR 與 SSIM 指標上與原始基線幾乎無異，成功在嚴苛的硬體限制下，為大規模 DiT 影像生成找出了效能與精度的最佳平衡。

透過挖掘模型架構的計算不對稱性與降噪冗餘性，CoCoDiff 證明了演算法層級的排程與快取，能有效突破底層硬體頻寬的物理限制。