近日，CMU 助理教授贾志豪（Zhihao Jia）团队创新玩法，推出了一个名为「Mirage Persistent Kernel（MPK）」的编译器，可以自动将 LLM 转化为优化的巨型内核（megakernel），从而将 LLM 推理延迟降低 1.2 到 6.7 倍。 在这种设计中，系统仅启动一个 GPU 内核来执行整个模型 —— 从逐层计算到 GPU 间通信 —— 整个过程无需中断。这种方法提供了以下几个关键的性能优势： 消除内核启动开销：通过避免重复的内核调用，即使是在多 GPU 环境下，也能消除内核启动开销；实现跨层软件 pipeline 允许内核在计算当前层的同时，开始为下一层加载数据；重叠计算与通信：由于巨型内核可以同时执行计算操作和 GPU 间通信，从而隐藏通信延迟。 现有的高级 ML 框架 —— 如 PyTorch、Triton 和 TVM，它们本身并不支持端到端巨型内核生成。此外，现代 LLM 系统由各种不同的专用内核库构建而成：用于通信的 NCCL 或 NVSHMEM，用于高效注意力计算的 FlashInfer 或 FlashAttention，以及用于自定义计算的 CUDA 或 Triton。 那么能否通过编译自动化这个过程呢？受到这个问题的启发，来自 CMU、华盛顿大学、加州大学伯克利分校、英伟达和清华大学的团队开发出了 MPK—— 一个编译器和运行时系统，它能自动将多 GPU 的 LLM 推理转换为高性能的巨型内核。MPK 释放了端到端 GPU 融合的效能优势，同时只需要开发者付出极小的手动努力。 MPK 的一个关键优势在于：通过消除内核启动开销，并最大程度地重叠跨层的计算、数据加载和 GPU 间通信，实现了极低的 LLM 推理延迟。 除了单 GPU 优化，MPK 还将计算与 GPU 间通信融合进一个单一的巨型内核。 这种设计使得 MPK 能够最大程度地重叠计算与通信。因此，MPK 相对于当前系统的性能提升随着 GPU 数量的增加而增大，使其在多 GPU 部署场景下尤为高效。 Part 1：MPK 编译器，其将 LLM 的计算图转化为优化的任务图；Part 2：MPK 运行时系统，该系统在单个巨型内核内执行任务图，以实现高吞吐量与低延迟。 LLM 的计算过程通常表示为计算图，其中每个节点对应一个计算算子（如矩阵乘法、注意力机制）或集合通信原语（如 all-reduce），边表示算子间的数据依赖关系。现有系统通常为每个算子启动独立的 GPU 内核。 然而，这种「单算子单内核」的执行模型难以实现 pipeline 优化，因为依赖