PaperlayerView exact revision on arXiv

Source-linked AI summary

mKernel: Fast Multi-GPU, Multi-Node Fused Kernels

Ziming Mao, Yihan Zhang, Shawn Wei Chew, Shuang Ma, Costin Raiciu, Yang Zhou, Scott Shenker, Ion Stoica

arXiv:2609.13585v1cs.DCcs.AIcs.LGcs.NI

TL;DR

Communication is a substantial cost in distributed model training: in production MoE training, it accounts for 43.6% of the forward pass and 32% of end-to-end training time. mKernel fuses computation, intra-node NVLink communication, and inter-node RDMA at tile granularity, using SM specialization, hierarchical data movement, and portable host-assisted GPU-initiated communication. On two 16-GPU H200 clusters, mKernel achieves up to 1.72× on GEMM+AllReduce and 1.88× on Ring Attention over unfused baselines.

  • Problem

    Communication is a substantial cost in distributed model training: in production MoE training, it accounts for 43.6% of the forward pass and 32% of end-to-end training time.

    Source: §1 · p. 1

  • Method

    mKernel fuses computation, intra-node NVLink communication, and inter-node RDMA at tile granularity, using SM specialization, hierarchical data movement, and portable host-assisted GPU-initiated communication.

    Source: §1 · pp. 1–2

  • Results

    On two 16-GPU H200 clusters, mKernel achieves up to 1.72× on GEMM+AllReduce and 1.88× on Ring Attention over unfused baselines.

    Source: §7 · p. 13

  • Takeaways & Limitations

    The same kernels coordinate NVLink and inter-node RDMA on InfiniBand and EFA, and IBGDA provides little additional performance over host-assisted GPU-initiated communication.

    Source: §1 · p. 2

Abstract

from arXiv · show

Communication has become a bottleneck in distributed training and inference of large models. Overlapping communication with computation at the granularity of kernels, on separate streams, reduces only part of this communication cost. Fused kernels often have better performance by transmitting each output tile as soon as it is produced, but existing fused kernels are largely confined to a single NVLink domain. We present mKernel, a library of multi-GPU, multi-node fused kernels that overlap computation, intra-node NVLink communication, and inter-node RDMA at tile granularity. mKernel partitions the streaming multiprocessors (SMs) of a persistent kernel into compute and communication roles, and an on-GPU controller tunes the SM partition adaptively at run time, since the best SM partition varies with the kernel and the input shape. It structures data movement hierarchically so that data traversing the inter-node network is minimized. Finally, it drives the network from the GPU through a lightweight command queue and host proxy implemented directly on RDMA verbs, which allows the same kernels to run on any network backend (e.g. InfiniBand and on AWS EFA); we observe, surprisingly, that GPUDirect Async (IBGDA) yields little additional benefit over host-assisted GPU-initiated communication. We implement five kernels spanning tensor, sequence, and expert parallelism. On two 16-GPU H200 clusters, mKernel achieves speedups of up to 1.72x on GEMM+AllReduce and $1.88\times$ on Ring Attention.

1 Introduction

Communication consumes a substantial share of distributed model execution, while existing overlap often stops at kernel boundaries or within a single node. mKernel coordinates computation, NVLink communication, and inter-node RDMA at tile granularity and reports speedups across tensor and sequence parallelism workloads.

  • Communication accounts for 43.6% of the forward pass and 32% of end-to-end training time in production MoE training, while inter-device communication reaches up to 47% of execution time across popular MoE models and frameworks.The introduction also notes that accelerator compute throughput is growing faster than network bandwidth.Source: §1 · p. 1
  • Prior approaches overlap communication at chunk or tile granularity, but most focus on one node or fuse only selected stages; mKernel coordinates computation, intra-node NVLink, and inter-node RDMA within one fused kernel.Figure 1 contrasts kernel-boundary chunk release with mKernel's tile-level readiness signal.Source: §1 · p. 1, Figure 1 · p. 2
  • mKernel combines compute and communication SM roles, hierarchical data movement, portable host-assisted GPU-initiated RDMA, and runtime SM partitioning based on measured progress and remaining work.The same communication path supports InfiniBand and AWS EFA; IBGDA yields little performance difference from host-assisted GPU-initiated communication.Source: §1 · pp. 1–2
  • Up to 1.41× on AllGather+GEMM, 1.72× on GEMM+AllReduce, and 1.88× on Ring Attention versus cuBLAS or FlashAttention followed by NCCL.The evaluation covers five kernels on two clusters of 2 nodes × 8 H200 GPUs.Source: §1 · p. 2

