Mirage Persistent Kernel: A Compiler and Runtime for Mega-Kernelizing Tensor Programs

文章标题: Mirage Persistent Kernel: 一个用于对张量程序进行超级内核化(Mega-Kernelizing)的编译器和运行时
作者/机构: Xinhao Cheng, Zhihao Zhang, Yu Zhou, Jianan Ji, Jinchen Jiang, Zepeng Zhao, Ziruo Xiao, Zihao Ye, Yingyi Huang, Ruihang Lai, Hongyi Jin, Bohan Hou, Mengdi Wu, Yixin Dong, Anthony Yip, Zihao Ye, Songting Wang, Wenqin Yang, Xupeng Miao, Tianqi Chen, Zhihao Jia
机构: 卡内基梅隆大学, 清华大学, 英伟达, 密歇根大学, 独立研究员, 普渡大学

A1 主要贡献

本文介绍了一种名为 Mirage Persistent Kernel (MPK) 的系统，它是首个能将多GPU模型推理自动转换为单个高性能超级内核（mega-kernel）的编译器和运行时系统。

核心问题: 现有机器学习系统普遍采用“每个算子一个内核”（kernel-per-operator）的执行模型，这种模型存在以下几个关键限制：
1. 无法实现跨算子软件流水线：现代GPU在同一流上连续启动的内核之间会隐式插入内核屏障（kernel barrier），这强制依赖的算子必须严格串行执行，阻碍了跨算子的软件流水线优化。
2. 无法实现细粒度的计算-通信重叠：由于依赖关系仅在粗粒度的算子级别捕获，因此运行时必须等待整个算子完成后才能启动依赖的通信或计算任务。例如，矩阵乘法后的all-reduce操作必须等待整个矩阵乘法完成，即使all-reduce的每个部分只依赖于矩阵乘法结果的一个子集。
3. 内核启动开销和灵活性受限：每次推理迭代需要启动成百上千个内核，尽管CUDA Graphs可以缓解启动开销，但其本质上是静态的，难以适应模型推理中常见的动态工作负载（如控制流、张量形状或数据依赖的变化）。

研究目标与创新点: 为了克服上述限制，MPK提出将所有计算和通信融合成一个单一的超级内核。其核心思想是在单个流式多处理器（SM）的粒度上表示计算和GPU间通信，而不是整个GPU。为此，MPK引入了以下创新：

SM级图表示（tGraph）：MPK引入了一种名为tGraph的SM级图表示。tGraph的节点表示在单个SM上运行的任务，边则编码了任务间的细粒度依赖关系。这种表示方法暴露了额外的并行性，并使得跨算子软件流水线和细粒度内核重叠等优化成为可能，而这些在传统的kernel-per-operator方法中是不可行的。
MPK编译器：该编译器接收一个张量程序和推理配置作为输入，并自动将其计算图转换为一个高度优化的、针对特定推理配置和GPU架构的SM级tGraph。编译器通过事件融合、图归一化和图线性化等多种优化来减少同步开销并最大化tGraph的性能。此外，MPK还利用现有的超优化技术为每个任务自动生成高性能的CUDA实现。
核内并行运行时：MPK使用一个完全嵌入在超级内核内部的核内并行运行时来执行SM级的tGraph。该运行时将GPU的SMs划分为执行任务的“工作者”（workers）和维护任务间依赖关系并调度任务的“调度器”（schedulers）。运行时采用事件驱动、完全异步的执行模型，并通过结合即时（just-in-time）和提前（ahead-of-time）分派的混合任务启动策略，以最小化运行时开销同时保持SM间的动态负载均衡。

主要成果: MPK被实现为一个PyTorch内核后端，用户只需几行代码即可将PyTorch模型超级内核化。评估表明，与现有的kernel-per-operator LLM服务系统（如SGLang和vLLM）相比，MPK在单GPU和多GPU部署上均能将端到端推理延迟降低高达1.7倍，将LLM推理性能推向硬件极限。

A3 背景知识

2.1 GPU编程模型

GPU计算组织形式。在GPU上，计算被组织成内核（kernel），每个内核代表一个以单程序多数据（SPMD）方式在众多核心上并发执行的函数。一个内核由一个线程块（thread block）网格组成，每个线程块被分配到一个流式多处理器（SM）上，并包含多个操作单个数据元素的线程。

GPU内存与同步模型。每个线程拥有私有的寄存器文件，而同一线程块内的线程通过低延迟的共享内存进行数据交换和集体操作。所有内核的输入和输出都存储在GPU设备内存中。CUDA编程模型不支持内核内跨线程块的直接同步，因为现代GPU架构缺乏此类协调的硬件机制。因此，跨算子的依赖关系（例如，一个矩阵乘法必须在另一个开始前完成）是通过内核屏障（kernel barrier）来强制执行的，GPU运行时会自动在同一流上连续启动的内核之间插入这些屏障。

