mKernel: Unleashing Multi-GPU, Multi-Node Fused Kernel Power for GPU-Driven Communication

GPU communication often acts as a significant speed constraint in real-world AI workloads. Citing data from the mKernel initiative, researchers note that communication can occupy 43.6% of the forward pass and 32% of total training duration. For common Mixture-of-Experts (MoE) architectures, transmitting data between devices may take as much as 47% of total runtime. A team from UC Berkeley’s UCCL project has introduced mKernel, a library of persistent CUDA kernels designed to merge NVLink traffic within a node, RDMA transfers across nodes, and processing tasks into one unified operation.

The Challenge: CPU-Managed Communication

Traditionally, handling data exchange between GPUs follows a CPU-managed pattern. The CPU manages the control flow, invoking libraries such as NCCL or NVSHMEM to execute operations like AllReduce or AllGather. Processing and data transfer occur on independent CUDA streams and are synchronized only at kernel boundaries.

The researchers highlight two core limitations with this traditional setup:

(1) CPU capabilities fail to keep pace with advancing GPU power. Consider a GB300 NVL72 rack, which packs 72 Blackwell Ultra GPUs alongside 36 Grace CPUs, providing 720 PFLOP/s of FP8/FP6 performance, 1.44 EFLOP/s of FP4 Tensor Core throughput, and 130 TB/s of rack-wide NVLink bandwidth. At these extreme speeds, even minor delays from host-side coordination — like a cudaLaunchKernel call, a CPU verifying writes are complete, or a stream-synchronization event — manifest as stalls in the processing pipeline.

(2) CPU-mediated systems can only synchronize processing and communication at coarse, kernel-level intervals. More detailed, chunk-level coordination is impossible to achieve from the host processor.

The proposed solution is GPU-managed communication: the GPU autonomously handles data transfers, integrating communication operations directly into the same kernel executing the math. While many current fused-kernel tools are confined to one node or one GPU, mKernel is engineered specifically for multi-node environments.

mKernel’s Functionality

mKernel provides a set of persistent CUDA kernels. Each kernel combines intra-node NVLink transfers, inter-node RDMA, and core computation within a single execution unit.

Unified Multi-GPU, Multi-node Execution: Both NVLink communication and cross-node RDMA traffic are integrated into the same persistent kernel.

Detailed Intra-kernel Synchronization: Processing and data exchange happen concurrently at a tile-by-tile or chunk-by-chunk level, improving efficiency for both local and remote GPU communication.

Persistent Kernel with Specialized Streaming Multiprocessors (SMs): Compute Thread Arrays (CTAs) are designated for specific functions: compute, intra-comm, inter-send, and inter-reduce. The allocation of SMs to each function can be customized based on the workload shape.

GPU-Initiated Networking via libibverbs: mKernel executes direct GPU-initiated RDMA writes without relying on NCCL or NVSHMEM. This custom-built backend is engineered for peak performance and compatibility with various networking hardware types.

The Five Integrated Operations

Testing Environment

The mKernel team tested their library on two separate clusters, each consisting of 2 nodes with 8 H200 GPUs; the key difference between the clusters lies in their inter-node networking fabric:

mKernel was measured against several established tools, including NCCL, Triton-distributed, Flux, Mercury, MagiAttention, Transformer-Engine, and ring-flash-attention. The developers mention that comprehensive scaling evaluations are currently proceeding.

Supported Backends and System Requirements

mKernel accommodates two primary networking backends:

Both backends utilize the identical host API and the same GPU-side kernel. Only the proxy/session implementation varies (session.h for CX7, session_efa.h for EFA). Requirements: NVIDIA Hopper GPUs (default build targets sm_90a), CUDA 12.9, Python with PyTorch. The CX7 backend necessitates libibverbs development libraries. The EFA backend necessitates AWS EFA software installed with libfabric, libibverbs, efadv, and EFA headers, typically situated at EFA_HOME=/opt/amazon/efa.

Marktechpost’s Visual Explainer