2025 年，由 HyperAI 超神經(jīng)主辦的 Meet AI Complier 技術(shù)沙龍已經(jīng)行至第 7 期，在社區(qū)小伙伴和多位行業(yè)專家的支持下，我們?cè)诒本?、上海、深圳等地建立了多個(gè)據(jù)點(diǎn)，為開發(fā)者和愛好者提供交流平臺(tái)，揭開創(chuàng)新技術(shù)的神秘面紗，直面一線開發(fā)者的應(yīng)用反饋，共享技術(shù)落地的實(shí)戰(zhàn)經(jīng)驗(yàn)，聆聽多角度的創(chuàng)新思維。

關(guān)注微信公眾號(hào)「HyperAI 超神經(jīng)」，后臺(tái)回復(fù)關(guān)鍵字「0705 AI 編譯器」，即可獲取確認(rèn)授權(quán)的講師演講 PPT 。

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

HyperAI 超神經(jīng)在不違原意的前提下，對(duì)鄭思澤老師的演講分享進(jìn)行了整理匯總，以下為演講實(shí)錄。

分布式訓(xùn)練的現(xiàn)實(shí)挑戰(zhàn)

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

當(dāng)前主流的硬件互聯(lián)方式包括 NVLink 、 PCIe 以及跨節(jié)點(diǎn)的網(wǎng)絡(luò)通信。在理想條件下，H100 的 NVLink 單向帶寬可以達(dá)到 450GB/s，但在國(guó)內(nèi)大多數(shù)部署中，更常見的實(shí)際是 H800，其單向帶寬只有約 200GB/s，整體通信能力和拓?fù)鋸?fù)雜度大幅下降。我們?cè)陧?xiàng)目中遇到的一個(gè)明顯挑戰(zhàn)便是由于帶寬不足與通信拓?fù)洳粚?duì)稱帶來的系統(tǒng)性能瓶頸。

不同 GPU 集群的帶寬

針對(duì)此，早期的分布式優(yōu)化往往依賴大量手動(dòng)實(shí)現(xiàn)的通信算子，包括 Tensor 并行、 Pipeline 并行、 Data 并行等策略，均需要精心編寫底層通信邏輯。常見做法是調(diào)用如 NCCL 、 ROCm CCL 等通信庫(kù)，但這類方案往往缺乏通用性和可移植性，開發(fā)成本和維護(hù)代價(jià)都較高。

在分析現(xiàn)有系統(tǒng)瓶頸時(shí)，我們總結(jié)了 3 個(gè)關(guān)鍵事實(shí)：

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

首先是基礎(chǔ)硬件條件帶來的限制。如果用 H100 訓(xùn)練大模型的話，計(jì)算延遲往往顯著高于通信延遲，因此無需特別關(guān)注計(jì)算與通信的重疊調(diào)度。但在當(dāng)前 H800 的環(huán)境下，通信延遲被明顯拉長(zhǎng)。我們?cè)u(píng)估過，在某些場(chǎng)景下，近一半的訓(xùn)練時(shí)間會(huì)被通信延遲空消耗掉，導(dǎo)致整體 MSU（Model Scale Utilization）顯著下降。若不進(jìn)行通信與計(jì)算的 overlap（重疊）優(yōu)化，系統(tǒng)將面臨嚴(yán)重的資源浪費(fèi)問題。

硬件帶寬延遲

在中小規(guī)模下，這種損耗尚可接受；但一旦模型擴(kuò)展到數(shù)千張卡級(jí)別，比如在 MegaScale 或 DeepSeek 的訓(xùn)練實(shí)踐中，那么累計(jì)的資源損失將達(dá)到百萬甚至千萬美元級(jí)別，這對(duì)企業(yè)而言是非常現(xiàn)實(shí)的成本壓力。