内核屏障对优化的限制。虽然内核屏障简化了依赖管理，但它们也阻碍了关键的GPU优化，例如跨内核软件流水线和细粒度的内核重叠。

软件流水线。GPU架构日益异构化，集成了张量核心和张量内存加速器（TMA）等专用加速器。由于TMA的加载和存储指令是异步执行的，数据移动可以与张量核心和CUDA核心的计算同时进行。为了充分利用这些加速器，需要软件流水线技术——该技术将计算和数据移动的独立阶段在任务的多次迭代中交错进行，以最大化硬件利用率。现有系统实现了任务内流水线（如图2a所示），其中单个任务被分解为多次迭代，使得TMA、张量核心和CUDA核心可以同时为不同迭代执行数据传输、矩阵计算和辅助操作。然而，内核屏障将流水线限制在单个任务内部，无法实现任务间流水线，并引入了导致硬件资源未被充分利用的流水线气泡（pipeline bubbles）。

细粒度内核重叠。内核屏障也排除了重叠使用不同硬件资源（例如计算和通信）的内核的机会，因为它们在整个内核的粒度上强制执行依赖关系，而不是在单个数据单元的粒度上。图3a展示了大型语言模型（LLM）中的一个常见模式，其中MatMul算子后跟着一个AllGather算子。现有系统通常将它们作为两个独立的内核启动，这要求AllGather内核的所有线程块必须等到前一个MatMul内核的所有线程块都完成后才能开始。实际上，AllGather和MatMul之间的数据依赖关系是更细粒度的：由于AllGather执行逐元素操作，它的每个线程块仅依赖于单个MatMul线程块的输出。这种依赖结构使得细粒度内核重叠成为可能，不同的SM可以并行执行MatMul和AllGather，只要细粒度的数据依赖关系得到维护即可。正如先前工作【41，{NanoFlow}: Towards optimal large language model serving throughput, 2025, 19th USENIX Symposium on Operating Systems Design and Implementation (OSDI 25)】所指出的，这种重叠允许系统同时利用现代GPU上的计算和通信带宽。然而，实现这种重叠需要SM之间在亚内核粒度上进行适当的同步（如图3b所示），而这是传统内核屏障所不支持的。

图3: MPK 与现有方法在支持任务间细粒度内核重叠方面的比较。数据依赖（图3b中的黑箭头）确保了正确性。

2.2 内核融合

内核融合的概念与优势。内核融合通过将多个在相同数据上顺序执行的GPU内核合并成一个语义上等效的单一内核，从而消除了内核屏障。内核融合通过避免实例化中间结果、减少设备内存访问和消除内核启动开销来提高性能。

现有内核融合的局限性。内核融合已在张量程序编译器中被广泛采用。PyTorch、JAX和TVM等框架使用基于规则的启发式方法来融合相邻的内核【14，JAX: composable transformations of Python+NumPy programs, 2018】【15，TVM: end-to-end optimization stack for deep learning, 2018, CoRR】【25，Tensors and Dynamic neural networks in Python with strong GPU acceleration, 2017, https://pytorch.org】，而像Mirage和TASO这样的系统则通过编译器超优化自动发现融合规则【22 ，Taso: Optimizing deep learning computation with automatic generation of graph substitutions, 2019, Proceedings of the 27th ACM Symposium on Operating Systems Principles】【33，Mirage: A {MultiLevel} superoptimizer for tensor programs, 2025, 19th USENIX Symposium on Operating Systems Design and Implementation (OSDI 25)】。然而，现有的编译器只能融合小规模的局部算子组，因为生成一个能够忠实实现复杂张量程序的单一内核在计算上非常困难，并且通常是不可行的。

图4: MPK 编译器将一个内核级计算图转换为一个优化的 SM 级 tGraph。MM、AT 和 AR 分别表示 MatMul、Attention 和 AllReduce 任务。

超级内核范式及其挑战。超级内核（mega-kernel）范式将内核融合推向极致，它将一个张量程序的所有计算和通信都融合到一个持久化内核中，并使用设备内存同步原语来协调跨SM的执行。尽管具有性能优势，但目前的ML编译器（如PyTorch、Triton、JAX和TVM）不支持超级内核编译。现有的超级内核是由GPU专家为特定模型手工制作的。例如，FlashDMoE将混合专家（mixture-of-experts）计算和GPU间通信融合成一个单一内核【13，Flashdmoe: Fast distributed moe in a single kernel, 2025】，而Spector等人则为LLAMA-1B手动设计并实现了一个低延迟的超级内核【9，Designing a Low-Latency Megakernel for Llama1B, 2025, https://hazyresearch.stanford.edu/blog/ 2025-05-27-no-bubbles】【30，Llama: Open and efficient foundation language models, 2023, arXiv preprint arXiv:2302.13971】。这些手动方法需要大量的工程努力和深厚的GPU专业知识。相比之下，MPK采用一种基于编译器的方法，自动将张量程序转换为优化的超级内核，从而无需人工操作。

