2025 年，由 HyperAI 超神經主辦的 Meet AI Complier 技術沙龍已經行至第 7 期，在社區小伙伴和多位行業專家的支持下，我們在北京、上海、深圳等地建立了多個據點，為開發者和愛好者提供交流平臺，揭開創新技術的神秘面紗，直面一線開發者的應用反饋，共享技術落地的實戰經驗，聆聽多角度的創新思維。

關注微信公眾號「HyperAI 超神經」，后臺回復關鍵字「0705 AI 編譯器」，即可獲取確認授權的講師演講 PPT 。

在「Triton-distributed：原生 Python 編程實現高性能通信」主題演講中，來自字節跳動的 Seed Research Scientist 鄭思澤詳細解析了 Triton-distributed 在大模型訓練中的通信效率突破、跨平臺適配能力，以及如何通過 Python 編程實現通信與計算的深度融合。分享結束后，現場迅速進入提問高峰，圍繞 FLUX 框架、 Tile 編程模型、 AllGather 與 ReduceScatter 優化等細節展開的探討層出不窮，討論聚焦核心技術難點與實踐經驗，切實促進了理論與應用的結合。

HyperAI 超神經在不違原意的前提下，對鄭思澤老師的演講分享進行了整理匯總，以下為演講實錄。

分布式訓練的現實挑戰

在當前大模型迅速演進的背景下，無論是訓練還是推理，分布式系統已成為不可或缺的一環。我們也在這一方向上開展了編譯器層面的探索，并已開源該項目，命名為 Triton-Distributed 。

當前主流的硬件互聯方式包括 NVLink 、 PCIe 以及跨節點的網絡通信。在理想條件下，H100 的 NVLink 單向帶寬可以達到 450GB/s，但在國內大多數部署中，更常見的實際是 H800，其單向帶寬只有約 200GB/s，整體通信能力和拓撲復雜度大幅下降。我們在項目中遇到的一個明顯挑戰便是由于帶寬不足與通信拓撲不對稱帶來的系統性能瓶頸。

不同 GPU 集群的帶寬

針對此，早期的分布式優化往往依賴大量手動實現的通信算子，包括 Tensor 并行、 Pipeline 并行、 Data 并行等策略，均需要精心編寫底層通信邏輯。常見做法是調用如 NCCL 、 ROCm CCL 等通信庫，但這類方案往往缺乏通用性和可移植性，開發成本和維護代價都較高。

在分析現有系統瓶頸時，我們總結了 3 個關鍵事實：

Fact 1：硬件帶寬受限，通信延遲成為瓶頸

首先是基礎硬件條件帶來的限制。如果用 H100 訓練大模型的話，計算延遲往往顯著高于通信延遲，因此無需特別關注計算與通信的重疊調度。但在當前 H800 的環境下，通信延遲被明顯拉長。我們評估過，在某些場景下，近一半的訓練時間會被通信延遲空消耗掉，導致整體 MSU（Model Scale Utilization）顯著下降。若不進行通信與計算的 overlap（重疊）優化，系統將面臨嚴重的資源浪費問題。

硬件帶寬延遲

在中小規模下，這種損耗尚可接受；但一旦模型擴展到數千張卡級別，比如在 MegaScale 或 DeepSeek 的訓練實踐中，那么累計的資源損失將達到百萬甚至千萬美元級別，這對企業而言是非常現實的成本壓力。

推理場景同樣如此。 DeepSeek 早期推理部署使用了多達 320 張卡，盡管后續進行了壓縮和優化，但通信延遲依舊是分布式系統不可回避的核心問題。因此，如何在程序層面有效調度通信與計算、提升整體效率，成為我們必須正面應對的關鍵課題。

Fact 2: 通信開銷高，直接影響 MFU 表現

在當前的大模型訓練和推理中，通信開銷始終是一個主要瓶頸。我們觀察到，無論底層使用的是 NVLink 、 PCIe，還是不同代際的 GPU（比如 A100 、 H800），通信所占比例都非常高。特別是在國內的實際部署中，由于帶寬限制更明顯，通信延遲會直接拖慢整體效率。

對于大模型訓練來說，這種高頻的跨卡通信會顯著拉低系統的 MFU 。因此，優化通信開銷對于提升訓練與推理性能，都是非常關鍵的提升點，也是我們重點關注的方向之一。

不同 GPU 配置的通信與計算占比

