近日，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 优化，因为依赖关系是在整个内核的粗粒度层面强制执行的，而非实际数据单元层面。 典型案例如矩阵乘法（matmul）后接 all-reduce 操作：现有系统中，all-reduce 内核必须等待整个 matmul 内核完成。而实际上，all-reduce 的每个数据分块仅依赖 matmul 输出的局部结果。这种逻辑依赖与实际依赖的错配，严重限制了计算与通信的重叠潜力。 下图 2 展示了 MPK 编译器将 PyTorch 定义的 LLM 计算图转化为优化细粒度任务图，最大化暴露并行性。右侧展示次优方案 —— 其引入不必要的数据依赖与全局屏障，导致跨层流水线优化机会受限。 为了解决此问题，MPK 引入的编译器可将 LLM 计算图自动转化为细粒度任务图。该任务图在子内核级别显式捕获依赖关系，实现更激进的跨层流水线优化。 任务（矩形表示），代表分配给单个 GPU 流式多处理器（SM）的计算 / 通信单元。事件（圆形表示），表示任务间的同步点。触发机制，每个任务发出指向触发事件的边，该事件在关联任务全部完成后激活。依赖机制，每个任务接收来自依赖事件的边，表明事件激活后任务立即启动。 任务图使 MPK 能够发掘计算图中无法实现的 pipeline 优化机会。例如，MPK 可以构建优化任务图 —— 其中每个 all-reduce 任务仅依赖于生成其输入的对应 matmul 任务，从而实现分块执行与计算通信重叠。 MPK 包含内置 GPU 运行时系统，可在单个 GPU 巨型内核内完整执行任务图。这使得系统能在推理过程中无需额外内核启动的情况下，实现任务执行与调度的细粒度控制。 获取任务：从队列中提取下一待执行任务。执行计算：运行任务（如矩阵乘法 / 注意力机制 / GPU 间数据传输）。事件触发：任务完成后通知触发事件。循环执行：重复上述过程。 调度决策由 MPK 的分布式调度单元处理，每个调度单元运行于单个线程束（warp）上。由于每个流式多处理器（SM）可以容纳多个线程束，因此单 SM 最多可并发运行 4 个调度单元。每个调度单元维护激活事件队列，并持续执行以下操作： 下图 3 展示了 MPK 的执行时间线，其中每个矩形代表一个在工作单元上运行的任务；每个圆圈代表一个事件。当一个任务完成时，它会递增其对应触发事件的计数器。当事件计数器达到预设阈值时，该事件被视为已激活，并被加入调度单元的事件队列。随后，调度单元会启动所有依赖于该事件的下游任务。 由于所有的调度和任务切换都发生在单一内核上下文内，任务间的开销极低，通常仅需 1-2 微秒，从而能够高效地执行多层、多 GPU 的 LLM 工作负载。 团队对 MPK 的愿景是使巨型内核编译既易于使用又具备高性能。目前，你只需几十行 Python 代码（主要用于指定巨型内核的输入和输出）即可将一个 LLM 编译成一个巨型内核。此方向仍有广阔的探索空间，目前正在积极攻关的一些关键领域包括如下： 支持现代 GPU 架构。下一个里程碑是将 MPK 扩展到支持下一代架构，例如 NVIDIA Blackwell。一个主要挑战在于如何将线程束专业化，这是新型 GPU 的一项关键优化技术，与 MPK 的巨型内核执行模型相集成。处理工作负载动态性。MPK 目前构建的是静态任务图，这限制了它处理动态工作负载（如 MoE 模型）的能力。团队正在开发新的编译策略，使 MPK 能够在巨型内核内部支持动态控制流和条件执行。高级调度与任务分配。MPK 在任务级别解锁了新的细粒度调度能力。虽然当前的实现使用简单的轮询调度在流式多处理器（SM）之间分配任务，但团队看到了在高级调度策略（如优先级感知或吞吐量优化策略）方面令人兴奋的机会，可应用于诸如延迟服务等级目标（SLO）驱动的服务或混合批处理等场景。 团队相信，MPK 代表了在 GPU 上编译和执行 LLM 推理工作负载方式的根本性转变，并热切期待与社区合作，共同推动这一愿景向前发展。