3 SM级图表示

tGraph表示法。本节介绍了tGraph，这是一种在单个流式多处理器（SM）粒度上表达张量程序计算的表示方法。这种细粒度的表示暴露了额外的并行性，并使得跨算子软件流水线和细粒度内核重叠等优化成为可能，而这些在现有的“每个算子一个内核”执行模型中都是无法支持的。

tGraph的结构。图4展示了一个tGraph示例，其中每个节点代表一个任务或一个事件。每个任务（以蓝色或橙色矩形表示）表示在一个SM上执行的计算或通信单元。每个事件（以绿色圆形表示）代表任务间的同步。在图中，任务和事件交替出现：每个任务只有指向其触发事件的出边和来自其依赖事件的入边。当一个任务的所有依赖事件都被激活时，该任务就准备好执行，并在完成后通知其触发事件。一个事件一旦收到了所有与之关联的任务的通知，就会被激活。

tGraph的优势。这种结构以比传统计算图更精细的粒度捕获依赖关系。例如，多GPU LLM服务通常涉及一个MatMul算子后跟一个AllReduce算子（图4a）。由于粗粒度的内核屏障同步了整个内核，现有系统通常会顺序执行这些算子。相比之下，SM级任务图可以表示精确的任务级依赖关系：由于AllReduce执行逐元素通信，其每个任务仅依赖于一个对应的MatMul任务。通过在依赖的任务对之间插入细粒度的事件，MPK可以重叠计算密集型的MatMul任务和通信密集型的AllReduce任务，从而最大化整体GPU利用率。

tGraph的等价表示与优化。多个tGraph可能表示同一个计算图。图4c展示了一个替代但次优的tGraph，其中事件仅捕获算子级别的依赖关系，类似于传统的内核屏障。第4节将描述MPK如何通过推断精确的数据依赖来生成高性能的任务图，以最大化并发性并最小化同步开销。

与CUDA Graphs的比较。tGraph可以被看作是CUDA Graphs的底层扩展，两者共享一些结构上的相似性。与CUDA Graphs一样，tGraph是静态实例化的，并编码了操作之间的显式依赖关系。然而，CUDA Graphs仅在内核级别捕获依赖关系，而tGraph则在单个SM任务和亚内核事件的粒度上操作。CUDA Graphs主要描述内核启动顺序，并依赖流语义进行同步，这限制了重叠和融合只能在内核边界内进行。相比之下，tGraph显式地建模了任务内和跨算子的依赖关系，从而实现了跨SM的细粒度同步以及在单个内核内计算和通信的重叠。这种设计使得MPK能够利用CUDA Graphs或内核级执行模型无法访问的并行性。

A2 方法细节

本节介绍了MPK编译器，它以一个计算图和相关的推理配置作为输入，并生成一个针对目标配置和底层GPU架构进行优化的tGraph。图5概述了端到端的编译工作流程。

4.1 tGraph生成

算子分解。MPK编译器通过划分算子的输出张量，将输入计算图中的每个算子分解为一组任务，使得所有任务计算输出的不相交子集，从而可以在SM上并行执行。大多数张量代数算子可以沿多个输出维度进行划分；例如，矩阵乘法的输出张量可以沿行和列维度进行切片，以暴露并行化机会。

分区策略的选择。分区策略的性能取决于问题规模和目标GPU架构。为了找到一个有效的策略，MPK选择一种能够最小化从设备内存到共享内存数据加载的分区策略，因为访问设备内存比访问共享内存或在CUDA核心或张量核心上执行计算要昂贵得多。默认情况下，MPK生成的任务数量与SM数量成正比，以促进执行期间SM间的负载均衡。MPK还提供了一个接口，用户可以通过设置每个输出维度上期望的并行度来轻松指定自定义分区策略。

依赖分析。MPK使用事件来捕获任务之间的依赖关系。对于共享一个张量的任意两个算子，MPK会枚举这两个算子的所有任务对。如果任务t1产生的输出区域与任务t2消耗的输入区域重叠，MPK就会为任务对(t1,t2)引入一个事件e。这个事件作为一个同步点，表明t2在t1完成生产所需数据之前不能开始执行。相应地，MPK会在生成的tGraph中插入两条边(t1, e)和(e,t2)。这种细粒度的依赖分析保留了所有的生产者-消费者依赖关系，同时暴露了独立任务之间的最大并行性。

