We introduce 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 introduces an SM-level graph representation that captures data dependencies at the granularity of individual streaming multiprocessors (SMs), enabling cross-operator software pipelining, \rev{fine-grained overlap of computation and communication, and other optimizations that are infeasible under the conventional kernel-per-operator execution model}. The MPK compiler lowers tensor programs into optimized SM-level task graphs and generates fast CUDA implementations for each task, while the MPK in-kernel parallel runtime executes these tasks within a single persistent 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, achieving up to 1.7$\times$ lower end-to-end inference latency and pushing LLM inference performance close to the limits of the underlying hardware. MPK is publicly available at https://github.com/mirage-project/mirage.
翻译:我们提出Mirage持久化核(MPK),这是首个能够自动将多GPU模型推理转换为单个高性能兆核的编译器与运行时系统。MPK引入了一种SM级图表示方法,以单个流式多处理器(SM)为粒度捕获数据依赖关系,从而支持跨算子软件流水线、计算与通信的细粒度重叠,以及传统“每算子单核”执行模型下不可行的其他优化。MPK编译器将张量程序降级为优化的SM级任务图,并为每个任务生成高效的CUDA实现,而MPK内核内并行运行时则通过跨SM的分散式调度在单个持久化兆核内执行这些任务。这些组件共同在最小化开发者工作量的前提下提供端到端内核融合,同时保留现有编程模型的灵活性。评估结果表明,MPK显著优于现有的每算子单核大语言模型(LLM)服务系统,将端到端推理延迟降低高达1.7倍,并使LLM推理性能逼近底层硬件极限。MPK已在https://github.com/mirage-project/mirage开源。