推理場(chǎng)景同樣如此。 DeepSeek 早期推理部署使用了多達(dá) 320 張卡，盡管后續(xù)進(jìn)行了壓縮和優(yōu)化，但通信延遲依舊是分布式系統(tǒng)不可回避的核心問題。因此，如何在程序?qū)用嬗行д{(diào)度通信與計(jì)算、提升整體效率，成為我們必須正面應(yīng)對(duì)的關(guān)鍵課題。

Fact 2: 通信開銷高，直接影響 MFU 表現(xiàn)

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

對(duì)于大模型訓(xùn)練來說，這種高頻的跨卡通信會(huì)顯著拉低系統(tǒng)的 MFU 。因此，優(yōu)化通信開銷對(duì)于提升訓(xùn)練與推理性能，都是非常關(guān)鍵的提升點(diǎn)，也是我們重點(diǎn)關(guān)注的方向之一。

不同 GPU 配置的通信與計(jì)算占比

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

目前在分布式系統(tǒng)中，可編程性與性能之間仍存在較大鴻溝。過去我們更關(guān)注單卡編譯器的優(yōu)化能力，比如如何在一張卡上發(fā)揮出色的性能；但當(dāng)我們擴(kuò)展到單機(jī)多卡，甚至跨節(jié)點(diǎn)的分布式系統(tǒng)后，情況就更加復(fù)雜。

一方面，分布式通信涉及大量底層技術(shù)細(xì)節(jié)，如 NCCL 、 MPI 、拓?fù)浣Y(jié)構(gòu)，且分散在各種專用庫(kù)中，使用門檻較高。很多時(shí)候，開發(fā)者需要手動(dòng)實(shí)現(xiàn)通信邏輯、手動(dòng)調(diào)度計(jì)算與同步，開發(fā)成本和出錯(cuò)概率都很高。另一方面，如果有工具能自動(dòng)處理分布式下復(fù)雜的通信調(diào)度和算子優(yōu)化，就可以幫助開發(fā)者顯著降低開發(fā)門檻，提升分布式系統(tǒng)的可用性和可維護(hù)性，這正是我們?cè)?Triton-Distributed 中希望解決的問題之一。

不同通信編程方式的特點(diǎn)

基于前面提到的 3 個(gè)現(xiàn)實(shí)問題，我們?cè)?Triton-Distributed 中提出了 3 個(gè)核心方向：

首先，推動(dòng)通信與計(jì)算的重疊（overlapping）機(jī)制。在通信開銷日益突出的分布式場(chǎng)景下，我們希望盡可能調(diào)度出計(jì)算與通信的并行窗口，提升系統(tǒng)整體效率。

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

最后，我們認(rèn)為，這些優(yōu)化應(yīng)通過編譯器完成，而不是依賴開發(fā)者手動(dòng)編寫高度定制化的 CUDA 實(shí)現(xiàn)。讓分布式系統(tǒng)的開發(fā)更抽象、更高效，是我們努力的方向。

Triton-distributed 架構(gòu)解析：原生 Python 實(shí)現(xiàn)高性能通信

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

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

因此，我們認(rèn)為通信-計(jì)算融合的優(yōu)化能力，應(yīng)由編譯器層承擔(dān)，以應(yīng)對(duì)復(fù)雜模型結(jié)構(gòu)和多樣硬件環(huán)境，避免重復(fù)手工實(shí)現(xiàn)帶來的開發(fā)負(fù)擔(dān)。

兩層通信原語(yǔ)抽象

在我們的編譯器設(shè)計(jì)中，采用了兩層通信原語(yǔ)（primitives）抽象結(jié)構(gòu)，以兼顧上層優(yōu)化表達(dá)能力和底層部署的可落地性。

第一層是偏高層的原語(yǔ)，主要在 tile 粒度上完成計(jì)算調(diào)度，并提供面向通信的抽象接口。它以 rank 間的 push/get 操作作為通信抽象，并通過 tag 標(biāo)識(shí)機(jī)制區(qū)分每一次通信行為，方便調(diào)度器追蹤數(shù)據(jù)流與依賴關(guān)系。

第二層則更貼近底層實(shí)現(xiàn)，采用了一套類似于 Open Shared Memory 標(biāo)準(zhǔn)（OpenSHMEM）的原語(yǔ)體系。這一層主要用于映射到現(xiàn)有的通信庫(kù)或硬件后端，實(shí)現(xiàn)真實(shí)的通信行為。