事件融合。MPK应用两种互补的事件融合形式——后继集（successor-set）融合和前驱集（predecessor-set）融合——来消除冗余的同步点并简化构建的tGraph。对于一个事件e，我们定义了两个函数：InTasks(e)，即触发e的任务集合；以及OutTasks(e)，即依赖于e的任务集合。这些函数使我们能够描述何时多个事件表现出相同的依赖结构，从而可以被融合。

后继集融合。首先，后继集融合会合并那些作为相同消费者任务集先决条件的事件。因为这些任务在所有此类事件被激活之前都不能开始执行，所以将它们分开表示并不能提供额外的调度灵活性。
定义 4.1 (后继集融合)。对于一个tGraph中的任意两个事件e1和e2，当且仅当OutTasks(e1)=OutTasks(e2)时，可以应用后继集融合。MPK会从图中移除事件e1和e2，并引入一个融合后的事件e'，其InTasks(e') = InTasks(e1) ∪ InTasks(e2)且OutTasks(e') = OutTasks(e1)。
例如，在图5(b)中，后继集事件融合将事件e10和e14融合成一个新的事件（即图5(c)中的e4），因为它们都是任务O1的先决条件。

前驱集融合。其次，前驱集融合会合并那些依赖于相同生产者任务集的事件。因为这类事件是同时被触发的，将它们作为独立的同步节点维护会引入不必要的图复杂性。
定义 4.2 (前驱集融合)。对于一个tGraph T中的任意两个事件e1和e2，当且仅当InTasks(e1)=InTasks(e2)时，可以应用前驱集融合。MPK会从T中移除事件e1和e2，并引入一个融合后的事件e'，其InTasks(e')=InTasks(e1)且OutTasks(e')=OutTasks(e1) ∪ OutTasks(e2)。
例如，前驱集融合将图5(c)中的事件e4、e5、e6和e7融合成一个单一的新事件（新tGraph中的e4），因为所有这些事件都依赖于任务A1和A2。

tGraph表示的挑战与解决方案。MPK必须解决的一个核心挑战是如何表示任务和事件之间的依赖关系。由于MPK在SM上并行执行任务和更新事件，运行时需要一种统一且廉价的表示方式，以避免昂贵的间接索引。这里出现了两个挑战。首先，一个任务可能依赖并触发任意数量的事件。一个直接的方法是为每个任务预留最大数量的依赖事件和触发事件的空间，但这会导致显著的内存开销。其次，经过事件融合后，一个事件可能触发任意数量的任务。通过为每个事件分配最大扇出（fan-out）的空间来表示这些出边同样代价高昂。MPK引入了两种技术来解决这些挑战：tGraph归一化和线性化。

tGraph归一化。MPK通过tGraph归一化来解决第一个挑战，该过程将输入的tGraph转换为功能上等效的形式，其中每个任务最多只有一个依赖事件和一个触发事件。tGraph归一化是通过执行两种类型的转换来实现的，将每个任务的最大扇入（fan-in）和扇出（fan-out）减少到最多一个。首先，当输入tGraph中的一个任务T0触发多个事件e1, ..., ek时，MPK通过引入一个新事件e'和k个空的虚拟任务T1, ..., Tk来转换tGraph，每个虚拟任务不执行任何计算且依赖于e'。转换后，任务t0只触发e'，而每个新引入的任务Ti只触发原始事件ei中的一个，如图6a所示。这种转换确保了每个任务只有一个触发事件。图5(e)展示了MPK如何应用这种转换将A1和A2的触发事件数量减少到一个。

图5: MPK 编译器工作流。在(b)中，Q、K、V、A、O和R分别表示分解查询投影、键投影、值投影、注意力、输出投影和RMSNorm算子产生的任务集。在(e)中的D1和D2是在tGraph归一化期间插入的虚拟任务，以保证每个任务只有一个触发事件。最后，(f)展示了MPK如何线性化tGraph并将其结果结构存储在GPU内存中，其中任务和事件都遵循统一的、规范的表示。

图6: MPK 执行图转换以归一化任意 tGraph，确保每个任务最多有一个依赖事件和一个触发事件。

tGraph归一化的扇入转换。其次，当输入tGraph中的一个任务T0依赖于多个事件e1, ..., ek时，MPK通过引入一个新事件e'和k个空的虚拟任务T1, ..., Tk来转换tGraph，每个虚拟任务不执行任何计算但会触发e'。经过此转换后，任务T0和每个新引入的任务Ti都只依赖于一个事件，如图6b所示。

tGraph归一化的开销分析。当tGraph中的任务具有多个扇入和扇出事件时，tGraph归一化会因添加额外的任务和事件而引入开销。这种情况通常只在原始计算图具有可并行执行的算子时发生。例如，图5(d)中的任务A1和A2有两个扇出事件，因为原始计算图中的RMSNorm和输出投影算子都依赖于注意力算子并且可以并行运行。在实践中，我们观察到归一化开销可以忽略不计（在我们的评估中始终低于1%），因为现实世界的模型是“深”的（许多顺序算子）而不是“宽”的（许多可并行化的算子）。

Algorithm 1 MPK’s tGraph linearization algorithm. It is guaranteed that each task is enqueued into T once and that each event is enqueued into E once. Lines 5-7 ensure that all tasks depending on an event are consecutive in T .

Output: A list of tasks T such that for each event e ∈ G , the set of tasks e
launches are consecutive in T .
1: T ← ∅
2: E ← {e ∈G|e.counts =0} ▷ Enqueue all events with no dependent tasks
3: while E is not empty do
4: e← E.dequeue()
5: for all task t ∈ G do
6: if t.dependent_event=e then
7: T.enqueue(t)
8: e′ ← t.trigger_event
9: if all tasks triggering e′ are in T then
10: E.enqueue(e′)
11:
12: return T

tGraph线性化。仅靠tGraph归一化无法解决第二个挑战：在融合之后，一个事件可能仍然需要触发大量任务（例如，图5(e)中的事件e5触发了四个任务），这需要额外的存储空间来记录它们的索引。MPK使用一种基于广度优先搜索（BFS）的算法（见算法1）来解决这个问题，通过线性化tGraph。这种线性化确保了由同一事件触发的所有任务在最终的任务排序中被分配到连续的索引。因此，一个事件的扇出（fan-out）可以仅使用第一个和最后一个任务的索引来紧凑地编码，从而无需存储显式的依赖任务列表，同时保留了所有的依赖语义。

线性化tGraph的存储结构。图5(f)展示了MPK如何将线性化后的tGraph存储在GPU设备内存中。对于每个任务，MPK只记录其依赖事件和触发事件的索引。对于每个事件，MPK存储该事件被激活所需的触发次数；一旦激活，运行时将启动所有索引在该事件的第一个和最后一个任务索引范围内的任务。

4.2 任务实现生成

任务实现的自动生成。除了构建tGraph，MPK还需要为每个任务生成一个在GPU SM上执行的设备函数。MPK利用先前的工作，采用编译器超优化的方法来为每个任务自动生成高性能的实现。具体来说，MPK在线程块级别而不是内核级别执行超优化。每个计算任务都与一个参考的PyTorch实现相关联，MPK利用Mirage超优化器【33，Mirage: A {MultiLevel} superoptimizer for tensor programs, 2025, 19th USENIX Symposium on Operating Systems Design and Implementation (OSDI 25)】自动搜索一个最优的线程块图，该图被发送到Mirage编译器以生成CUDA实现。该CUDA实现包括所有SM内部的优化，如软件流水线、寄存器重用和布局优化，以最小化bank冲突。

5 核内并行运行时

核内运行时设计。MPK采用一个核内并行运行时，在单个超级内核内跨所有SM执行tGraph。这种设计消除了内核启动开销，并能对调度、同步和执行顺序进行细粒度控制。一旦启动，该超级内核会持续管理计算和通信，直到推理工作负载完成。

工作者与调度器。为了支持这种执行模型，MPK将GPU的SM划分为工作者（workers）和调度器（schedulers）。每个工作者在一个物理SM上运行，并维护一个独立的任务队列。工作者执行一个轻量级循环，重复地从队列中取出任务，执行相关的计算或通信，并通过通知任务的触发事件来表示任务完成。这种设计确保了工作者被充分利用，同时实现了跨算子的异步执行。

图7: MPK的事件驱动执行模型。圆形表示事件，蓝色（或橙色）矩形分别表示计算（或通信）任务。从事件到任务的边对应任务启动，而从任务到事件的边表示该任务触发相关事件。AT、MM和AR分别指注意力、矩阵乘法和AllReduce。

调度器组织。调度器以warp粒度组织，每个SM托管四个调度器warp。每个调度器维护一个事件队列，并重复轮询新激活的事件，将相应的任务分派给工作者。工作者和调度器的分配在内核启动时固定，并与GPU的物理SM数量匹配，避免了内核内部任何动态的角色切换开销。

章节内容。本节的其余部分详细介绍了核内运行时架构。§5.1描述了MPK的事件驱动执行模型。§5.2介绍了两种互补的任务启动机制，分析了它们的权衡，并解释了MPK如何将它们结合起来以实现低延迟和负载均衡的执行。§5.3描述了MPK为进一步减少开销和提高吞吐量而采用的其他系统优化。

5.1 事件驱动执行

事件驱动执行模型。MPK使用事件驱动模型来执行tGraph。每个tGraph都以一个指定的开始事件（例如图7中的e0）开始，该事件没有先决条件。这个事件最初被加入到调度器的事件队列中。当调度器（例如s1）从队列中取出该事件时，它会启动所有依赖于该事件的任务（例如AT1, ..., AT4）。每个启动的任务被分派给一个工作者，该工作者执行任务，并在完成后通知与该任务关联的触发事件。

事件激活与执行传播。当一个事件的所有先决条件都已完成，从而集体触发了该事件达到所需次数时，该事件就被激活。当一个事件被激活时，它会被加入到调度器的事件队列中，从而使运行时能够继续通过tGraph传播执行。通过这种方式，事件成为驱动任务执行的机制，实现了细粒度的异步执行。

5.2 混合任务启动

任务入队方式。一个任务可以通过即时（Just-In-Time, JIT）或提前（Ahead-Of-Time, AOT）的方式被加入到工作者的任务队列中。在JIT模式下，调度器只有在任务的依赖事件被完全激活后，才会将任务分配给一个工作者；因此，任务在分配后可以立即开始执行。在AOT模式下，运行时在任务的前驱事件被激活之前，就预先将任务加入到一个工作者的队列中。工作者在依赖事件被完全激活之前不能执行该任务，因此会在本地等待依赖关系得到满足。

JIT与AOT的权衡。JIT和AOT方法各有优势。一方面，JIT任务启动使MPK能够适应工作负载不平衡。例如，现代LLM中的注意力算子由于数据依赖的序列长度，其执行时间变化很大——长序列长度的请求比短序列的请求需要更长的时间来完成。这种变化使得静态分配效率低下。在JIT模式下，MPK只有在触发下游任务（如MatMul或AllReduce）的注意力任务完成后才启动它们。先完成其注意力任务的工作者可以执行更多的下游任务，从而改善端到端延迟并平衡SM间的负载，如图7所示。

JIT的延迟开销。另一方面，JIT任务启动由于额外的工作者-调度器通信而涉及更高的延迟。图8展示了在由任务T1触发的事件e之后启动任务T2的差异。在JIT模式下（图8a），执行T1的工作者必须通知一个调度器，该调度器随后取出事件e，启动T2，并将其加入到一个工作者的任务队列中。接收方工作者在执行前必须先取出T2。这个链条需要两次同步（工作者→调度器和调度器→工作者）。相比之下，在AOT模式下（图8b），T2已经被加入到一个预先分配的工作者的队列中，该工作者只需等待事件激活。因此，AOT只需要一次同步（工作者→工作者，通过事件触发），从而减少了每个任务的启动延迟。

MPK的混合启动策略。MPK使用混合任务启动来结合两种方法的优点。在tGraph构建期间，编译器根据算子的执行时间是否依赖于数据并可能导致运行时不平衡，将其分类为JIT或AOT。具有数据依赖执行时间的算子（如注意力）被标记为JIT，并且所有下游任务保持为JIT，直到遇到一个全局屏障（即一个必须由所有上游任务触发的事件）。屏障消除了不平衡，使得后续算子适合AOT。所有其余算子都被标记为AOT以最小化分派开销。这些标签在算子粒度上应用：由同一算子生成的所有任务共享相同的模式。

混合队列的执行逻辑。工作者维护两个队列——一个用于JIT任务，一个用于AOT任务。工作者总是优先处理JIT任务，因为它们可以立即执行。当一个工作者耗尽其JIT队列时，它会检查第一个AOT任务的依赖事件是否已完全激活；如果是，工作者就执行该AOT任务。这种设计确保了工作者只有在没有就绪工作时才会阻塞。

调度器与AOT任务的处理。调度器只处理与JIT相关的事件。所有AOT任务在执行开始前就已被预先加入队列。由于算子通常会生成与工作者数量成正比的任务数量（§4），MPK以轮询（round-robin）方式将AOT任务分配给各个工作者，以保持负载均衡。预先启动AOT任务减少了调度器的负载并摊销了分派开销，而JIT任务启动则动态地维护了SM间的工作负载平衡。

5.3 运行时优化

本小节介绍了一系列旨在最小化MPK运行时开销的优化措施。

分页式共享内存抽象。在CUDA和Triton等传统GPU编程模型中，共享内存是每个线程块私有的高速片上内存。共享内存仅在内核的生命周期内存在——一旦内核完成，其共享内存会自动释放。在当前的“每个算子一个内核”方法中，每个内核都假定独占访问整个共享内存以符合此编程模型。然而，这种设计阻碍了跨任务软件流水线（§2），该流水线需要将后续任务的数据加载与当前任务的计算重叠，因为两个任务都需要访问共享内存。为了实现这种流水线，MPK引入了分页式共享内存抽象。共享内存被划分为多个固定大小的页，任务实现被修改为在这些页上操作，而不是假定整体分配。任务可以根据其共享内存占用获取一个或多个页，并且在不再需要时必须释放这些页。一旦一个任务释放了它的任何页，它就不再被允许获取更多的共享内存，从而确保了简化的调度单调使用模式。当当前任务发出释放信号时，MPK可以立即为下一个任务预分配可用的页并开始数据预取。这种设计在超级内核执行模型中实现了细粒度的、按需分配的共享内存资源。

跨任务软件流水线。为了在同一工作者上执行的任务之间实现软件流水线（§2.1），MPK将每个任务分解为一个预加载阶段和一个计算阶段。预加载阶段发出数据传输指令，从设备内存中获取所需张量的一个块到共享内存，从而初始化任务内软件流水线，但不执行任何计算。当两个条件成立时，MPK会机会性地将当前任务T1的计算阶段与后续任务T2的预加载阶段重叠：（1）T1已经发出了其所有的数据传输指令，并且（2）有足够的共享内存页可用于T2的预加载阶段。需要注意的是，这样的流水线不会干扰T1的执行，因为MPK插入了适当的SM内同步屏障，以确保T2的内存传输不会与T1正在进行的数据传输冲突。

预取任务描述。每个工作者在设备内存中维护JIT和AOT两个任务队列。每个任务都关联一个任务描述，该描述编码了其输入张量、输出张量和配置参数；在我们当前的实现中，每个描述占用352字节（§6.1）。为了减少入队/出队延迟并隐藏设备内存访问成本，MPK采用了一种轻量级的预取机制，在需要之前将即将到来的任务描述取回到共享内存中。

A4 实验环境与结果

实验环境

实现细节:
* 后端集成: MPK作为PyTorch的内核后端实现。通过torch.compile(backend=MPK)，一个PyTorch程序可以被编译成MPK超级内核，返回一个可调用的PyTorch函数。
* 代码库: MPK当前实现包含约4万行C++，8.4万行CUDA和1万行Python代码。
* 核内运行时: 使用CUDA编写，通过设备内存中的信号量协调工作者和调度器。
* 编译器: MPK编译器用C++和Python实现，集成了Mirage超优化器【33，Mirage: A {MultiLevel} superoptimizer for tensor programs, 2025, 19th USENIX Symposium on Operating Systems Design and Implementation (OSDI 25)】来生成优化的CUDA任务实现，并使用NVSHMEM【10，NVIDIA OpenSHMEM Library (NVSHMEM) Documentation, 2025, https://docs.nvidia.com/nvshmem/api/index.html】支持核内GPU间通信。

动态工作负载支持:
* 为了支持LLM服务中的动态性，MPK集成了连续批处理（continuous batching）【36，Orca: A distributed serving system for {Transformer-Based} generative models, 2022, 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22)】和分页注意力（paged attention）【23，Efficient memory management for large language models, 2023, Proceedings of the ACM Symposium on Operating Systems Principles (SOSP)】。请求调度和KV缓存元数据更新逻辑都在超级内核内部的一个任务中执行。
* 为处理动态批量大小，MPK为代表性的批量大小（2的幂次，直至最大批次）生成多个专用的tGraph。运行时，调度器根据当前批次大小选择合适的图。