Fact 3: 可編程性與性能之間的鴻溝

目前在分布式系統中，可編程性與性能之間仍存在較大鴻溝。過去我們更關注單卡編譯器的優化能力，比如如何在一張卡上發揮出色的性能；但當我們擴展到單機多卡，甚至跨節點的分布式系統后，情況就更加復雜。

一方面，分布式通信涉及大量底層技術細節，如 NCCL 、 MPI 、拓撲結構，且分散在各種專用庫中，使用門檻較高。很多時候，開發者需要手動實現通信邏輯、手動調度計算與同步，開發成本和出錯概率都很高。另一方面，如果有工具能自動處理分布式下復雜的通信調度和算子優化，就可以幫助開發者顯著降低開發門檻，提升分布式系統的可用性和可維護性，這正是我們在 Triton-Distributed 中希望解決的問題之一。

不同通信編程方式的特點

基于前面提到的 3 個現實問題，我們在 Triton-Distributed 中提出了 3 個核心方向：

首先，推動通信與計算的重疊（overlapping）機制。在通信開銷日益突出的分布式場景下，我們希望盡可能調度出計算與通信的并行窗口，提升系統整體效率。

其次，需要針對大模型的計算與通信模式進行深度融合與適配。比如在模型中常見的 AllReduce 、 Broadcast 等通信 pattern，我們嘗試將其與計算的 pattern 做融合，從而減少同步等待、壓縮執行路徑。

最后，我們認為，這些優化應通過編譯器完成，而不是依賴開發者手動編寫高度定制化的 CUDA 實現。讓分布式系統的開發更抽象、更高效，是我們努力的方向。

Triton-distributed 架構解析：原生 Python 實現高性能通信

我們希望在分布式訓練中實現 overlapping，但真正落地并不容易。概念上，overlapping 是指通過多個 stream 并發執行計算和通信，以掩蓋通信延遲。這在算子之間無依賴的場景下較為容易，但像 Tensor Parallel（TP）或 Expert Parallel（EP）中，必須先完成 AllGather 才能進行 GEMM，二者處于關鍵路徑，重疊難度很大。

目前常見方法包括：一是將任務劃分為多個 Micro-Batch，借助 Batch 間的獨立性實現 overlap；二是在單個 Batch 內以更細粒度（如 tile 粒度）進行切分，通過 kernel fusion 達到并行效果。我們在 Flux 中也探索了這類切分與調度機制，同時，大模型訓練中的通信模式高度復雜。例如 DeepSeek 在做 MoE 時需自定義 All-to-All 通信，以兼顧帶寬和負載均衡；又如低延遲推理和量化場景下，NCCL 等通用庫難以滿足性能要求，往往需要手寫通信 kernel，這些都提高了定制化成本。

因此，我們認為通信-計算融合的優化能力，應由編譯器層承擔，以應對復雜模型結構和多樣硬件環境，避免重復手工實現帶來的開發負擔。

兩層通信原語抽象

在我們的編譯器設計中，采用了兩層通信原語（primitives）抽象結構，以兼顧上層優化表達能力和底層部署的可落地性。

第一層是偏高層的原語，主要在 tile 粒度上完成計算調度，并提供面向通信的抽象接口。它以 rank 間的 push/get 操作作為通信抽象，并通過 tag 標識機制區分每一次通信行為，方便調度器追蹤數據流與依賴關系。

第二層則更貼近底層實現，采用了一套類似于 Open Shared Memory 標準（OpenSHMEM）的原語體系。這一層主要用于映射到現有的通信庫或硬件后端，實現真實的通信行為。

此外，在多 rank 的場景中，我們還需要引入 barrier 與 signal 控制機制，用于跨 rank 的同步。比如需要通知其他 rank 我方的數據已寫入完畢，或等待某個 rank 的數據準備就緒時，這類同步信號就非常關鍵。

編譯流程圖

編譯器架構與語義建模

在編譯棧方面，我們的整體流程仍然基于原始的 Triton 編譯框架。從源碼開始，Triton 會先將用戶代碼轉為抽象語法樹（AST），再翻譯成 Triton IR 。而在我們構建的 Triton-Distributed 中，我們對原有的 Triton IR 做了擴展，新增了一套面向分布式語義的 IR 層。這套分布式 IR 中，引入了對同步操作的語義建模，例如 wait 和 notify，用于描述 rank 之間的通信依賴關系；同時，我們還設計了一套面向 OpenSHMEM 的語義接口，以支持更底層的通信調用。

