arXiv:2611.16248 · v1

MegaOps: A Kernel Overlapping Framework with PyTorch-like Modularity and Megakernel Performance

首次提交 2026-11-12 11:43

版本更新 2026-11-12 11:43

雷达观测 2026-11-13 14:16

摘要

Modern ML frameworks execute models using the kernel-per-operator paradigm, which runs each operator as a separate GPU kernel. This paradigm is modular, since developers can compose various kernel implementations and profile each kernel in isolation. However, this paradigm enforces dependencies at the granularity of entire kernels, which places a barrier between dependent kernels. These barriers cause streaming multiprocessors to idle, serialize computation and communication, and prevent a kernel from prefetching its weights while the preceding kernel runs. Megakernels remove kernel barriers by fusing the operators of an entire layer or model into a single kernel, but they give up the modularity of the kernel-per-operator paradigm. In this work, we ask whether a kernel-per-operator ML framework can perform on par with megakernels. Our key insight is that removing kernel barriers does not require fusing kernels. We develop MegaOps, a framework that launches dependent kernels on separate CUDA streams and enforces each dependency on individual tiles of the intermediate tensor that the producer kernel writes and the consumer kernel reads. Since kernels synchronize on data tiles instead of on each other's thread blocks, any two kernels that adopt the same data tiling can be composed without modifying either kernel. We evaluate MegaOps on dense Llama3 and Qwen3 models on single H100, A100, and L40 GPUs and on multi-GPU servers with 4$\times$A100 and 8$\times$H100. Our MegaOps framework matches the throughput of hand-tuned ThunderKittens megakernels, and achieves 1.09$\times$-1.10$\times$ the decode throughput and 0.98$\times$-1.10$\times$ the prefill throughput of the best-performing baseline, while reusing existing Triton, cuBLAS, and cuDNN kernels.