2 Background and Motivation

The background describes a hierarchy in which intra-node NVLink and inter-node network communication have different capabilities. Efficient fused kernels must balance tile-level overlap against communication costs, transport semantics, synchronization, and workload-dependent SM allocation.

2.1 Communication hierarchy of GPU clusters

  • Across clusters, GPU memory and intra-node NVLink provide local communication, while GPUDirect RDMA transfers payloads between GPU memories across nodes; kernels must allocate SMs to computation, local communication, and network coordination.The evaluated clusters use InfiniBand or AWS EFA and connect eight GPUs per node through NVLink and NVSwitch.Source: §2.1 · p. 3

2.2 Overlapping communication and computation in fused kernels

  • Tile-level fusion lets communication consume output tiles as produced and computation consume input tiles as they arrive, while additional communication SMs can reduce transfer time but slow computation.Execution time also includes pipeline fill/drain and synchronization costs.Source: §2.2 · p. 3

2.3 Challenges of designing efficient multi-GPU, multi-node fused kernels

  • The bandwidth gap and message-oriented RDMA path motivate local aggregation or replication, early network transfers, and grouping ready tiles into chunks to amortize per-message costs.In rail-optimized topologies, inter-node data is exchanged between GPUs with the same local index while NVLink handles local movement.Source: §2.3 · pp. 3–4
  • Different transport ordering guarantees and GPU memory visibility require readiness and completion mechanisms that let persistent kernels safely consume remote data without global synchronization after every transfer.InfiniBand reliable connections preserve write order, whereas EFA's SRD transport does not guarantee the same ordering.Source: §2.3 · p. 4
  • The best communication-versus-compute SM allocation varies with kernel type, input shape, workload mix, and within-kernel imbalance, making static profiling costly to maintain.The reported intra-node sweeps span best allocations of 2–64 communication SMs.Source: §2.3 · p. 4

3 Design

mKernel coordinates compute, local NVLink communication, and inter-node RDMA inside persistent kernels. It bridges tile-level progress across network messages, adjusts SM roles as work progresses, and uses a host proxy for portable network submission.

3.1 Bridging intra- and inter-node communication granularity

  • NVSwitch performs local reduction or broadcast before rail peers exchange data, while independent tile and network-chunk sizes balance early transfers against message overhead.For GEMM+AllReduce, four locally reduced tiles form a 256 KiB chunk that can be sent once its tiles are ready.Source: §3.1 · pp. 4–5

3 Design

  • Compute, intra-node communication, inter-node submission, and receive blocks coordinate tile readiness through GPU memory, NVLink, a host proxy, and the NIC.The optional controller publishes a target number of communication blocks based on progress counters.Source: §3 · p. 4, Figure 2 · p. 5

3.2 SM partioning and adaptive tuning

  • The adaptive controller estimates role costs from cycle counters and sets a communication-block target to equalize predicted completion times.Blocks can change roles at task boundaries; a compute block waiting for communication input may temporarily assist communication.Source: §3.2 · p. 6

3.3 Portable inter-node communication in fused kernels

  • The host proxy batches ready commands while preserving per-chunk transfers and notifications; tested GPUDirect Async provides little performance gain over host-assisted communication.InfiniBand and EFA use different notification paths, and the proxy may submit partial batches rather than wait for a full batch.Source: §3.3 · p. 7

4 Implementation

The five implemented kernels use persistent scheduling, hierarchical communication, and readiness flags, with transfer timing and tile order tailored to each kernel.

4 Implementation

  • AllGather+GEMM and Ring Attention initiate inter-node transfers at launch by submitting the input shard or KV slice to rail peers.This starts inter-node movement early in those schedules.Source: §4 · p. 8
  • NVSwitch handles local aggregation and replication; for GEMM+AllReduce, each GPU sends one eighth of the output across the network.Local reduction precedes the inter-node exchange.Source: §4 · p. 8
  • AllGather+GEMM processes its local shard, then remaining shards within the node, and finally remote shards.The tile order follows input availability.Source: §4 · p. 8

5 Evaluation

Across tensor, expert, and sequence parallelism, mKernel is evaluated against unfused and fused baselines on two H200 clusters. Results include speedups across kernels and show that SM-allocation preferences vary by workload and shape.

