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× on Ring Attention.
Figures & tables
Figure 1: Three schedules for computation and communication. (a) Intra-node fusion followed by a separate inter-node collective. (b) Two streams overlap communication of completed two-tile chunks with computation of subsequent chunks; each producer-kernel boundary releases a chunk. (c) mKernel schedules both communication tiers within the kernel, exposing tile-level readiness signal. We note that existing compute kernels already operate on tile granularity, so GEMM efficiency is not affected. Numbers identify output tiles; widths are schematic, not measured.
Level
Link
Bandwidth
SM roles
Transfer mechanism
Within a GPU
HBM3e
4.8 TB/s
compute
loads/stores, TMA
Within a node
NVLink + NVSwitch
450 GB/s
local communication
TMA, NVSwitch
Across nodes (IB)
InfiniBand (CX7)
50 GB/s
send and receive
NIC RDMA
Across nodes (EFA)
AWS EFA (SRD)
50 GB/s
send and receive
NIC RDMA
Table 1: Communication hierarchy of two separate H200 clusters: one uses InfiniBand (IB), and the other uses AWS EFA. Both share the within-GPU and within-node hierarchy shown here. Crossing a node boundary changes which SMs coordinate data movement. NVLink and network bandwidths are per GPU, per direction.
Figure 2: Architecture of mKernel (two nodes; one GPU per node shown in detail). Each GPU executes a persistent kernel whose thread blocks are assigned to roles. \scriptsize1⃝ Compute blocks write completed tiles to GPU memory and set readiness flags. \scriptsize2⃝ Intra-node communication blocks exchange tiles with the other seven GPUs over NVLink, using NVSwitch for broadcast and reduction. \scriptsize3⃝ Inter-node send blocks enqueue transfer commands in host memory. \scriptsize4⃝ A host proxy thread submits batches of commands to the NIC. \scriptsize5⃝ The NIC transfers data to the rail peer on node 1; the NIC or receiving proxy signals arrival, depending on the transport. \scriptsize6⃝ Receive blocks observe the flag and consume the data. The optional controller ( Section 3.2 ) reads progress counters and publishes a target number of communication blocks (dashed).
Figure 3: Static sweep of the SM partition (single node, 8 H100 GPUs). Each curve gives the latency of one problem size as a function of the number of communication SMs, normalized to the best fixed SM partition for that size (dashed line). The star markers at the right of each panel give the latency of the adaptive controller for the same sizes.
ConnectX-7 (InfiniBand)
AWS EFA (SRD)
Transport order
ordered per RC connection
unordered
Notification path
data write, then flag write
write with immediate; proxy sets flag
Table 2: Inter-node notification paths (EFA’s completion-based variant shown). Transport completion and GPU memory visibility are separate requirements.
Figure 4: GPUDirect Async (IBGDA) relative to host-assisted GPU-initiated communication on the ConnectX-7 testbed of Section 5.1 (2 nodes × 8 H200), with identical kernels. Bars give the throughput ratio, with GPUDirect Async in the numerator; the dashed line marks parity.
Kernel
Data dependency
Intra-node
Inter-node
Chunk
AllGather + GEMM (TP)
comm. → compute
multicast broadcast of each shard
shard sent once to each node; ring forwarding beyond two nodes
128 rows
GEMM + ReduceScatter (TP)
compute → comm.
TMA atomic add into the owner’s buffer
exchange of per-node partial sums
2–32 tiles
GEMM + AllReduce (TP)
compute → comm.
NVSwitch reduction; multicast broadcast of the result
exchange of 1/8 of the output per GPU
4 tiles (256 KiB)
MoE Dispatch + GEMM (EP)
comm. → compute
TMA pull of tokens from peers
token buffer copied to the rail peer
16 tokens; 512 KB
Ring Attention (SP)
independent within a step
TMA store of KV to the next GPU
KV slice sent once to each node
128-row KV tiles
Table 3: Fused kernels implemented in mKernel . TP, SP, and EP denote tensor, sequence, and expert parallelism.
Testbed
GPUs
Inter-node network
NICs per node
ConnectX-7
2 × 8 H200
InfiniBand
8 × 400 Gb/s ConnectX-7
AWS EFA
2 × 8 H200
AWS SRD
16 × 200 Gb/s EFA
Table 4: Testbeds. Both contain 2 nodes × 8 H200 GPUs connected by NVLink/NVSwitch within each node, with different inter-node networks.
Figure 5: Throughput on the AWS EFA testbed (2 nodes × 8 H200; TFLOPS per GPU, higher is better). mKernel is shown in orange. The evaluated DeepEP+DeepGEMM implementation does not support EFA; its bars show ConnectX-7 measurements for reference.
Figure 6: Throughput on the ConnectX-7 InfiniBand testbed (2 nodes × 8 H200; TFLOPS per GPU). A red × indicates that a baseline failed at that problem size.
Figure 7: Adaptive SM partition tuning . (a) GEMM+AllReduce ( M=N=32 K, K=4 K) shifts its target toward communication as computation nears completion. (b) MoE Dispatch+GEMM (128K total tokens) increases the dispatch-SM target while computation and communication make similar relative progress. Gray dashed lines mark the best fixed partitions (16 and 45 SMs respectively).
The rapid growth of large-scale machine learning (ML) has made distributed training across multiple GPUs a fundamental component of modern ML systems. As model sizes and computational throughput continue to increase, communication overhead has become a dominant bottleneck in multi-GPU training, particularly when computation and communication are executed sequentially. This work explores concurrent execution of computation and collective communication using two portable runtime controls: shared-memory-driven occupancy shaping for computation kernels and elevated scheduling priority for communication kernels. Our approach regulates computation-kernel residency through per-block shared-memory allocation, leaving sufficient on-chip resources for communication kernels to make progress. In addition, assigning higher priority to communication streams ensures steady communication progress once resources become available. Experiments on NVIDIA A40, A100, H100, and AMD MI250X GPUs demonstrate that the proposed method enables effective computation-communication overlap and reduces total execution time by up to 25.5 percent, without modifying vendor libraries or kernel implementations.
Minyu Cui, Miquel Pericas
Chalmers University of Technology and University of Gothenburg, Gothenburg, Sweden
Improving GPU kernel efficiency is crucial for efficient AI systems. Recent work has explored leveraging large language models (LLMs) for GPU kernel generation and optimization. However, existing LLM-based kernel optimization pipelines typically rely on opaque, implicitly learned heuristics within the LLMs to determine optimization strategies. This leads to inefficient trial-and-error and weakly interpretable optimizations. Our key insight is to replace implicit heuristics with expert optimization knowledge that is aware of task trajectories. Specifically, we present KernelClaw, a multi-agent framework with a dual-level memory architecture. KernelClaw operates by coordinating agents with long-term memory of reusable expert knowledge and short-term memory to prevent repetitive backtracking. On KernelBench Levels 1-3, KernelClaw achieves a 100% success rate and average speedups of 6.98x, 7.89x, and 1.94x over Torch Eager on Levels 1, 2, and 3, respectively, outperforming prior baselines. Code is available at https://github.com/0satan0/KernelMem/.
Qitong Sun, Jun Han, Tianlin Li +6
School of Computer Science and Engineering, Beihang University, China · Zhejiang Lab, China · Nanyang Technological University, Singapore
We present KernelArc, a multi-agent framework for autonomous GPU kernel optimization across heterogeneous workloads. Strategy-specialized agents run in parallel and coordinate through conclusions-only shared memory, a deterministic benchmark guard, and read-only cross-agent state with plateau-triggered drafting. We evaluate KernelArc on NVIDIA H100 and B200 GPUs using category-representative SOL-ExecBench workloads. The resulting implementations span custom BF16 GEMM, static cuBLASLt Expert-API configuration tables, fused mixture-of-experts backward, shape-gated decoder-layer fusion, native NVFP4 grouped-query attention, and paged prefill attention. In the public SOL-ExecBench leaderboard snapshot recorded on August~20, 2026, KernelArc ranked first on every representative L1, L2, Quantization, and FlashInfer task evaluated. The trajectories support the paper's central motivation: shared multi-agent search can broaden exploration and reach stronger incumbents within a fixed candidate budget, while the value of individual coordination features depends on the kernel and optimization stage.
Joyjit Kundu, Ben Stoffelen, Kaili Wang +2
AILabs, Interuniversity Microelectronics Centre (IMEC)