在實際代碼生成階段，這些語義可以映射為對底層通信庫的外部調用（external call）。我們通過 LLVM 中間層，直接將這些調用鏈接到 OpenSHMEM 提供的 bitcode 版本的庫（而非源碼），以實現跨 rank 的高效共享內存通信。這種方式繞過了 Triton 不支持源碼直接接入 external lib 的限制，使得共享內存相關的調用可以在編譯期順利完成符號解析與鏈接。

編譯棧圖

高層原語與底層執行的映射機制

在 Triton-distributed 中，我們設計了一套覆蓋高層抽象與底層控制的通信原語體系。以 consumer_tile_wait 為例，開發者只需聲明需要等待的 tile ID，系統會自動根據當前算子語義（如 AllGather）推導出通信目標的具體 rank 和 offset，完成同步邏輯。高層原語屏蔽了具體數據來源與信號傳遞的細節，提升了開發效率。

相比之下，底層原語則提供了更細粒度的控制能力。開發者需要手動指定 signal 指針、作用域（GPU 或 system）、內存語義（acquire 、 release 等）及預期值。這種機制雖然更復雜，但適用于對通信延遲和調度精準性要求極高的場景。

高層原語與底層原語

高層次的原語大致分為兩類：信號控制和數據控制。在信號控制的語義中，我們主要定義了 3 類角色：producer 、 consumer 和 peer，它們之間通過讀寫 signal 實現同步，類似于分布式通信中的握手機制。對于數據傳輸，Triton-distributed 提供了 push 與 pull 兩種原語，分別對應主動將數據發送到遠端卡，或從遠端拉取數據到本地卡。

所有底層通信原語均遵循 OpenSHMEM 標準，當前已支持 NVSHMEM 和 ROCSHMEM 。高層與底層原語之間具備明確的映射關系，編譯器負責將簡潔的接口自動轉換為底層的同步與傳輸指令。通過這套機制，Triton-distributed 既保留了通信調度的高性能能力，也大幅降低了分布式編程的復雜度。

在 Triton-distributed 中，高層通信原語（如 notify 和 wait）的設計目標是以簡潔語義描述跨卡同步需求，同時由編譯器負責將其翻譯為對應的底層執行邏輯。以 notify 為例，它與 wait 構成同步語義的一對：前者用于發送通知，后者用于等待數據準備完成。開發者只需指定 tile ID，系統即可根據算子類型與通信拓撲，自動推導出通信目標、信號偏移等底層細節。

具體的底層實現會因部署環境而異。例如在 8 卡 GPU 的場景中，這類同步可通過線程內的 _syncthreads() 與 atomic_dd 實現；在跨機部署中，則依賴于如 NVSHMEM 或 ROCSHMEM 提供的 signal_up 等原語完成等效操作。這些機制共同構成了高層語義與底層原語之間的映射關系，具有良好的通用性和可擴展性。

高層語義與底層原語之間的映射關系

以一個 GEMM ReduceScatter 的通信場景為例：假設系統中有 4 張 GPU，每個 tile 的目標位置由預先計算的元信息（如每個 rank 的 tile 分配量、 barrier 數量）決定。開發者只需在 Triton 編寫的 GEMM kernel 中添加一條 notify 語句，而 ReduceScatter kernel 端則用 wait 來同步接收數據。

整個過程可在 Python 中表達，也支持雙 stream 啟動的 kernel 模式，通信邏輯清晰且易于調度。這一機制不僅提高了跨卡通信編程的可表達性，也大幅降低了底層實現的復雜度，為分布式大模型的高效訓練與推理提供了強有力的基礎能力支持。

多維度的 Overlapping 優化：從調度機制到拓撲感知

雖然 Triton-distributed 已經提供了相對簡潔的高層通信原語接口，但在實際編寫和優化 kernel 的過程中，仍存在一定的技術門檻。我們觀察到，盡管原語設計具備良好的表達能力，但真正能夠靈活運用并深入優化的用戶仍然有限。本質上，通信優化仍是一項強依賴工程經驗和調度理解的工作，目前仍需由開發者手動控制。圍繞這個問題，我們總結出一些關鍵優化路徑，以下為 Triton-distributed 中的典型實現策略。

Push vs Pull：數據流向與 barrier 數控制約

