mKernel：多GPU、多節點融合核心庫，實現GPU驅動通訊

GPU通訊開銷是生產級AI工作負載中可衡量的瓶頸。根據mKernel專案引用的資料，通訊可佔用前向傳播的43.6%和端到端訓練時間的32%。在流行的混合專家（MoE）模型中，裝置間通訊可佔總執行時間的47%。加州大學伯克利分校的UCCL專案團隊釋出了mKernel，這是一個持久CUDA核心庫，將節點內NVLink通訊、節點間RDMA和計算融合到單個核心中。

問題：主機驅動通訊

多GPU通訊的標準模型是主機驅動的：CPU執行控制路徑並呼叫NCCL或NVSHMEM等庫。該庫發出跨GPU的集合操作（如AllReduce、AllGather等）。計算和通訊在獨立的CUDA流上執行，並在核心邊界處重疊。

研究團隊指出該方法的兩個問題：

（1）CPU未隨GPU計算能力擴充套件。GB300 NVL72機架整合了72個Blackwell Ultra GPU和36個Grace CPU，提供720 PFLOP/s FP8/FP6、1.44 EFLOP/s FP4 Tensor Core效能和130 TB/s的全對全機架內NVLink頻寬。在這些速度下，微秒級的主機編排開銷（如cudaLaunchKernel呼叫、CPU側“所有寫入完成”檢查、流間事件）直接表現為流水線氣泡。

（2）主機驅動系統在粗粒度核心邊界處重疊計算和通訊。無法從主機側實現瓦片或塊級別的細粒度重疊。

替代方案是GPU驅動通訊：GPU自身觸發傳輸，通訊與計算融合到同一核心中。大多數現有的融合核心庫在單個節點或單個GPU內執行。mKernel針對多節點情況。

mKernel的功能

mKernel是一個持久CUDA核心庫。每個核心將節點內NVLink通訊、節點間RDMA和密集計算融合到單個核心中。

• 多GPU+多節點，一個核心：節點內NVLink和節點間RDMA都位於同一持久核心內。
• 細粒度核心內重疊：計算和通訊在瓦片/塊粒度上重疊，涵蓋節點內和節點間GPU通訊。
• 持久核心與SM專用化：CTA自行分配角色：計算、內部通訊、傳送、歸約。每個角色專用的SM數量可按形狀調整。
• 基於libibverbs的GPU驅動網路：mKernel使用GPU發起的RDMA寫入，不依賴NCCL或NVSHMEM。通訊後端從頭編寫以最大化效能並支援異構網路裝置。

五種融合核心

| 核心 | 融合內容 | 描述 | |------|----------|------| | AllGather + GEMM | AllGather → GEMM | 每個rank持有A的一個分片。當rank透過NVLink/RDMA收集其他rank的分片時，本地GEMM在瓦片到達時立即消費。 | | GEMM + AllReduce | GEMM → AllReduce | 在單次啟動中計算C = A @ B並跨所有rank歸約部分輸出。輸出瓦片一旦產生就被推入歸約樹。 | | MoE Dispatch + GEMM | All-to-All分發 → 分組GEMM | 將MoE令牌路由到其專家rank（節點內NVLink + 節點間all-to-all），並在同一核心中執行按專家分組的GEMM。令牌到達後立即處理，無需暫存緩衝區往返。 | | Ring Attention | Ring KV交換 → FlashAttention | 跨rank的序列並行注意力。每一步在環上旋轉KV塊，同時本地FlashAttention消費先前接收的塊。計算和環傳送/接收在單個持久核心內併發執行。 | | GEMM + ReduceScatter | GEMM → ReduceScatter | 計算C = A @ B並對輸出進行reduce-scatter。每個輸出瓦片一旦產生就被歸約並轉發到其所屬的rank。 |

評估設定

研究團隊在兩個2節點×8-H200叢集上評估mKernel，它們僅在節點間網路方面不同：

• AWS EFA：2節點×8 H200，節點內NVLink，節點間AWS EFA/SRD，每節點16×200 Gb/s EFA。
• ConnectX-7：2節點×8 H200，節點內NVLink，節點間InfiniBand，每節點8×400 Gb/s NVIDIA ConnectX-7。

mKernel與NCCL、Triton-distributed、Flux、Mercury、MagiAttention、Transformer-Engine和ring-flash-attention進行了基準測試。團隊指出更大規模的基準測試仍在進行中。

後端與要求

mKernel支援兩個網路後端：

| 後端 | 宏 | 傳輸 | 執行環境 | |------|------|------|----------| | CX7 | INTERNODE_BACKEND_IBVERBS | libibverbs RC | ConnectX-7 / InfiniBand / RoCE | | EFA | INTERNODE_BACKEND_EFA | libibverbs + efadv (SRD) | AWS p5/p5e (H200, EFA) |

兩個後端共享相同的主機側API和相同的GPU核心。僅代理/會話實現不同（CX7使用session.h，EFA使用session_efa.h）。要求：NVIDIA Hopper GPU（預設構建目標sm_90a）、CUDA 12.9、帶有PyTorch的Python。CX7後端需要libibverbs開發標頭檔案和庫。EFA後端需要AWS EFA安裝，包含libfabric、libibverbs、efadv，預設EFA_HOME=/opt/amazon/efa。

路線圖與關鍵要點

mKernel目前提供五種融合核心，支援ConnectX-7和AWS EFA後端。未來計劃包括：

• 完整異構加速器/NIC支援，包括拓撲感知發現、放置、路由
• 節點間megakernel：將多個融合步驟摺疊成一個覆蓋transformer層的單個megakernel
• Blackwell GPU支援

關鍵要點：

• mKernel將節點內NVLink、節點間RDMA和計算融合到單個持久CUDA核心中。
• 根據引用資料，MoE模型中通訊開銷最高佔執行時間的47%。
• 包含五種核心：AllGather+GEMM、GEMM+AllReduce、MoE Dispatch+GEMM、Ring Attention和GEMM+ReduceScatter。
• GPU發起的RDMA直接透過libibverbs實現，不依賴NCCL或NVSHMEM。
• 目前需要Hopper GPU（sm_90a）和ConnectX-7或AWS EFA網路；Blackwell支援已在路線圖中。

更多詳情請檢視倉庫和技術細節。