Mirage Persistent Kernel: A Compiler and Runtime for Mega-Kernelizing Tensor Programs

TL;DR

Mirage Persistent Kernel (MPK) tackles the bottlenecks of kernel-per-operator inference on multi-GPU systems by introducing a SM-level $t$Graph that encodes fine-grained task dependencies and enables cross-operator software pipelining. The MPK compiler converts tensor programs into optimized $t$Graphs and emits high-performance CUDA code for each task, while an in-kernel runtime executes all tasks inside a single mega-kernel using decentralized scheduling and an event-driven model. Key contributions include event fusion, graph normalization, and linearization within the compiler, plus a paged shared-memory abstraction and hybrid JIT/AOT task launching in the runtime to minimize overhead and maximize throughput. The results show up to $1.7\times$ latency reduction and up to $10\times$ multi-GPU throughput gains over kernel-per-operator baselines, with MoE optimizations and broad model-GPU coverage, demonstrating that end-to-end mega-kernel inference is practical and beneficial for real-world LLM serving. MPK maintains compatibility with existing ML frameworks (e.g., PyTorch) by providing a backend that can mega-kernelize models with minimal code changes, unlocking near-hardware-limits performance while preserving developer usability.

Abstract

We introduce Mirage Persistent Kernel (MPK), the first compiler and runtime system that automatically transforms multi-GPU model inference into a single high-performance megakernel. MPK introduces an SM-level graph representation that captures data dependencies at the granularity of individual streaming multiprocessors (SMs), enabling cross-operator software pipelining, fine-grained kernel overlap, and other previously infeasible GPU optimizations. The MPK compiler lowers tensor programs into highly optimized SM-level task graphs and generates optimized CUDA implementations for all tasks, while the MPK in-kernel parallel runtime executes these tasks within a single mega-kernel using decentralized scheduling across SMs. Together, these components provide end-to-end kernel fusion with minimal developer effort, while preserving the flexibility of existing programming models. Our evaluation shows that MPK significantly outperforms existing kernel-per-operator LLM serving systems by reducing end-to-end inference latency by up to 1.7x, pushing LLM inference performance close to hardware limits. MPK is publicly available at https://github.com/mirage-project/mirage.

Figures (13)

• Figure 1: An overview of MPK.
• Figure 2: Comparing how MPK and existing approaches support intra- and cross-task pipelining.
• Figure 3: Comparing how MPK and existing approaches support fine-grained kernel overlap between tasks. Data dependencies (black arrows in \ref{['fig:mpk-overlap']}) ensure correctness.
• Figure 4: The MPK compiler transforms a kernel-level computation graph into an optimized SM-level $t$Graph. MM, AT, and AR denote MatMul, Attention, and AllReduce tasks, respectively.
• Figure 5: The MPK compiler workflow. In (b), $Q$, $K$, $V$, $A$, $O$, and $R$ denote the set of tasks produced by decomposing the query projection, key projection, value projection, attention, output projection, and RMSNorm operators, respectively. $D_1$ and $D_2$ in (e) are dummy tasks inserted during $t$Graph normalization to guarantee that each task has exactly one triggering event. Finally, (f) shows how MPK linearizes the $t$Graph and stores the resulting structure in GPU memory, where both tasks and events follow a uniform, canonical representation.

Theorems & Definitions (2)

• Definition 4.1: Successor-set fusion
• Definition 4.2: Predecessor-set fusion