在通信與計算的重疊優化中，Triton-distributed 提供了 push 和 pull 兩種數據傳輸方式。雖然它們在語義上僅僅是「主動發送」與「被動拉取」的方向差異，但在實際的分布式執行中，其性能表現和調度控制能力卻存在明顯不同。

以 barrier 數量為例，pull 模式通常需要設置 2 個 barrier：一個用于確保本地數據在被對方拉取前已經準備好，另一個則用于保護該數據在整個通信周期內不會被本地任務修改，從而防止數據不一致或讀寫沖突。而 push 模式則只需要在數據寫入遠端后設置一個 barrier，用以同步所有設備即可，整體控制更簡單。

但 pull 模式也有其優勢，它允許本地節點主動控制數據拉取順序，從而更精確地調度通信時機與計算重疊關系。當我們希望最大化 overlap 效果、實現通信與計算的并行性時，pull 提供了更高的靈活性。

總體來看，如果主要目標是提升 overlap，則推薦使用 pull；而在一些純通信任務中，如單獨的 AllGather 或 ReduceScatter kernel，push 模式因其實現簡潔、開銷更小而更為常見。

Push 模式與 Pull 模式流程圖

Swizzling 調度：按數據局部性動態調整順序

通信與計算的重疊不僅依賴于原語選擇，還與調度策略密切相關。其中，Swizzling 是一種基于拓撲感知的調度優化手段，旨在減少跨卡計算過程中的執行空閑。在分布式視角下，可以將每張 GPU 卡視為一個獨立的執行單元。由于每張卡初始持有的數據片段不同，若所有卡從相同 tile 索引開始計算，部分 rank 將不得不等待數據就緒，導致執行階段出現長時間的空閑，從而拉低整體計算效率。

Swizzling 的核心思想是：根據每張卡本地已有數據的位置動態調整起始計算偏移。例如，在 AllGather 場景中，每張卡可以優先處理自身持有的數據，同時發起對遠端 tile 的拉取，實現通信與計算的并發調度。若所有卡一律從 tile 0 開始處理，只有 rank 0 能立即開始計算，其余 rank 則將因等待數據而產生串行延遲。

更復雜的情形如跨機 ReduceScatter 場景中，Swizzling 策略還需結合網絡拓撲進行設計。以兩臺節點（node）為例，合理的調度方式是：優先計算對方節點所需的數據，盡早觸發跨機 point-to-point 通信；而在傳輸過程中，再并行計算本地節點所需數據，最大化通信與計算的 overlap 效果。

目前，這類調度優化仍由編程者控制，以避免編譯器在通用優化中犧牲關鍵性能路徑。我們也意識到，理解 Swizzling 等細節對開發者有一定門檻。未來，我們希望通過提供更多實際案例和模板代碼，幫助開發者更快掌握分布式算子開發模式，逐步構建起開放、高效的 Triton-distributed 編程生態。

Swizzling 優化流程

非完美分塊調度：跨 rank tile 的處理優先級

在實際的大模型訓練與推理場景中，算子的輸入 shape 往往并不規整，尤其是在 token 長度不固定的情況下，tile 分塊也難以保持整齊劃一。這種非完美分塊（Imperfect Tiling）會導致部分 tile 橫跨多個 rank，即同一個 tile 的數據分布在多個設備上，增加了調度與同步的復雜性。

以 AllGather GEMM 為例，假設某個 tile 同時包含了本地和遠端的數據。如果從這個 tile 開始計算，則必須等待遠端數據先完成傳輸，進而引入額外的 bubble，影響整體計算的并行性。更優的做法是：跳過這個跨 rank tile，優先處理完全本地可用的數據，將等待遠端輸入的 tile 調度至最后執行，從而實現通信與計算的最大重疊。

而在 ReduceScatter 場景中，調度順序則應反向處理。由于跨 rank tile 的計算結果需要盡早發送給遠端，最佳策略是：優先處理那些被遠端節點依賴的 tile，以便盡早完成跨機數據發送，減少遠端的依賴。

非完美分塊調度

MoE 下的 Dynamic Sorting 策略

在 MoE（Mixture-of-Experts）模型中，token 需要根據路由結果被分發至多個 expert，通常伴隨 All-to-All 通信與 Group GEMM 計算。為了提升通信與計算的重疊效率，Triton-distributed 引入了 Dynamic Sorting，按計算任務對通信數據的依賴強度進行分階段調度，優先處理數據依賴較少的部分。

