This presentation introduces two related projects from CMU: Mirage, a multi-level superoptimizer for tensor programs, and MPK (Mirage Persistent Kernel), which compiles entire large language models (LLMs) into single "mega kernels."
Mangi Wu first details Mirage, explaining its holistic approach to optimizing GPU kernel execution. It addresses the complexity of GPU hardware and the suboptimal performance of current systems that launch separate kernels for individual operators. Mirage uses a novel multi-level graph representation (mugraph) and an exhaustive search strategy guided by "abstract expressions" to discover and generate highly optimized, custom CUDA kernels. It supports algebraic transformations, schedule transformations, and custom kernel generation, verified by a probabilistic equivalence checker. The approach significantly reduces engineering effort while outperforming existing systems by up to 3.3 times.
Shinhao Chang then presents MPK, a follow-up project aimed at compiling entire LLMs into a single mega kernel to overcome limitations of the "kernel per layer" approach in LLM serving. These limitations include poor inter-layer pipelining, difficulty overlapping computation and communication, and high kernel launch overheads due to dynamic workloads. MPK leverages a "task graph" for fine-grained dependency management and an internal runtime system with GPU-resident schedulers and workers to handle dynamic workloads efficiently within a single kernel. This approach, built upon Mirage's superoptimization, achieves better load balancing, maximizes resource utilization, and delivers LLM serving performance improvements of 1.2 to 6.7 times over existing systems, pushing towards hardware limits. Both projects are open-source, offering adaptable and high-performance solutions for modern deep learning workloads.
Summarizing the Mega Kernel Compiler Project
0:00 Project Overview: Introduction to the mega kernel compiler project from CMU, featuring two talks on Mirage and Mirage Persistent Kernel (MPK) by Mangi Wu and Shinhao Chang.
0:46 Mirage - Problem Statement: GPUs are complex, making manual high-performance kernel design difficult. Existing systems launching separate kernels lead to suboptimal performance due to missed optimizations (kernel fusion, algebraic transformations, local memory reuse).
0:03:01 Manual Effort: Achieving high performance often requires hundreds of lines of C++ (CUDA) or Python (Triton) code, demanding significant manual effort.
0:03:46 Mirage - Solution: Mirage is a multi-level superoptimizer that uses a holistic approach for optimizations across the GPU hierarchy, supporting algebraic and schedule transformations, and custom kernel generation.
0:05:05 Mirage - Key Benefits: Requires less engineering effort (few lines of Python code), achieves better performance (up to 3.3x speedup), and is easily adaptable to new architectures and models.
0:05:42 MuGraph IR: Mirage uses a multi-level graph representation (mugraph) that captures the GPU's memory hierarchy (kernel, thread block, and thread levels) to represent computations.
0:09:25 Exhaustive Search & Pruning: Mirage performs an exhaustive search to generate possible mugraphs, guided by "expression-guided pruning" using "abstract expressions" to manage the large search space.
0:13:48 Abstract Expressions: These abstract away index details in algebraic expressions, making them easier for an automated theorem prover to reason about for pruning, significantly improving search scalability.
0:25:43 Probabilistic Equivalence Verifier: Ensures the functional correctness of generated mugraph candidates using random testing in a finite field to avoid numerical issues.
00:27:23 MuGraph Optimizer: Performs post-verification optimizations like tensor layouts, memory planning, and operator scheduling, which do not affect the mugraph's correctness.
0:29:25 Case Study (QGA on H100): Mirage can generate highly customized kernels, like those for group query attention, achieving 2.2 times faster performance on H100 by leveraging new inter-thread block communication features.
0:30:51 Mirage Status: Open-source project, with a paper presented at OSDI.
0:43:28 MPK - Mega Kernel for LLM Serving (Follow-up): Shinhao Chang introduces MPK, aiming to compile entire large models into a single mega kernel for LLM serving.
0:45:51 Limitations of "Kernel Per Layer": Lack of inter-layer pipelining (kernels act as static barriers), difficulty overlapping compute and communication (coarse-grain data dependencies), and high kernel launch overheads for dynamic workloads.
0:48:11 Mega Kernel Advantages: Provides better inter-layer pipelining, allows overlapping compute and communication (fine-grain data dependencies), enables better load balancing, and minimizes kernel launch overhead for dynamic workloads.
0:51:44 Mega Kernel Challenges: Managing complex data dependencies, handling various forms of dynamism (continuous batching, mixed prefill/decode phases, speculative decoding), and optimizing for end-to-end performance.
0:53:22 MPK Solutions: Uses a "task graph" to manage dependencies, an internal runtime system to handle dynamism, and Mirage's superoptimizer to generate highly optimized tasks.
0:54:02 MPK Benefits: Reduces engineering effort (Python code for model architecture), delivers high performance (1.2 to 6.7 times faster than existing service systems like TorchServe), and offers day-zero support for new models.
0:55:16 Task Graph: Interleaves "tasks" (units of workload on a single Streaming Processor, compute or communication) and "events" (synchronization mechanisms among tasks) at a finer granularity than CUDA graphs.
0:58:28 MPK Runtime System: Employs schedulers (run on warp, manage event queues) and workers (run on SM, execute tasks) to gradually launch tasks and handle events while maintaining data dependencies, all within the mega kernel.
1:02:08 Performance (Q13B on A100): MPK reduces per-token latency from 14.5ms to 12.5ms, pushing software performance towards the theoretical hardware limit of ~10ms.
1:03:00 MPK Status: Open-source project, integrated with Mirage.
1:09:50 Register Spillage: Mega kernels can face register spillage issues; MPK addresses this by separating schedulers and workers into different kernels or by careful register allocation.
1:11:17 Tooling: Currently, there are only statistical results during search; no visualizer or logger for the entire search and transpilation process is available yet, but a mugraph visualizer exists.