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支持已在路線圖中。

更多詳情請查看倉庫和技術細節。