這種排序方式確保了每一階段的計算都能以盡可能低的通信阻塞開始，從而在 All-to-All 與 Group GEMM 之間實現更好的 overlap 效果。整體調度從數據依賴最少的 tile 開始，逐步擴展至依賴復雜的數據塊，最大程度提升了執行并發性。

Dynamic Sorting

面向硬件的通信加速

Triton-distributed 還支持結合特定硬件能力進行通信優化，尤其是在使用 NVSwitch 架構時，可利用其內置的 SHARP Accelerator 執行低延遲的通信計算。該模塊可在交換芯片內完成如 Broadcast 、 AllReduce 等操作，實現數據在傳輸路徑中的聚合加速，減少延遲與帶寬消耗。相關指令已集成進 Triton-distributed，具備相應硬件的用戶可直接調用，構建更高效的通信 kernel 。

AOT 編譯優化：降低推理延遲開銷

Triton-distributed 引入了 AOT（Ahead-of-Time，提前編譯）機制，專門針對推理場景中對延遲極度敏感的需求進行優化。 Triton 默認采用 JIT（Just-In-Time compilation，即時編譯）編譯方式，函數首次執行時存在顯著的編譯與緩存開銷。

AOT 機制則允許用戶在運行前將函數預編譯為字節碼，推理階段直接加載執行，避免了 JIT 編譯過程，從而有效降低了編譯及緩存帶來的延遲。基于此，Triton-distributed 對 AOT 機制進行了擴展，現已支持分布式環境中的 AOT 編譯與部署，進一步提升了分布式推理的性能表現。

性能實測與案例復現

我們對 Triton-distributed 在多平臺、多任務場景下的性能進行了全面測試，涵蓋 NVIDIA H800 、 AMD GPU 、 8 卡 GPU 與跨機集群，并對比了 PyTorch 、 Flux 等主流分布式實現方案。

在 8 卡 GPU 上，Triton-distributed 在 AG GEMM 和 GEMM RS 任務中相較 PyTorch 實現有顯著加速，相比手工優化的 Flux 方案也取得更優性能，得益于 Swizzling 調度、通信 offload 和 AOT 編譯等多重優化。同時在 AMD 平臺上對比 PyTorch + RCCL 的組合，雖然整體加速幅度略小，但同樣取得顯著優化，限制主要來自測試硬件算力偏弱和非 switch 拓撲。

在 AllReduce 任務中，Triton-distributed 在從小到大的多種消息尺寸下，在我們測試的硬件配置中，相比 NCCL 均有明顯加速，平均加速約 1.6 倍。在 Attention 場景中，我們主要測試了 gather-KV 類型的 attention 操作。相較于 PyTorch Touch 的原生實現，Triton-distributed 在 8 卡 GPU 上的性能可達約 5 倍提升；同時也優于開源的 Ring Attention 實現，提升幅度約為 2 倍。

跨機測試方面，AG GEMM 提速 1.3 倍，GEMM RS 提速 1.4 倍，表現略低于 Flux，但在 shape 靈活性和可擴展性上更具優勢。我們還測試了高速推理場景下的單 token decoding，在 1M token context 下延遲可控制在 20–30 微秒，兼容 NVLink 與 PCIe 。

此外，我們對 DeepEP 中的分布式調度邏輯進行了功能復現，主要對齊其 All-to-All 路由與上下文分發策略。在 64 卡以內的場景下，Triton-distributed 的性能與其基本持平，部分配置下略有提升。

最后，我們還提供了基于 Qwen-32B 的 prefill 與 decode Demo，支持在 8 卡 GPU 上部署運行，實測可獲得約 1.2 倍的推理加速效果。

打造開放的分布式編譯生態

目前我們正面臨定制化 overlapping 場景的挑戰，過去主要依賴手工優化解決，工作量大且成本高。為此，我們提出并開源了分布式的 Triton-distributed 框架。雖然它是基于 Triton 實現的，但其實不論各家公司使用何種編譯器，或者底層通信庫如何，都能將其集成進來，打造一個開放的分布式生態。

在國內乃至全球，這一領域仍較為空白。我們希望借助社區的力量，吸引更多開發者參與進來，無論是在語法設計、性能優化，還是支持更多類型的硬件設備方面，共同推動技術進步。最后，我們取得了不錯的性能表現，相關示例也全部開源，歡迎大家積極提 issue 交流，也期待更多小伙伴加入，共創未來！