This video features two presentations from a Carnegie Mellon University research group on their "mega kernel" compiler project. The first talk introduces Mirage, a multi-level super-optimizer for GPU tensor programs. Mirage addresses the limitations of traditional compilers by performing a holistic optimization across the GPU's memory hierarchy. It uses a multi-level graph representation (M-Graph) and employs an exhaustive search strategy, pruned by a novel "abstract expression" technique and an automated theorem prover, to discover optimal computational graphs. These graphs can combine algebraic transformations, scheduling changes, and custom kernel fusion. Generated candidates are verified for correctness using a probabilistic method.
The second talk presents MPK (Mirage Persistent Kernel), a follow-up project that uses Mirage to compile an entire large language model (LLM) into a single, persistent "mega kernel." This approach overcomes the drawbacks of the standard kernel-per-layer execution model, such as synchronization barriers and poor resource utilization. MPK uses a fine-grained task graph to manage dependencies and implements a dynamic, task-based runtime system within the GPU kernel itself, with dedicated scheduler and worker SMs. This enables advanced inter-layer pipelining, overlapping of computation and communication, and efficient handling of dynamic workloads, pushing LLM serving performance close to theoretical hardware limits.
Mirage and MPK: A Super-Optimizer for Tensor Programs and Mega-Kernel LLM Compilation
00:00:50 The Challenge with GPU Kernels: Modern machine learning frameworks launch separate, pre-compiled kernels for each operation (e.g., matrix multiplication, softmax). This approach misses key optimization opportunities like kernel fusion and algebraic transformations, leading to suboptimal performance and requiring significant manual effort to create high-performance, fused kernels.
00:03:48 Introducing Mirage: A multi-level super-optimizer designed to holistically optimize tensor programs. It automatically performs algebraic transformations, schedule transformations, and generates custom kernels, requiring only a few lines of Python code from the user.
00:05:47 The M-Graph Representation: Mirage uses a multi-level Intermediate Representation (IR) called M-Graph, which mirrors the GPU's memory hierarchy. It consists of a kernel graph (device memory), a thread block graph (shared memory), and a thread graph (registers), allowing it to represent and optimize computations across all hardware levels.
00:09:25 Optimization via Exhaustive Search: Unlike traditional compilers that rely on a limited set of transformation rules, Mirage performs an exhaustive search to generate all possible computation graphs up to a bounded size. This allows it to discover novel optimization strategies.
00:11:49 Pruning with Abstract Expressions: To manage the massive search space, Mirage uses an innovative pruning technique. It converts the target computation into a simplified "abstract expression" (ignoring index details) and uses a theorem prover to check if a partially generated graph is a valid sub-expression. If not, the entire search branch is pruned, dramatically improving scalability.
00:25:40 Correctness Verification: Generated candidate kernels are not guaranteed to be correct. Mirage uses a probabilistic verifier that runs the original and generated programs with random inputs in a finite field to formally check for functional equivalence with a high degree of confidence.
00:43:32 From Fused Kernels to a Mega Kernel: The second talk introduces MPK (Mirage Persistent Kernel), a project that compiles an entire LLM into a single, persistent "mega kernel." This avoids the limitations of the conventional kernel-per-layer approach in LLM serving.
00:45:51 Limitations of the Kernel-per-Layer Model: The standard approach creates implicit synchronization barriers between layers, prevents inter-layer pipelining, makes it difficult to overlap computation and communication, and struggles with the dynamic workloads common in LLM serving (e.g., continuous batching).
00:53:22 The MPK System Architecture: MPK addresses these challenges using three key techniques:
Task Graph: A fine-grained representation where nodes are "tasks" (work for a single SM) and "events" (synchronization points), allowing for precise dependency management.
In-Kernel Runtime: A dynamic runtime system that runs entirely within the mega kernel. A few SMs act as "schedulers," dispatching tasks to the remaining "worker" SMs.
Super-Optimizer: Mirage is used to generate the highly optimized code for each task in the graph.
01:02:08 Pushing Performance to Hardware Limits: By eliminating kernel launch overhead and enabling advanced pipelining and resource overlap, the MPK mega kernel significantly reduces per-token latency. The results show MPK outperforms existing serving systems (like vLLM) by 1.2 to 6.7 times, approaching the theoretical performance limits imposed by the GPU's memory bandwidth.