硬件与软件配置:
* GPU: 实验在三代NVIDIA GPU上进行：A100, H100, 和 B200。
* SM分配: 在所有实验中，保留4个SM用于调度器（共16个调度器warp），其余SM分配给工作者。
* 共享内存: 共享内存页大小固定为32KB，这使得A100、H100和B200上每个SM分别有5、7、7个共享内存页。
* 模型: 评估了五种广泛部署的LLM，涵盖了不同的参数规模和架构家族，包括密集模型和混合专家（MoE）模型。
* 实验设置: 采用离线、批处理推理设置，请求的prompt长度固定为64，解码长度为1024个token，以消除请求到达模式变化带来的影响，从而纯粹比较系统性能。所有模型使用bfloat16精度。

Table 1: 评估中 MPK 的配置。
表1: 评估中 MPK 的配置。

实验结果

端到端性能比较

实验内容: 将MPK与两个先进的LLM服务系统SGLang和vLLM进行端到端吞吐量比较。这两个基线系统都采用kernel-per-operator方法，并依赖于包括FlashInfer、FlashAttention、cuBLAS、cuTLASS和Triton在内的多种专用内核库。
实验结果:
* 在单批次推理中，MPK在所有模型和硬件上的性能提升了1.0倍至1.7倍。对于较小的模型和较新的GPU，性能提升更为显著。
* 原因分析：(1) kernel-per-operator方法的内核启动开销更高；(2) 这些系统存在流水线气泡，因为内核抽象阻止了跨任务流水线；(3) SGLang和vLLM在CPU上进行页面分配和请求调度，增加了CPU-GPU同步延迟。这些开销在模型变小、硬件变强时占比更大。
* MPK非常适合低延迟场景，例如在A100 GPU上运行Qwen3-8B时，MPK将每token解码延迟从14.5ms降低到12.5ms，接近约10ms的理论下限。
* 与vLLM和SGLang相比，MPK的易用性更高，它通过编译器自动将PyTorch模型超级内核化，而无需手动优化。