此外，在多 rank 的場(chǎng)景中，我們還需要引入 barrier 與 signal 控制機(jī)制，用于跨 rank 的同步。比如需要通知其他 rank 我方的數(shù)據(jù)已寫入完畢，或等待某個(gè) rank 的數(shù)據(jù)準(zhǔn)備就緒時(shí)，這類同步信號(hào)就非常關(guān)鍵。

編譯流程圖

編譯器架構(gòu)與語(yǔ)義建模

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

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

編譯棧圖

高層原語(yǔ)與底層執(zhí)行的映射機(jī)制

在 Triton-distributed 中，我們?cè)O(shè)計(jì)了一套覆蓋高層抽象與底層控制的通信原語(yǔ)體系。以 consumer_tile_wait 為例，開發(fā)者只需聲明需要等待的 tile ID，系統(tǒng)會(huì)自動(dòng)根據(jù)當(dāng)前算子語(yǔ)義（如 AllGather）推導(dǎo)出通信目標(biāo)的具體 rank 和 offset，完成同步邏輯。高層原語(yǔ)屏蔽了具體數(shù)據(jù)來源與信號(hào)傳遞的細(xì)節(jié)，提升了開發(fā)效率。

相比之下，底層原語(yǔ)則提供了更細(xì)粒度的控制能力。開發(fā)者需要手動(dòng)指定 signal 指針、作用域（GPU 或 system）、內(nèi)存語(yǔ)義（acquire 、 release 等）及預(yù)期值。這種機(jī)制雖然更復(fù)雜，但適用于對(duì)通信延遲和調(diào)度精準(zhǔn)性要求極高的場(chǎng)景。

高層原語(yǔ)與底層原語(yǔ)

高層次的原語(yǔ)大致分為兩類：信號(hào)控制和數(shù)據(jù)控制。在信號(hào)控制的語(yǔ)義中，我們主要定義了 3 類角色：producer 、 consumer 和 peer，它們之間通過讀寫 signal 實(shí)現(xiàn)同步，類似于分布式通信中的握手機(jī)制。對(duì)于數(shù)據(jù)傳輸，Triton-distributed 提供了 push 與 pull 兩種原語(yǔ)，分別對(duì)應(yīng)主動(dòng)將數(shù)據(jù)發(fā)送到遠(yuǎn)端卡，或從遠(yuǎn)端拉取數(shù)據(jù)到本地卡。

所有底層通信原語(yǔ)均遵循 OpenSHMEM 標(biāo)準(zhǔn)，當(dāng)前已支持 NVSHMEM 和 ROCSHMEM 。高層與底層原語(yǔ)之間具備明確的映射關(guān)系，編譯器負(fù)責(zé)將簡(jiǎn)潔的接口自動(dòng)轉(zhuǎn)換為底層的同步與傳輸指令。通過這套機(jī)制，Triton-distributed 既保留了通信調(diào)度的高性能能力，也大幅降低了分布式編程的復(fù)雜度。

在 Triton-distributed 中，高層通信原語(yǔ)（如 notify 和 wait）的設(shè)計(jì)目標(biāo)是以簡(jiǎn)潔語(yǔ)義描述跨卡同步需求，同時(shí)由編譯器負(fù)責(zé)將其翻譯為對(duì)應(yīng)的底層執(zhí)行邏輯。以 notify 為例，它與 wait 構(gòu)成同步語(yǔ)義的一對(duì)：前者用于發(fā)送通知，后者用于等待數(shù)據(jù)準(zhǔn)備完成。開發(fā)者只需指定 tile ID，系統(tǒng)即可根據(jù)算子類型與通信拓?fù)?，自?dòng)推導(dǎo)出通信目標(biāo)、信號(hào)偏移等底層細(xì)節(jié)。

