Mirage Persistent Kernel: A Compiler and Runtime for Mega-Kernelizing Tensor Programs
Fuente:
arXiv
Gespeichert in:
| Hauptverfasser: | , , , , , , , , , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Veröffentlicht: |
2025
|
| Schlagworte: | |
| Online-Zugang: | |
| Tags: |
Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
|
| _version_ | 1866914221721124864 |
|---|---|
| author | Cheng, Xinhao Zhang, Zhihao Zhou, Yu Ji, Jianan Jiang, Jinchen Zhao, Zepeng Xiao, Ziruo Ye, Zihao Huang, Yingyi Lai, Ruihang Jin, Hongyi Hou, Bohan Wu, Mengdi Dong, Yixin Yip, Anthony Ye, Zihao Wang, Songting Yang, Wenqin Miao, Xupeng Chen, Tianqi Jia, Zhihao |
| author_facet | Cheng, Xinhao Zhang, Zhihao Zhou, Yu Ji, Jianan Jiang, Jinchen Zhao, Zepeng Xiao, Ziruo Ye, Zihao Huang, Yingyi Lai, Ruihang Jin, Hongyi Hou, Bohan Wu, Mengdi Dong, Yixin Yip, Anthony Ye, Zihao Wang, Songting Yang, Wenqin Miao, Xupeng Chen, Tianqi Jia, Zhihao |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2512_22219 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Mirage Persistent Kernel: A Compiler and Runtime for Mega-Kernelizing Tensor Programs Cheng, Xinhao Zhang, Zhihao Zhou, Yu Ji, Jianan Jiang, Jinchen Zhao, Zepeng Xiao, Ziruo Ye, Zihao Huang, Yingyi Lai, Ruihang Jin, Hongyi Hou, Bohan Wu, Mengdi Dong, Yixin Yip, Anthony Ye, Zihao Wang, Songting Yang, Wenqin Miao, Xupeng Chen, Tianqi Jia, Zhihao Distributed, Parallel, and Cluster Computing Machine Learning Programming Languages 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. |
| title | Mirage Persistent Kernel: A Compiler and Runtime for Mega-Kernelizing Tensor Programs |
| topic | Distributed, Parallel, and Cluster Computing Machine Learning Programming Languages |
| url | https://arxiv.org/abs/2512.22219 |