5.1 Experimental setup

  • The evaluation uses two clusters, each with 2 nodes × 8 H200 GPUs, and compares mKernel with unfused and supported fused implementations.One cluster uses ConnectX-7 and the other EFA; the evaluated DeepEP implementation does not support EFA.Source: Table 4 · p. 9, §5.1 · p. 9

5.2 Tensor parallelism

  • 1.41× on EFA and 1.34× on ConnectX-7 are the maximum AllGather+GEMM speedups over cuBLAS+NCCL, with mKernel faster at every tested size.Inter-node transfers start first, and each remote shard is sent once per destination node.Source: §5.2 · p. 9
  • 1.53× on EFA and 1.72× on ConnectX-7 are mKernel’s peak GEMM+AllReduce speedups.The schedule reduces within NVSwitch, sends each GPU’s one eighth of the output across the network, and overlaps the three phases with GEMM.Source: §5.2 · p. 9
  • For M ≥16K, GEMM+ReduceScatter is 1.02–1.27× faster, and it outperforms Triton-distributed at most input sizes.Source: §5.2 · p. 9

5.3 Expert parallelism

  • On EFA, MoE Dispatch+GEMM achieves 3.1–4.5× speedups over all-to-all followed by GEMM, a comparison that also reflects expert-computation organization.The evaluated DeepEP+DeepGEMM bars are ConnectX-7 measurements because that implementation does not support EFA.Source: §5.3 · p. 9, Figure 5 · p. 10, Figure 6 · p. 11

5.4 Sequence parallelism

  • 1.88× over the unfused baseline on ConnectX-7 is the reported peak speedup for Ring Attention.On EFA, Ring Attention reaches up to 1.78× over the unfused baseline.Source: §5.4 · p. 11

5.5 Adaptive SM partitioning

  • The best measured communication allocation ranges from 2 to 64 SMs across workloads, and the worst measured AllReduce partition is approximately 25× slower than the best.Too few communication SMs can stall computation or build a communication backlog; too many reduce compute resources.Source: §5.5 · p. 11
  • Without per-shape tuning, the adaptive controller achieves a geometric mean of 1.18× across 21 configurations compared to the static best fixed SM partition.The allocation target also changes during execution for GEMM+AllReduce and MoE Dispatch+GEMM.Source: §5.5 · p. 12, Figure 7 · p. 12

6 Related Work

Related systems address tiled computation, fusion, overlap scheduling, communication interfaces, and persistent execution. mKernel focuses on explicit schedules for inter-node fused kernels.

Kernel frameworks and distributed fusion

  • GPU abstractions and systems expose tiled computation, communication fusion, or compiler optimization; mKernel uses explicit schedules within persistent kernels to coordinate NVLink, RDMA chunks, and SM roles.The cited systems include CUTLASS, Triton, ThunderKittens, ParallelKittens, FLUX, TileLink, Triton-distributed, and Mercury.Source: §6 · p. 12

Overlap scheduling and SM allocation

  • Prior work overlaps operations or selects profiled compute–communication configurations, while mKernel updates its target SM partition during execution using measured progress and remaining work.The related work includes operation-level overlap, finer-grained scheduling, and workload-dependent resource allocation.Source: §6 · p. 12

Communication interfaces and topology; Persistent kernels and megakernels

  • Communication systems provide device-side operations, proxies, or topology-aware collectives; mKernel combines a compact command queue and direct verbs with fused schedules on InfiniBand and EFA.Persistent-kernel systems also exist, while mKernel focuses on computation and communication within each distributed kernel.Source: §6 · pp. 12–13

7 Conclusion

mKernel fuses multi-GPU, multi-node computation and communication at tile granularity. On two 16-GPU H200 clusters, it outperforms unfused baselines on GEMM+AllReduce and Ring Attention.

7 Conclusion

  • 1.72× on GEMM+AllReduce and 1.88× on Ring Attention are mKernel speedups over the unfused baselines.Measured on two 16-GPU H200 clusters.Source: §7 · p. 13
  • mKernel is a multi-GPU, multi-node fused-kernel library that overlaps computation and communication at tile granularity.Source: §7 · p. 13
  • mKernel combines SM specialization, hierarchical data movement, and portable host-assisted GPU-initiated communication.Source: §7 · p. 13
  • The library implements five kernels spanning tensor, sequence, and expert parallelism.Source: §7 · p. 13
  • Its intra-node adaptive mode dynamically tunes compute and communication SM allocations, reducing the overhead of tuning kernels for each kernel type and input shape.Source: §7 · p. 13
Loading 2609.13585v1…