具體的底層實(shí)現(xiàn)會(huì)因部署環(huán)境而異。例如在 8 卡 GPU 的場(chǎng)景中，這類同步可通過線程內(nèi)的 _syncthreads() 與 atomic_dd 實(shí)現(xiàn)；在跨機(jī)部署中，則依賴于如 NVSHMEM 或 ROCSHMEM 提供的 signal_up 等原語(yǔ)完成等效操作。這些機(jī)制共同構(gòu)成了高層語(yǔ)義與底層原語(yǔ)之間的映射關(guān)系，具有良好的通用性和可擴(kuò)展性。

高層語(yǔ)義與底層原語(yǔ)之間的映射關(guān)系

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

整個(gè)過程可在 Python 中表達(dá)，也支持雙 stream 啟動(dòng)的 kernel 模式，通信邏輯清晰且易于調(diào)度。這一機(jī)制不僅提高了跨卡通信編程的可表達(dá)性，也大幅降低了底層實(shí)現(xiàn)的復(fù)雜度，為分布式大模型的高效訓(xùn)練與推理提供了強(qiáng)有力的基礎(chǔ)能力支持。

多維度的 Overlapping 優(yōu)化：從調(diào)度機(jī)制到拓?fù)涓兄?/h2>

雖然 Triton-distributed 已經(jīng)提供了相對(duì)簡(jiǎn)潔的高層通信原語(yǔ)接口，但在實(shí)際編寫和優(yōu)化 kernel 的過程中，仍存在一定的技術(shù)門檻。我們觀察到，盡管原語(yǔ)設(shè)計(jì)具備良好的表達(dá)能力，但真正能夠靈活運(yùn)用并深入優(yōu)化的用戶仍然有限。本質(zhì)上，通信優(yōu)化仍是一項(xiàng)強(qiáng)依賴工程經(jīng)驗(yàn)和調(diào)度理解的工作，目前仍需由開發(fā)者手動(dòng)控制。圍繞這個(gè)問題，我們總結(jié)出一些關(guān)鍵優(yōu)化路徑，以下為 Triton-distributed 中的典型實(shí)現(xiàn)策略。

Push vs Pull：數(shù)據(jù)流向與 barrier 數(shù)控制約

在通信與計(jì)算的重疊優(yōu)化中，Triton-distributed 提供了 push 和 pull 兩種數(shù)據(jù)傳輸方式。雖然它們?cè)谡Z(yǔ)義上僅僅是「主動(dòng)發(fā)送」與「被動(dòng)拉取」的方向差異，但在實(shí)際的分布式執(zhí)行中，其性能表現(xiàn)和調(diào)度控制能力卻存在明顯不同。

以 barrier 數(shù)量為例，pull 模式通常需要設(shè)置 2 個(gè) barrier：一個(gè)用于確保本地?cái)?shù)據(jù)在被對(duì)方拉取前已經(jīng)準(zhǔn)備好，另一個(gè)則用于保護(hù)該數(shù)據(jù)在整個(gè)通信周期內(nèi)不會(huì)被本地任務(wù)修改，從而防止數(shù)據(jù)不一致或讀寫沖突。而 push 模式則只需要在數(shù)據(jù)寫入遠(yuǎn)端后設(shè)置一個(gè) barrier，用以同步所有設(shè)備即可，整體控制更簡(jiǎn)單。

但 pull 模式也有其優(yōu)勢(shì)，它允許本地節(jié)點(diǎn)主動(dòng)控制數(shù)據(jù)拉取順序，從而更精確地調(diào)度通信時(shí)機(jī)與計(jì)算重疊關(guān)系。當(dāng)我們希望最大化 overlap 效果、實(shí)現(xiàn)通信與計(jì)算的并行性時(shí)，pull 提供了更高的靈活性。

總體來看，如果主要目標(biāo)是提升 overlap，則推薦使用 pull；而在一些純通信任務(wù)中，如單獨(dú)的 AllGather 或 ReduceScatter kernel，push 模式因其實(shí)現(xiàn)簡(jiǎn)潔、開銷更小而更為常見。

Push 模式與 Pull 模式流程圖

Swizzling 調(diào)度：按數(shù)據(jù)局部性動(dòng)態(tài)調(diào)整順序