图9: 在 A100、H100 和 B200 GPU 上比较 MPK 与现有系统在五个模型上的性能。所有性能数据均相对于 MPK 进行归一化（越高越好）。每个 MPK 条上方的数值表示其相对于性能最佳的现有系统的加速比。由于模型超出了单个 A100 的内存容量，因此省略了 Qwen3-30B-A3B 在 A100 上的结果。

混合专家（MoE）模型案例研究

实验内容: 评估MPK针对MoE模型的两项特定优化：(1) 混合工作负载均衡器和 (2) 融合的gather-GEMM实现。
* 混合工作负载均衡器: 静态方法在token路由倾斜时会导致负载不平衡，而全动态方法同步开销大。MPK的混合策略在编译时静态划分任务，在运行时任务根据全局MoE元数据动态调整工作负载分配，避免了细粒度同步开销。
* 融合的gather-GEMM: MPK将基于TMA的gather操作替换为异步的token级拷贝，并将其直接集成到GEMM任务的数据加载阶段，消除了独立的gather内核和额外的调度开销。
实验结果:
* 如图10所示，在Qwen3-30B-A3B模型上，混合均衡策略在所有批量大小下都优于纯静态分区。
* 与SGLang相比，MPK的混合MoE方法取得了1.07倍到1.18倍的加速。融合的gather-GEMM也带来了稳定的性能提升。

