将多GPU推理融合为单个高性能巨内核,显著降低延迟。
MPK: A Compiler and Runtime for Mega-Kernelizing Tensor Programs
- 以SM为粒度建模数据依赖,实现跨算子软件流水
- 实测推理延迟降低1.7倍,逼近硬件性能极限
- 无需重构代码,适合现有LLM服务系统快速升级
我们提出镜像持久内核(MPK),首个可自动将多GPU模型推理转换为单个高性能巨内核的编译器与运行时系统。MPK引入SM级图表示,以流式多处理器(SM)为粒度捕捉数据依赖,支持跨算子软件流水、细粒度计算与通信重叠等传统按算子执行模型无法实现的优化。MPK编译器将张量程序降级为优化的SM级任务图,并生成各任务的高效CUDA实现;其内核级并行运行时在单一持久巨内核中,通过SM间去中心化调度执行这些任务。三者协同实现端到端内核融合,仅需极小开发成本,同时保持现有编程模型灵活性。评估表明,MPK显著优于现有按算子执行的LLM服务系统,端到端推理延迟最高降低1.7倍,使LLM推理性能接近底层硬件极限。MPK已开源:https://github.com/mirage-project/mirage。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。