Mirage Persistent Kernel: Fusing Multi-GPU LLM Inference into One Kernel Launch

The Problem: Kernel Launch Overhead in Multi-GPU LLM Inference

LLM inference on multiple GPUs typically launches hundreds or thousands of small kernels, one for each operation like attention, normalization, and linear layers. Each launch carries overhead, and inter-GPU communication adds synchronization points. Mirage Persistent Kernel (MPK) attacks this by compiling the entire inference computation into a single megakernel: one fused GPU kernel that performs all necessary computation and communication within a single launch. The README claims a 1.2x to 6.7x reduction in inference latency. The target user is an engineer who needs low-latency serving for large models, has access to multiple GPUs, and is willing to trade compile-time and manual tuning for runtime speed.

How MPK Works: Persistent Kernels and Schedulers

MPK is a compiler and runtime system. The core idea is a persistent kernel: a kernel that stays resident on the GPU across the entire inference loop, avoiding the cost of repeated launches. The runtime uses a task graph where each fused operation is a task. The PersistentKernel API takes parameters like num_workers, num_local_schedulers, and num_remote_schedulers. These must sum to the number of physical SMs: the README states that num_workers plus (num_local_schedulers + num_remote_schedulers) / 4 must match the SM count. This constraint is unusual and requires you to know your GPU's SM count. The kernel also requires two meta tensors: step, an integer array tracking the decoding step, and tokens, a tensor of shape [num_requests, seq_length] that stores prompts and generated tokens. The runtime increments step after each decoding iteration. This design suggests MPK is tightly coupled to autoregressive decoding, not arbitrary inference workloads.

Getting Started: Installation and the Qwen3 Demo

Installation is from source only, as of the latest README update. The commands are: git clone --recursive --branch mpk https://www.github.com/mirage-project/mirage, then cd mirage, then pip install -e . -v, and finally export MIRAGE_HOME=$(pwd). The README notes that binary wheels are in progress but not yet available. To test, you run a demo script for Qwen3-8B: python demo/qwen3/demo.py runs with native Triton and FlashInfer kernels, while python demo/qwen3/demo.py --use-mirage compiles and executes the megakernel. Adding --profiling visualizes the execution timeline. The demo script is the only example provided, so you must adapt it for other models. The compilation step, mpk.compile(), is explicit, and execution is a simple call to mpk(). This is a low-level API: you define tensors, attach PyTorch tensors, and specify grid_dim and block_dim for each fused layer.

The API: Defining Tensors and Fused Layers

The API is Python but feels like CUDA. You instantiate a PersistentKernel with world_size and mpi_rank for multi-GPU setup. You attach existing PyTorch tensors with mpk.attach_input(torch_tensor=..., name=...). You allocate new tensors with mpk.new_tensor(dims=..., dtype=mi.bfloat16, name=..., io_category='cuda_tensor' or 'nvshmem_tensor'). The nvshmem_tensor category is required for remote GPU access, such as during all-reduce. Then you chain fused operations. The README shows rmsnorm_linear_layer, which fuses RMSNorm and a Linear layer. You must specify grid_dim and block_dim for each operation, which control the number of thread blocks and threads per block. The README suggests that the total number of thread blocks should be a multiple of the number of workers to avoid outliers. This is a manual tuning knob that a higher-level framework would hide.

Limitations and Wrong-Tool Cases

MPK is not a drop-in replacement for your inference stack. The most obvious limitation is the source-only build. The README explicitly says binary wheels are not yet available, so you must compile from source, which can be a barrier in production environments. The SM count constraint is another limitation: if your GPU has an SM count that does not satisfy the worker/scheduler formula, you cannot use the default configuration. The API is also narrow: it provides fused operations like rmsnorm_linear_layer, but the README does not list a full set of operations. If your model uses an operation not covered, you may need to write custom CUDA or extend the framework. Finally, MPK is designed for multi-GPU inference; for single-GPU models, the overhead of managing remote schedulers and nvshmem tensors may be unnecessary, and the latency gains are likely smaller.

Alternative: Native Triton and FlashInfer Kernels

The demo script itself provides the alternative: running the same model with native Triton and FlashInfer kernels instead of MPK. The command python demo/qwen3/demo.py without --use-mirage uses these standard kernels. The difference in approach is that Triton and FlashInfer launch multiple kernels, one per operation, and rely on the framework's runtime to schedule them. MPK instead compiles the entire graph into a single persistent kernel. This is a fundamental design trade-off: MPK reduces launch overhead but requires a compilation step and manual scheduling parameters. For workloads where the kernel launch overhead is not the bottleneck, the alternative may be simpler and more portable. The README's latency claim of 1.2x to 6.7x is relative to this baseline, but the exact conditions are not specified.

Maintenance, License, and Upgrade Path

The project is under active development, with the default branch named mpk and a recent release tagged osdi26-ae-v1, which is the artifact for the OSDI 2026 paper. The last push was in June 2026, indicating ongoing work. The license is Apache-2.0, which permits commercial use, modification, and distribution, but you should review the full license text for any patent or trademark clauses. The upgrade path is not documented in the README; there is a roadmap linked in the repository issues, but no specific upgrade instructions. Given the pre-release nature, expect breaking changes between versions. The README mentions that binary wheels are being worked on, so future releases may simplify installation, but as of now, you must build from source and track the mpk branch.

Editorial conclusion

Adopt MPK if you run multi-GPU LLM inference on NVIDIA hardware, need low latency, and can tolerate a source-only build and a kernel API that requires manual tensor and grid configuration. Skip it if you need a stable, pip-installable package, support for non-NVIDIA GPUs, or a high-level interface that hides scheduling details. Before adopting, verify that your model architecture matches the fused operations MPK provides (e.g., rmsnorm_linear_layer), that your GPU's SM count can satisfy the worker/scheduler constraint, and that you are comfortable with the mpk branch's pre-release status. The project is actively developed, but the lack of binary wheels as of mid-2025 means you must build from source, so test the demo on your exact hardware first.

Official sources

1. Official documentation
2. Official README
3. Project repository
4. Release notes