mKernel: multi-GPU, multi-node fused kernels
Blog | Join Slack | Twitter/X | Roadmap | Quick Start | Open Letter
mKernel is under active development, including optimizing for larger scale, different GPUs, and network topologies. The goal is to have a library for commonly used multi-node/GPU distributed kernels.
| Kernel | What it fuses | Description |
|---|---|---|
| AllGather + GEMM | AllGather → GEMM | Each rank holds a shard of the activation A. While ranks gather peers' shards over NVLink/RDMA, the local GEMM consumes tiles as soon as they arrive — overlapping the gather with (A_full @ B) so the matmul starts before the collective finishes. |
| GEMM + AllReduce | GEMM → AllReduce | Computes C = A @ B and reduces partial outputs across all 16 ranks in one launch. Output tiles are pushed into the reduction tree the instant they're produced, hiding the AllReduce inside the GEMM tail. |
| MoE Dispatch + GEMM | All-to-All dispatch → grouped GEMM | Routes MoE tokens to their expert ranks (intra-node NVLink + inter-node all-to-all) and runs the per-expert grouped GEMM in the same kernel. Tokens are matmul'd as soon as they land, no staging buffer round-trip. |
| MoE Dispatch + FFN + Combine | All-to-All dispatch → grouped FFN → All-to-All combine | A full expert-parallel MoE layer in one kernel: routes tokens to their expert ranks, runs the per-expert FFN (gemm1 → SwiGLU → gemm2), then combines the weighted expert outputs back to each token's source rank. Dispatch, both GEMMs, and the combine overlap across intra-node NVLink and inter-node RDMA without leaving the kernel. |
| Ring Attention | Ring KV exchange → FlashAttention | Sequence-parallel attention across 16 ranks: each step rotates a KV chunk around the ring while the local FlashAttention consumes the previously-received chunk. Compute and the ring send/recv run concurrently inside a single persistent kernel. |
| GEMM + ReduceScatter | GEMM → ReduceScatter | Computes C = A @ B and reduce-scatters the output across ranks. Each output tile is reduced and forwarded to its owning rank as soon as it's produced, so the scatter overlaps the GEMM rather than following it. |
# Pick BACKEND=efa for AWS EFA, or BACKEND=cx7 for ConnectX-7 / InfiniBand.
make BACKEND=cx7 PYTHON=python3 all
# Two-node benchmark example. Run from node 0; node 1 is launched over SSH.
NODE0_IP=<node0-data-ip> \
NODE1_IP=<node1-data-ip> \
NODE1_SSH=<node1-ssh-target> \
bash bench/run.sh all bench 2
make plots
sm_90a.CUDA_HOME=/usr/local/cuda-12.9), override with CUDA_HOME=....PYTHON=/path/to/python.EFA_HOME=/opt/amazon/efa by default.torchrun, passwordless SSH from node 0 to peer nodes, and routable data-plane IPs in NODE*_IP.| Backend | Macro | Transport | Where it runs |
|---|---|---|---|
| CX7 | -DINTERNODE_BACKEND_IBVERBS | libibverbs RC | ConnectX-7 / InfiniBand / RoCE |
| EFA | -DINTERNODE_BACKEND_EFA | libibverbs + efadv (SRD) | AWS p5/p5e (H200, EFA) |
Both backends share the same host-side API and the same on-GPU kernel; only the proxy / session implementation differs (include/comm/internode/session.h for CX7, session_efa.h for EFA).
| Kernel | Plot |
|---|---|
| AllGather + GEMM |  |
| GEMM + AllReduce |  |
| MoE Dispatch + GEMM |  |
| MoE Dispatch + FFN + Combine |  |
| Ring Attention |  |
| GEMM + ReduceScatter |  |
| Kernel | Plot |
|---|---|
| AllGather + GEMM |  |
| GEMM + AllReduce |  |
| Ring Attention |  |
| GEMM + ReduceScatter |  |
The MMA code is adapted from ThunderKittens (HazyResearch). Many thanks to the TK authors.
@article{mao2026mkernel,
title={mKernel: Fast Multi-GPU, Multi-Node Fused Kernels},
author={Mao, Ziming and Zhang, Yihan and Chew, Shawn Wei and Ma, Shuang and Raiciu, Costin and Zhou, Yang and Shenker, Scott and Stoica, Ion},
journal={arXiv preprint arXiv:2609.13585},
year={2026}
}
MIT — see LICENSE.
1,264 followers · starred May 2026
352 followers · starred May 2026
2,136 followers · starred May 2026
541 followers · starred May 2026
mKernel: multi-GPU, multi-node fused kernels
Blog | Join Slack | Twitter/X | Roadmap | Quick Start | Open Letter
mKernel is under active development, including optimizing for larger scale, different GPUs, and network topologies. The goal is to have a library for commonly used multi-node/GPU distributed kernels.
| Kernel | What it fuses | Description |
|---|---|---|
| AllGather + GEMM | AllGather → GEMM | Each rank holds a shard of the activation A. While ranks gather peers' shards over NVLink/RDMA, the local GEMM consumes tiles as soon as they arrive — overlapping the gather with (A_full @ B) so the matmul starts before the collective finishes. |
| GEMM + AllReduce | GEMM → AllReduce | Computes C = A @ B and reduces partial outputs across all 16 ranks in one launch. Output tiles are pushed into the reduction tree the instant they're produced, hiding the AllReduce inside the GEMM tail. |
| MoE Dispatch + GEMM | All-to-All dispatch → grouped GEMM | Routes MoE tokens to their expert ranks (intra-node NVLink + inter-node all-to-all) and runs the per-expert grouped GEMM in the same kernel. Tokens are matmul'd as soon as they land, no staging buffer round-trip. |
| MoE Dispatch + FFN + Combine | All-to-All dispatch → grouped FFN → All-to-All combine | A full expert-parallel MoE layer in one kernel: routes tokens to their expert ranks, runs the per-expert FFN (gemm1 → SwiGLU → gemm2), then combines the weighted expert outputs back to each token's source rank. Dispatch, both GEMMs, and the combine overlap across intra-node NVLink and inter-node RDMA without leaving the kernel. |
| Ring Attention | Ring KV exchange → FlashAttention | Sequence-parallel attention across 16 ranks: each step rotates a KV chunk around the ring while the local FlashAttention consumes the previously-received chunk. Compute and the ring send/recv run concurrently inside a single persistent kernel. |
| GEMM + ReduceScatter | GEMM → ReduceScatter | Computes C = A @ B and reduce-scatters the output across ranks. Each output tile is reduced and forwarded to its owning rank as soon as it's produced, so the scatter overlaps the GEMM rather than following it. |
# Pick BACKEND=efa for AWS EFA, or BACKEND=cx7 for ConnectX-7 / InfiniBand.
make BACKEND=cx7 PYTHON=python3 all
# Two-node benchmark example. Run from node 0; node 1 is launched over SSH.
NODE0_IP=<node0-data-ip> \
NODE1_IP=<node1-data-ip> \
NODE1_SSH=<node1-ssh-target> \
bash bench/run.sh all bench 2
make plots
sm_90a.CUDA_HOME=/usr/local/cuda-12.9), override with CUDA_HOME=....PYTHON=/path/to/python.EFA_HOME=/opt/amazon/efa by default.torchrun, passwordless SSH from node 0 to peer nodes, and routable data-plane IPs in NODE*_IP.| Backend | Macro | Transport | Where it runs |
|---|---|---|---|
| CX7 | -DINTERNODE_BACKEND_IBVERBS | libibverbs RC | ConnectX-7 / InfiniBand / RoCE |
| EFA | -DINTERNODE_BACKEND_EFA | libibverbs + efadv (SRD) | AWS p5/p5e (H200, EFA) |
Both backends share the same host-side API and the same on-GPU kernel; only the proxy / session implementation differs (include/comm/internode/session.h for CX7, session_efa.h for EFA).
| Kernel | Plot |
|---|---|
| AllGather + GEMM | |
| GEMM + AllReduce | |
| MoE Dispatch + GEMM | |
| MoE Dispatch + FFN + Combine | |
| Ring Attention | |
| GEMM + ReduceScatter | |
| Kernel | Plot |
|---|---|
| AllGather + GEMM | |
| GEMM + AllReduce | |
| Ring Attention | |
| GEMM + ReduceScatter | |
The MMA code is adapted from ThunderKittens (HazyResearch). Many thanks to the TK authors.
@article{mao2026mkernel,
title={mKernel: Fast Multi-GPU, Multi-Node Fused Kernels},
author={Mao, Ziming and Zhang, Yihan and Chew, Shawn Wei and Ma, Shuang and Raiciu, Costin and Zhou, Yang and Shenker, Scott and Stoica, Ion},
journal={arXiv preprint arXiv:2609.13585},
year={2026}
}
MIT — see LICENSE.
1,264 followers · starred May 2026
352 followers · starred May 2026
2,136 followers · starred May 2026
541 followers · starred May 2026