通信與計(jì)算的重疊不僅依賴于原語(yǔ)選擇，還與調(diào)度策略密切相關(guān)。其中，Swizzling 是一種基于拓?fù)涓兄恼{(diào)度優(yōu)化手段，旨在減少跨卡計(jì)算過程中的執(zhí)行空閑。在分布式視角下，可以將每張 GPU 卡視為一個(gè)獨(dú)立的執(zhí)行單元。由于每張卡初始持有的數(shù)據(jù)片段不同，若所有卡從相同 tile 索引開始計(jì)算，部分 rank 將不得不等待數(shù)據(jù)就緒，導(dǎo)致執(zhí)行階段出現(xiàn)長(zhǎng)時(shí)間的空閑，從而拉低整體計(jì)算效率。

Swizzling 的核心思想是：根據(jù)每張卡本地已有數(shù)據(jù)的位置動(dòng)態(tài)調(diào)整起始計(jì)算偏移。例如，在 AllGather 場(chǎng)景中，每張卡可以優(yōu)先處理自身持有的數(shù)據(jù)，同時(shí)發(fā)起對(duì)遠(yuǎn)端 tile 的拉取，實(shí)現(xiàn)通信與計(jì)算的并發(fā)調(diào)度。若所有卡一律從 tile 0 開始處理，只有 rank 0 能立即開始計(jì)算，其余 rank 則將因等待數(shù)據(jù)而產(chǎn)生串行延遲。

更復(fù)雜的情形如跨機(jī) ReduceScatter 場(chǎng)景中，Swizzling 策略還需結(jié)合網(wǎng)絡(luò)拓?fù)溥M(jìn)行設(shè)計(jì)。以兩臺(tái)節(jié)點(diǎn)（node）為例，合理的調(diào)度方式是：優(yōu)先計(jì)算對(duì)方節(jié)點(diǎn)所需的數(shù)據(jù)，盡早觸發(fā)跨機(jī) point-to-point 通信；而在傳輸過程中，再并行計(jì)算本地節(jié)點(diǎn)所需數(shù)據(jù)，最大化通信與計(jì)算的 overlap 效果。

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

Swizzling 優(yōu)化流程

非完美分塊調(diào)度：跨 rank tile 的處理優(yōu)先級(jí)

在實(shí)際的大模型訓(xùn)練與推理場(chǎng)景中，算子的輸入 shape 往往并不規(guī)整，尤其是在 token 長(zhǎng)度不固定的情況下，tile 分塊也難以保持整齊劃一。這種非完美分塊（Imperfect Tiling）會(huì)導(dǎo)致部分 tile 橫跨多個(gè) rank，即同一個(gè) tile 的數(shù)據(jù)分布在多個(gè)設(shè)備上，增加了調(diào)度與同步的復(fù)雜性。

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

而在 ReduceScatter 場(chǎng)景中，調(diào)度順序則應(yīng)反向處理。由于跨 rank tile 的計(jì)算結(jié)果需要盡早發(fā)送給遠(yuǎn)端，最佳策略是：優(yōu)先處理那些被遠(yuǎn)端節(jié)點(diǎn)依賴的 tile，以便盡早完成跨機(jī)數(shù)據(jù)發(fā)送，減少遠(yuǎn)端的依賴。

非完美分塊調(diào)度

MoE 下的 Dynamic Sorting 策略

在 MoE（Mixture-of-Experts）模型中，token 需要根據(jù)路由結(jié)果被分發(fā)至多個(gè) expert，通常伴隨 All-to-All 通信與 Group GEMM 計(jì)算。為了提升通信與計(jì)算的重疊效率，Triton-distributed 引入了 Dynamic Sorting，按計(jì)算任務(wù)對(duì)通信數(shù)據(jù)的依賴強(qiáng)度進(jìn)行分階段調(diào)度，優(yōu)先處理數(shù)據(jù)依賴較少的部分。

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

Dynamic Sorting

面向硬件的通信加速

