近日，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 推理工作负载方式的底子性转变，并热切等待与社区合作，共同推动这一愿景向前发展。