图10: 在 B200 上比较 MPK 与现有系统在 Qwen3-30B-A3B 上的性能。每个值代表各方法的实际 MoE 运行时（微秒，越低越好），条上方的数字表示 MPK-Hybrid-MoE 相对于 SGLang-MoE 实现的加速比。

多GPU性能结果

实验内容: 在NVIDIA H100 DGX实例上评估MPK在多GPU环境下的可扩展性，采用张量模型并行。MPK将AllReduce等集合通信算子自动编译为异步的GPU间数据传输任务和本地归约任务。
实验结果:
* 与结合了手动优化内核、CUDA Graphs和torch.compile的PyTorch相比，MPK的超级内核执行将吞吐量提高了多达10倍。
* 与SGLang和vLLM等高度优化的服务系统相比，当扩展到8个H100 GPU时，MPK实现了1.1倍至1.4倍的加速。
* 性能增益来源于三个关键优化：(1) 将页面分配和请求调度集成到超级内核内，消除了CPU端的调度开销；(2) 异步执行模型重叠了计算和通信任务；(3) 消除内核屏障，实现了跨任务软件流水线。

图11: 在多 H100 GPU 上，使用张量并行对 Qwen3-1.7B 推理进行 MPK 与现有系统的比较。所有系统的性能均通过 MPK 进行归一化（越高越好）。