Triton-distributed 還支持結(jié)合特定硬件能力進(jìn)行通信優(yōu)化，尤其是在使用 NVSwitch 架構(gòu)時(shí)，可利用其內(nèi)置的 SHARP Accelerator 執(zhí)行低延遲的通信計(jì)算。該模塊可在交換芯片內(nèi)完成如 Broadcast 、 AllReduce 等操作，實(shí)現(xiàn)數(shù)據(jù)在傳輸路徑中的聚合加速，減少延遲與帶寬消耗。相關(guān)指令已集成進(jìn) Triton-distributed，具備相應(yīng)硬件的用戶可直接調(diào)用，構(gòu)建更高效的通信 kernel 。

AOT 編譯優(yōu)化：降低推理延遲開銷

Triton-distributed 引入了 AOT（Ahead-of-Time，提前編譯）機(jī)制，專門針對(duì)推理場(chǎng)景中對(duì)延遲極度敏感的需求進(jìn)行優(yōu)化。 Triton 默認(rèn)采用 JIT（Just-In-Time compilation，即時(shí)編譯）編譯方式，函數(shù)首次執(zhí)行時(shí)存在顯著的編譯與緩存開銷。

AOT 機(jī)制則允許用戶在運(yùn)行前將函數(shù)預(yù)編譯為字節(jié)碼，推理階段直接加載執(zhí)行，避免了 JIT 編譯過程，從而有效降低了編譯及緩存帶來的延遲?；诖?，Triton-distributed 對(duì) AOT 機(jī)制進(jìn)行了擴(kuò)展，現(xiàn)已支持分布式環(huán)境中的 AOT 編譯與部署，進(jìn)一步提升了分布式推理的性能表現(xiàn)。

性能實(shí)測(cè)與案例復(fù)現(xiàn)

我們對(duì) Triton-distributed 在多平臺(tái)、多任務(wù)場(chǎng)景下的性能進(jìn)行了全面測(cè)試，涵蓋 NVIDIA H800 、 AMD GPU 、 8 卡 GPU 與跨機(jī)集群，并對(duì)比了 PyTorch 、 Flux 等主流分布式實(shí)現(xiàn)方案。

在 8 卡 GPU 上，Triton-distributed 在 AG GEMM 和 GEMM RS 任務(wù)中相較 PyTorch 實(shí)現(xiàn)有顯著加速，相比手工優(yōu)化的 Flux 方案也取得更優(yōu)性能，得益于 Swizzling 調(diào)度、通信 offload 和 AOT 編譯等多重優(yōu)化。同時(shí)在 AMD 平臺(tái)上對(duì)比 PyTorch + RCCL 的組合，雖然整體加速幅度略小，但同樣取得顯著優(yōu)化，限制主要來自測(cè)試硬件算力偏弱和非 switch 拓?fù)洹?/p>

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

跨機(jī)測(cè)試方面，AG GEMM 提速 1.3 倍，GEMM RS 提速 1.4 倍，表現(xiàn)略低于 Flux，但在 shape 靈活性和可擴(kuò)展性上更具優(yōu)勢(shì)。我們還測(cè)試了高速推理場(chǎng)景下的單 token decoding，在 1M token context 下延遲可控制在 20–30 微秒，兼容 NVLink 與 PCIe 。

此外，我們對(duì) DeepEP 中的分布式調(diào)度邏輯進(jìn)行了功能復(fù)現(xiàn)，主要對(duì)齊其 All-to-All 路由與上下文分發(fā)策略。在 64 卡以內(nèi)的場(chǎng)景下，Triton-distributed 的性能與其基本持平，部分配置下略有提升。

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

打造開放的分布式編譯生態(tài)

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

在國(guó)內(nèi)乃至全球，這一領(lǐng)域仍較為空白。我們希望借助社區(qū)的力量，吸引更多開發(fā)者參與進(jìn)來，無論是在語(yǔ)法設(shè)計(jì)、性能優(yōu)化，還是支持更多類型的硬件設(shè)備方面，共同推動(dòng)技術(shù)進(jìn)步。最后，我們?nèi)〉昧瞬诲e(cuò)的性能表現(xiàn)，相關(guān)示例也全部開源，歡迎大家積極提 issue 交流，也期待更多小伙伴加入，共創(chuàng)未來！