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支持已在路线图中。

更多详情请查看仓库和技术细节。