MPK: A Compiler and Runtime for Mega-Kernelizing Tensor Programs

Mirage Persistent Kernel (MPK) is a novel compiler and runtime system that automatically transforms multi-GPU tensor programs into a single high-performance mega-kernel using SM-level graph representations to enable cross-operator pipelining and fine-grained overlap of computation and communication, thereby significantly reducing LLM inference latency compared to traditional kernel-per-operator approaches.

The Gist

Imagine you are running a massive, high-speed factory that builds complex structures (like AI models). In the current way of doing things, every single step of the construction process is handled by a different specialized team.

The Old Way: The "Stop-and-Go" Factory
Right now, most AI systems work like a factory where Team A builds a wall, then shouts "Done!" and stops. The factory manager (the computer's operating system) has to pause everything, check the paperwork, and then call in Team B to paint the wall. Once Team B finishes, they stop, and the manager calls Team C to install the windows.

This creates a lot of wasted time:

1. The "Stop" Signal: Every time a team finishes, there's a mandatory pause to make sure everyone is ready for the next team. This is like a red light at every intersection.
2. The Manager's Busy Work: The manager spends too much time running back and forth between teams, handing out new instructions, instead of letting the teams work.
3. No Overlap: Even if Team B only needs a small part of the wall that Team A just finished, Team B has to wait for the entire wall to be built before they can start. They can't start painting just the first brick while the rest of the wall is still being built.

The New Way: MPK (Mirage Persistent Kernel)
The paper introduces MPK, which is like turning that factory into a single, continuous, super-efficient assembly line run by one giant, self-managing team.

Here is how MPK changes the game, using simple analogies:

1. The "One Big Kernel" (The Mega-Factory)

Instead of calling in different teams one by one, MPK launches one single, massive team that stays on the job from start to finish. This team doesn't stop between steps. They don't wait for a manager to tell them what to do next; they just keep moving.

• The Benefit: It eliminates the "stop-and-go" delays. It's like a train that never stops at stations to switch engines; it just keeps rolling.

2. The "SM-Level Map" (The Detailed Blueprint)

The paper introduces a new way of drawing the blueprint, called a tGraph.

• Old Blueprint: "Build Wall A, then Paint Wall A." (Too big and vague).
• MPK Blueprint: "Worker 1 builds Brick 1, Worker 2 builds Brick 2, Worker 3 paints Brick 1 while Worker 4 builds Brick 3."
• The Magic: MPK breaks the work down to the level of individual workers (called SMs or Streaming Multiprocessors). It knows exactly which worker needs which piece of data and when. This allows different workers to do different jobs (like building and painting) at the exact same time, as long as they don't step on each other's toes.

3. The "Self-Driving" Runtime (The Smart Foreman)

Inside this giant team, there is a smart, self-managing system (the In-Kernel Runtime).

• No Middleman: Instead of a manager outside the factory shouting instructions, the workers talk directly to each other.
• Just-in-Time vs. Ready-to-Go:
  • If a task is tricky and depends on unpredictable things (like how long a sentence is in a chat), the system waits until the previous step is actually done before starting the next one (Just-in-Time).
  • If a task is predictable, the system lines it up and gets it ready before the previous step is even finished (Ahead-of-Time).
• The Result: The workers never sit idle waiting for instructions. They are always busy.

4. The "Paged Shared Memory" (The Shared Toolbox)

In the old factory, each team had their own locked toolbox. If Team A needed a hammer, they had to wait for Team B to finish and put the hammer back.

• MPK's Solution: They use a "Paged Shared Memory." Imagine a giant, shared toolbox where tools are organized into small, movable pages. As soon as a worker finishes with a tool, they put it back on the shelf, and the next worker can immediately grab it. This allows workers to grab materials for their next job while they are still finishing their current job.

Why Does This Matter?

The paper tested this system on Large Language Models (the brains behind AI chatbots) using powerful NVIDIA GPUs.

• Speed: MPK made the AI run 1.7 times faster than the best existing systems in some cases.
• Latency: It reduced the time it takes to generate a single word of text, bringing it very close to the physical limits of the hardware.
• Ease of Use: The best part? You don't need to be a factory expert to use it. You can take a standard AI model (written in PyTorch) and "MPK-ify" it with just a few lines of code. The system automatically figures out how to break the work down and manage the workers.

In Summary:
MPK takes a chaotic factory where teams stop and start constantly, and turns it into a smooth, continuous, self-managing assembly line. By breaking work down to the smallest possible pieces and letting the workers coordinate directly, it removes all the wasted time, making AI inference significantly faster and more efficient.

Technical Summary

Technical Summary: MPK - A Compiler and Runtime for Mega-Kernelizing Tensor Programs

1. Problem Statement

Current high-performance GPU inference systems for large language models (LLMs) and other tensor programs predominantly rely on a kernel-per-operator execution model. In this paradigm, each operator (e.g., matrix multiplication, attention, all-reduce) is executed as a distinct GPU kernel. While this approach simplifies programming, it introduces three critical limitations that hinder performance:

1. Implicit Kernel Barriers: Conventional GPU execution enforces a barrier between consecutive kernel launches on the same stream. This prevents cross-operator software pipelining, forcing dependent operators to execute strictly sequentially and creating "pipeline bubbles" where hardware resources (Tensor Cores, TMAs) remain idle.
2. Coarse-Grained Overlap: Dependencies are enforced at the granularity of entire operators. Consequently, fine-grained overlap between computation and communication is impossible. For instance, an AllReduce operation must wait for the entirety of a preceding MatMul to complete, even though specific fragments of the reduction depend only on specific fragments of the multiplication output.
3. Launch Overhead and Static Flexibility: Executing hundreds or thousands of kernels per inference iteration incurs significant launch overhead. While CUDA Graphs mitigate this, they are largely static; dynamic changes in control flow, tensor shapes, or data dependencies often require re-instantiating the graph, limiting flexibility for dynamic workloads like LLM serving.

Existing solutions, such as manual mega-kernels (e.g., FlashDMoE, Spector et al.), offer performance benefits but require extensive engineering effort and deep GPU expertise, failing to generalize across models or hardware.

2. Methodology

The authors present Mirage Persistent Kernel (MPK), the first compiler and runtime system that automatically transforms multi-GPU model inference into a single, high-performance mega-kernel. MPK achieves this through two primary components:

A. SM-Level Graph Representation (tGraph)

MPK introduces a novel graph representation called tGraph that captures computation and communication at the granularity of individual Streaming Multiprocessors (SMs), rather than the entire GPU.

• Nodes: Represent tasks (computation or communication) running on a single SM.
• Edges: Encode fine-grained dependencies between tasks via events.
• Advantage: This granularity exposes parallelism inaccessible to kernel-level models, enabling cross-operator software pipelining and fine-grained overlap of computation and communication.

Compiler Optimizations:
The MPK compiler transforms a standard tensor program into an optimized tGraph through several stages:

1. Operator Decomposition: Splits operators into disjoint sets of tasks based on output tensor partitioning (e.g., tiling matrix multiplication rows/columns) to match the number of available SMs.
2. Dependency Analysis: Enumerates producer-consumer pairs between tasks to insert synchronization events only where data overlaps exist.
3. Event Fusion: Merges redundant synchronization points using successor-set fusion (merging events with identical consumer sets) and predecessor-set fusion (merging events with identical producer sets) to reduce graph complexity.
4. Normalization: Transforms the graph so that every task has at most one dependent event and one triggering event. This is achieved by inserting dummy tasks and events to handle fan-in/fan-out, ensuring a uniform, low-overhead representation.
5. Linearization: Uses a BFS-based algorithm to assign contiguous indices to tasks triggered by the same event. This allows the runtime to encode event fan-out as a simple index range (first/last task ID) rather than an explicit list, significantly reducing device memory footprint.
6. Task Implementation Generation: Uses compiler superoptimization (Mirage) to automatically generate high-performance CUDA implementations for each individual task, including intra-SM optimizations like register reuse and shared memory layout.

B. In-Kernel Parallel Runtime

MPK executes the tGraph within a single persistent mega-kernel using a decentralized, event-driven architecture:

• Worker/Scheduler Partitioning: The GPU's SMs are statically partitioned into Workers (which execute tasks) and Schedulers (which manage event queues).
• Event-Driven Execution: Schedulers poll for activated events (where all prerequisites are met) and dispatch corresponding tasks to workers. This eliminates the need for host-side synchronization or repeated kernel launches.
• Hybrid Task Launch: To balance latency and load balancing, MPK employs two modes:
  • Just-In-Time (JIT): Tasks are launched only after their dependencies are met. This adapts to dynamic workloads (e.g., variable sequence lengths in attention) to prevent load imbalance.
  • Ahead-Of-Time (AOT): Tasks are pre-enqueued into worker queues before dependencies are met. This reduces dispatch latency for predictable operators.
• Paged Shared Memory: To enable cross-task software pipelining, MPK introduces a paged abstraction for shared memory. Tasks acquire and release fixed-size pages, allowing the next task to begin prefetching data into shared memory while the current task computes, without resource conflicts.

3. Key Contributions

1. First Automated Mega-Kernel System: MPK is the first system to automatically transform arbitrary multi-GPU tensor programs into a single mega-kernel, removing the need for manual kernel fusion.
2. SM-Level Abstraction: The introduction of the tGraph representation shifts dependency tracking from the operator level to the SM-task level, enabling optimizations (pipelining, fine-grained overlap) previously infeasible under the kernel-per-operator model.
3. End-to-End Compiler and Runtime: The system integrates a compiler that optimizes task graphs and generates CUDA code with a runtime that executes these tasks asynchronously within a single kernel, supporting dynamic workloads like continuous batching and paged attention.
4. Hybrid Execution Strategy: The combination of JIT and AOT task launching allows the system to handle both the dynamic nature of LLM inference (variable batch/sequence lengths) and the static efficiency of fixed operators.

4. Evaluation Results

The authors evaluated MPK on five widely used LLMs (including Qwen3 and Llama-3 variants) across three generations of NVIDIA GPUs (A100, H100, B200), comparing it against state-of-the-art kernel-per-operator systems (SGLang, vLLM) and vanilla PyTorch.

• Performance Gains: MPK achieves 1.0× to 1.7× lower end-to-end inference latency compared to SGLang and vLLM. The improvements are most pronounced for smaller models and newer hardware generations (e.g., B200), where kernel launch overhead and pipeline bubbles become more significant relative to computation time.
• Throughput: In multi-GPU settings (up to 8 H100s), MPK improves throughput by 1.1× to 1.4× over optimized baselines.
• Latency Limits: On an A100, MPK reduced per-token decoding latency for Qwen3-8B from 14.5 ms (vLLM/SGLang) to 12.5 ms, approaching the estimated hardware lower bound of ~10 ms.
• Ablation Studies:
  • Cross-Task Pipelining: Reduced task runtime by 1.2–1.3×.
  • Compute-Communication Overlap: Reduced per-iteration latency by 1.1×.
  • Kernel Launch Reduction: Eliminated the ~1.1 ms per-token overhead associated with 293 kernel launches per token in baseline systems.
• Scalability: MPK scales effectively across 2, 4, and 8 GPUs, maintaining speedups over baselines.

5. Significance and Claims

The paper claims that MPK represents a paradigm shift in GPU inference systems by unifying the execution pipeline into a single mega-kernel. Its significance lies in:

• Breaking the Kernel Barrier: By moving synchronization to the SM level, MPK unlocks hardware utilization potential that is currently inaccessible to kernel-per-operator systems.
• Automation vs. Manual Effort: Unlike previous mega-kernel approaches that required expert manual design, MPK automates the process, making high-performance mega-kernel execution accessible to standard PyTorch developers with minimal code changes.
• Hardware Proximity: The system pushes LLM inference performance close to the theoretical limits of the underlying hardware by minimizing software overheads (launch, scheduling, synchronization) and maximizing hardware overlap.

The authors conclude that MPK opens a new path toward high-performance, compiler-driven inference systems that preserve the flexibility of existing programming models while delivering near-hardware-limit performance.