消融研究

跨任务流水线:
* 实验内容: 在Qwen3-8B的最后一个线性层上评估跨任务流水线（在当前任务计算的同时预加载下一个任务的数据）的影响。
* 实验结果: 如图12所示，跨任务流水线将任务运行时间减少了1.2至1.3倍，其性能甚至超过了cuBLAS编译的内核。

计算-通信重叠:
* 实验内容: 通过禁用细粒度依赖（即仅在集合算子和先前任务之间使用单个事件来捕获粗粒度依赖），来评估计算与通信任务重叠的影响。实验在4个H100 GPU上对Qwen3-1.7B进行张量并行推理。
* 实验结果: 如图13所示，启用计算-通信重叠将每次迭代的延迟降低了1.1倍。

图13: 计算-通信重叠的消融研究。我们测量了在4个H100 GPU上使用张量并行的Qwen3-1.7B的每次迭代运行时间（越低越好），比较了有无计算-通信重叠的性能。

A7 补充细节

A5 结论

本文介绍了MPK，这是首个能将多GPU模型推理自动转换为完全融合的超级内核（mega-kernel）的编译器和运行时系统。通过引入SM级任务图和核内并行运行时，MPK克服了“每个算子一个内核”模型的根本限制，实现了算子间的软件流水线、计算与通信的细粒度重叠，并消除了内核启动和CPU侧的调度开销。我们的评估表明，MPK将LLM服务的延迟推向了硬件极限，并在各种模型和GPU代际上显著提高了吞吐量。通过将执行统一在单个超级内核内，同时保留现有ML框架的易用性，MPK为构建高性能、编译器驱动的推理系统开辟了一条新路径。

PaperCache

Mirage Persistent Kernel: A Compiler and Runtime for Mega-Kernelizing Tensor Programs

Mirage Persistent Kernel: A Compiler and Runtime for Mega-Kernelizing Tensor Programs

A1 主要贡献