Mirage Persistent Kernel：突破LLM推理极限，自动巨核化技术实现1.7倍性能飞跃

关键词：#MPK、#LLM推理、#MegaKernel、#SM级任务图、#多GPU优化、#跨算子优化

MPK 作为首个自动 Mega Kernel 化多 GPU LLM 推理的编译器-运行时系统 ，以 SM 级 tGraph 打破核间壁垒，让跨算子 软件流水线与细粒度计算-通信重叠从理论走向实用；无需修改模型代码，仅需数行 PyTorch 集成，它即可在 A100/H100/B200 上相比 SGLang、vLLM 实现 1.0-1.7 倍端到端推理提速 ，推动 LLM 推理性能逼近硬件物理极限；而其模型无关的设计使其不仅适用于 Transformer 架构，也为其他类型的深度神经网络提供了通用的巨核化编译与执行方案 。

Mirage Persistent Kernel: A Compiler and Runtime for Mega-Kernelizing Tensor Programs
https://arxiv.org/pdf/2512.22219
https://github.com/mirage-project/mirage

本文提出 Mirage Persistent Kernel（MPK），首个将多 GPU 模型推理自动转化为高性能巨核的编译器与运行时系统，旨在突破传统“算子核”执行模式的固有局限。传统模式依赖核间屏障同步，导致无法实现跨算子流水线、细粒度计算-通信重叠，且核启动开销显著。MPK 的核心创新是 SM 级 tGraph 表示，以流多处理器（SM）为粒度将张量程序分解为任务，通过细粒度事件捕捉数据依赖 ，解锁此前不可行的 GPU 优化。

MPK 包含两大核心组件：

编译器 将张量程序转化为优化 tGraph，通过算子分解、依赖分析、事件融合及图归一化/线性化，生成高效 CUDA 实现；
内核内并行运行时 采用分布式调度，将 SM 划分为工作器与调度器，通过事件驱动异步执行、JIT+AOT 混合任务启动，最大化 GPU 利用率。

图 1：MPK 概述。该图展示了 Mirage Persistent Kernel（MPK）的整体架构，清晰呈现了其核心组成部分 —— 编译器与核内并行运行时的协同工作流程，直观体现了 MPK 将张量程序转化为单一大核（mega-kernel）的关键过程。MPK 编译器负责将输入的张量程序和推理配置转化为 SM 级的优化任务图（tGraph），并生成高效 CUDA 实现；核内并行运行时则在单个大核中调度执行这些任务。这种架构突破了传统 “一算子一核” 模式的局限，为后续跨算子软件流水线、细粒度计算通信重叠等优化奠定了基础，是实现端到端核融合的关键框架支撑。

此外，MPK 还引入分页共享内存、跨任务软件流水线等优化，进一步降低开销。

在 A100、H100、B200 三代 NVIDIA GPU 上的评估显示，MPK 在 5 个主流 LLM（含稠密型与 MoE 模型）上，较 SGLang、vLLM 等系统单 GPU 推理提速 1.0-1.7 倍，多 GPU 场景提速 1.1-1.4 倍 ，显著降低端到端延迟。其仅需数行代码即可集成到 PyTorch，兼顾高性能与易用性，且模型无关，可支持任意 DNN 架构，为编译器驱动的高性能推理系统提供新范式。

图 9：该图对比了 MPK 与现有系统在 A100、H100、B200 三种 GPU 上对五种模型的性能，所有性能数据均以 MPK 为基准进行归一化（数值越高越好）；MPK 柱状图上方的数值表示其相对于性能最佳现有系统的加速比；因模型超出单 A100 内存容量，故省略 Qwen3-30B-A3B 在 A100 上的结果。对比对象为 SGLang 与 vLLM（当前 LLM 推理领域的先进系统，采用 “一算子一核” 模式，集成 FlashInfer、cuBLAS 等优化库）。实验覆盖 dense 与 MoE 模型，跨三代 GPU。结果显示，MPK 单批次推理性能提升 1.0-1.7 倍，小模型与新 GPU 上提升更显著 —— 因 “一算子一核” 模式的核启动开销、流水线气泡、CPU 调度延迟在这些场景更突出。例如，Qwen3-8B 在 A100 上，MPK 将解码延迟从 14.5ms（vLLM/SGLang）降至 12.5ms，接近硬件理论下限（约 10ms），证明 MPK 能突破传统模式局限，推动 LLM 推理接近硬件极限。

本文目录

关键问题
- 问题一：MPK 的“自动化”本质及其边界
- 问题二：内核内调度的开销与可扩展性风险
一、引言
- 1.1 MPK 编译器
- 1.2 核内并行运行时
- 1.3 评估结果
二、背景知识
- 2.1 GPU 编程模型
- 2.2 核函数融合
三、SM 级图表示
- 3.1 与 CUDA 图的对比
四、MPK 编译器
- 4.1 tGraph 生成
- 4.2 任务实现生成
五、核内并行运行时（In-Kernel Parallel Runtime）
- 5.1 事件驱动执行（Event-Driven Execution）
- 5.2 混合任务启动（Hybrid Task Launch）
- 5.3 运行时优化
六、评估
- 6.1 实现细节
- 6.2 实验设置
- 6.3 端到端结果
- 6.4 案例研究：混合专家模型
- 6.5 多 GPU 结果
- 6.6 消融实验
七、相关工作
- 7.1 手动设计内核（Manually Designed Kernels）
- 7.2 机器学习编译器（ML Compilers）
- 7.3 巨型内核（Mega-Kernels）
八、结论
参考文献

关键问题

问题一：MPK 的“自动化”本质及其边界

问题一：MPK“全自动”编译的依赖与泛化能力

MPK声称能够“全自动”地将多GPU模型推理转换为单一高性能巨核，但其核心编译优化（如超级优化、任务划分、硬件适配）是否严重依赖于对现有NVIDIA GPU架构的先验知识、手工调优规则以及静态建模？ 这种“自动化”在面对非NVIDIA硬件、非规则计算模式或极端动态负载时，其泛化能力和性能是否将显著受限？

MPK的“全自动”编译流程确实建立在若干关键的先验假设和领域特定优化之上，其能力边界可概括为以下几点：

超级优化依赖于历史规则与搜索空间：MPK为每个计算任务生成高性能CUDA实现时，依赖于其集成的Mirage超级优化器。该优化器通过搜索“线程块图”空间来寻找最优实现。这种搜索并非无中生有，而是基于对GPU执行模型（如内存层次、线程束调度）和已知高效模式（如软件流水线、共享内存复用）的深刻理解所构建的搜索空间。因此，其“自动生成”本质上是在结构化、预定义的优化空间内进行启发式搜索，其效果高度依赖于该空间是否涵盖了目标硬件的最优模式。
硬件适配是参数化而非根本性重构：MPK的编译器会针对目标GPU架构（如A100、H100、B200）生成专门的SM级任务图（tGraph）和配置参数。例如，共享内存页大小固定为32KB，调度器和工作线程的数量根据SM总数静态分配。

表1：评估中的MPK配置。该表展示了MPK在不同GPU上的关键配置参数，包括GPU型号、流多处理器（SM）数量、工作器（Worker）数量、调度器（Scheduler）数量。MPK的核内运行时将GPU的SM划分为工作器（执行任务）与调度器（管理事件与任务分配）。实验中，所有GPU均预留4个SM（16个调度器warp）作为调度器，剩余SM作为工作器。这种配置旨在平衡调度与计算资源，同时共享内存页大小固定为32KB，为页面化共享内存抽象提供基础。

这种适配是基于对不同NVIDIA GPU架构参数的静态映射，而非一种能从硬件ISA或微架构描述中自动推导出最优策略的通用机制。因此，其“硬件适配”主要针对连续几代NVIDIA GPU进行了调优，对于架构差异巨大的其他厂商GPU（如AMD CDNA或Intel GPGPU），当前的MPK设计很可能无法直接有效工作。

动态性支持通过预编译多版本实现：为支持LLM服务中的动态批次大小，MPK并非在单一内核内实现完全动态的资源配置，而是为多个代表性的批次大小（2的幂次）预先编译多个tGraph，运行时再根据当前批次进行选择。这是一种“离散化”的近似策略，而非处理任意形状的完全动态编译。对于训练或具有不可预测稀疏模式的工作负载，此方法可能不够灵活。

结论：MPK的“自动巨核化”是一项重大进步，但它并非一个无视硬件和模型特性的通用解决方案。它是一个高度工程化、针对NVIDIA GPU体系结构和稠密、可规则划分的张量计算模式进行了深度特化的编译器。其自动化在已知领域内效果显著，但将其直接应用于全新硬件架构或极度不规则的计算范式时，很可能需要扩展其优化规则库、重建设计空间，甚至修改其执行模型，否则性能将难以保证。

问题二：内核内调度的开销与可扩展性风险

MPK通过将任务调度器嵌入GPU内核内部，以消除核启动和CPU调度开销，从而获得性能提升。这是否仅仅是将系统的复杂度与瓶颈从主机侧转移到了设备侧？随着模型规模、SM数量或任务并发度的增长，这种设备内分布式调度机制（如事件轮询、队列竞争）是否会成为新的性能瓶颈，从而抵消“巨核融合”带来的收益？

MPK将调度复杂性移入GPU内核，这是一把双刃剑。论文数据显示了在当前实验规模下的显著收益，但该设计确实引入了新的潜在瓶颈和可扩展性疑问。

调度开销被精心最小化，但并未消失：MPK采用了多种优化来控制内核内调度开销：
- 混合任务启动（JIT+AOT）：对数据依赖性强、易导致负载不均衡的算子（如Attention）使用JIT调度以实现动态平衡；对执行时间稳定的算子使用AOT预分发以减少调度延迟。
- 轻量级同步原语：使用设备内存中的信号量（semaphore）和原子操作进行协调，避免昂贵的全局同步。
- 去中心化调度：多个调度器warp独立工作，无需全局协调，减少了争用。
- 任务描述预取：将任务描述从设备内存预取到共享内存，降低访问延迟。
  这些措施使得在评估所及的模型规模（最大30B参数）和并发度（批次≤16）下，调度开销相对于计算和通信重叠带来的收益而言是微小的。
可扩展性存在潜在风险，作者未充分压力测试：尽管有上述优化，随着系统规模扩大，以下风险可能浮现：
- 调度器争用：每个调度器warp负责一个事件队列。当大量任务在极短时间内完成并触发事件时（例如，一个广播事件被许多SM同时触发），对应的事件队列可能成为热点。论文未展示在极端并发条件（如数百个SM同时尝试激活事件）下的性能表现。
- 工作线程饿死：如果调度器的分发速度跟不上工作线程的消费速度，或者任务粒度极小导致调度频率很高，工作线程可能闲置。MPK的AOT预分配缓解了此问题，但对于完全JIT的动态任务链，风险仍存。
- 资源内部碎片化：Paged Shared Memory的固定页大小设计，在任务内存需求多样且动态的情况下可能导致碎片。论文未探讨极端场景下的内存分配效率。
- 内核复杂性爆炸：一个融合了所有算子的巨核，其指令缓存占用、寄存器压力等都可能随着模型复杂度增加而增长，可能影响SM的占用率和指令发射效率。

结论：MPK通过精妙的协同设计（编译器生成优化任务图 + 运行时轻量级调度）在当前常见的LLM服务场景下，成功地将调度开销控制在了极低水平，并实现了性能提升。然而，论文并未提供在超大规模模型（如万亿参数）、极致批量（如数百）或极端动态任务图下的可扩展性数据。因此，“调度开销内化是否会成为新瓶颈”这一问题，对于超出论文实验范围的场景，仍然是一个开放且合理的担忧。MPK的设计是迈向极致性能的重要一步，但其大规模极限仍需在未来工作中验证。

现有系统通常为每个算子分配一个专用的GPU核函数（kernel，即在GPU上执行的一段并行计算代码）。这些核函数可能由专家手工优化[17,35]，也可能由机器学习编译器自动生成[15,29,33]。然而，这种“每个算子对应一个核函数”的执行模型，严重限制了多项关键的跨算子GPU优化。

阻碍跨算子软件流水线：现代GPU会在同一流（stream，GPU中用于管理任务执行顺序的逻辑队列）上连续启动的核函数之间，隐式插入核函数屏障（kernel barrier，一种强制前序核函数所有线程执行完毕后，后续核函数线程才开始执行的同步机制）。这种机制虽然保证了数据依赖关系，却阻碍了跨算子软件流水线（software pipelining，一种通过重叠不同任务的执行阶段来提高硬件利用率的技术）的实现，迫使存在依赖关系的算子必须严格按顺序执行。英伟达（NVIDIA）近期推出的可编程依赖启动（PDL，Programmatic Dependent Launch）技术[11]允许同一流上的核函数部分重叠执行，但采用PDL需要大量工程投入，因为它从根本上改变了核函数的结构和控制流。
无法实现细粒度计算-通信重叠：“每个算子对应一个核函数”的方式无法实现细粒度的计算-通信重叠（compute-communication overlap，指让计算操作和数据传输操作同时进行，以隐藏通信延迟的优化手段）。由于依赖关系仅在粗粒度的算子层面捕获，运行时必须等待前一个算子完全执行完毕后，才能启动依赖于它的通信或计算操作。例如，当矩阵乘法（MatMul）算子后紧跟一个全归约（All-Reduce，一种用于多设备间数据同步的通信算子）算子时，全归约算子必须等待整个矩阵乘法完成才能启动——即便全归约的每个数据片段仅依赖于矩阵乘法输出的一个子集。要利用这类优化机会，需要以比单个核函数更细的粒度来表示和控制数据依赖关系。
依赖静态的CUDA图应对启动开销：该方式意味着每次推理迭代需要启动数百到数千个核函数。为降低启动开销，当前系统严重依赖CUDA图（CUDA Graphs，一种用于捕获GPU操作序列并以极低开销批量启动的技术）。但CUDA图在很大程度上是静态的：控制流、张量形状或数据依赖关系的任何变化，都需要重新实例化或修改已捕获的图，这对于模型推理中常见的动态工作负载（如可变的输入长度、批量大小）缺乏灵活性。

解决这些局限性的一个可行方案是将所有计算和通信融合为一个“巨型核函数”（mega-kernel，也称为持久核函数 persistent kernel）。在这种设计中，系统仅启动一个GPU核函数即可执行整个模型的计算——从层计算到GPU间通信——全程无中断。

巨型核函数在多个方面克服了“每个算子对应一个核函数”方式的局限性：
* 首先，它通过避免重复启动核函数消除了核函数启动开销；
* 其次，将所有算子计算融合到单个核函数中，支持跨算子软件流水线，可在当前算子计算仍在进行时预取下一个算子所需的数据；
* 最后，它支持计算与GPU间通信的细粒度重叠，通过并行执行更有效地隐藏通信延迟。

尽管具备这些优势，将机器学习模型自动转换为高性能巨型核函数仍然面临挑战。现有机器学习系统（如PyTorch[25]、Triton[29]、TVM[15]）均不支持端到端的巨型核函数生成。此外，这些系统依赖碎片化的专用库生态：用于通信的NCCL[4]或NVSHMEM[10]、用于注意力（Attention）计算的FlashInfer[35]或FlashAttention[17]，以及用于自定义计算的CUDA或Triton。这种碎片化使得难以将整个推理流水线统一到单个核函数中。

本文提出了幻影持久核函数（Mirage Persistent Kernel，MPK）——首个能将多GPU模型推理自动转换为高性能巨型核函数的编译器与运行时系统。MPK支持端到端核函数融合，且只需极少开发者投入：用户仅需几行代码即可将PyTorch模型“巨型核函数化”（mega-kernelize），相比在原生PyTorch中结合CUDA图和torch.compile（PyTorch的编译优化接口）运行模型，性能显著提升。MPK兼具巨型核函数的性能优势与现有机器学习框架的易用性。

MPK的核心思想是：以单个流式多处理器（SM）为粒度表示计算和GPU间通信，而非以整个GPU为粒度。MPK引入了一种SM级别的图表示（称为tGraph）：
* 图中的节点表示在单个SM上运行的任务。
* 边则编码任务间的细粒度依赖关系。

这种表示方式暴露了更多并行性，支持跨算子软件流水线、细粒度核函数重叠等现有“每个算子对应一个核函数”方式无法实现的优化。MPK通过图1所示的两个核心组件实现这一思想。

1.1 MPK 编译器

MPK编译器以张量程序和推理配置为输入，自动将张量程序的计算图转换为针对特定推理配置和GPU架构优化的SM级tGraph。编译器引入了多种优化手段（包括事件融合、图规范化、图线性化），以减少同步开销并最大化生成tGraph的性能。此外，MPK还利用现有超优化技术[33]为每个任务自动生成高性能CUDA实现，确保SM级执行的高效性。
* [33]是 OSDI 25 这篇：Mirage: A {Multi-Level} superoptimizer for tensor programs

1.2 核内并行运行时

MPK通过嵌入在巨型核函数内部的“核内并行运行时”（in-kernel parallel runtime）执行SM级tGraph，无需在模型执行过程中启动任何额外核函数，即可对任务执行和调度进行细粒度控制。为实现这一目标，运行时将GPU的SM划分为“工作器”（worker）和“调度器”（scheduler）两类角色：
* 工作器：维护专用任务队列，以先进先出（FIFO）顺序执行所有分配给它的任务；
* 调度器：维护任务间的依赖关系，在任务的前置条件满足时将其分配给工作器。

MPK运行时采用事件驱动的全异步执行模型，确保GPU资源得到最大程度利用。最后，MPK运行时还采用结合即时（JIT，Just-In-Time）调度和提前（AOT，Ahead-Of-Time）调度的混合任务启动策略，在最小化运行时开销的同时，保证SM间的动态负载均衡。

1.3 评估结果

我们将 MPK 实现为 PyTorch 的核函数后端：只需修改几行代码，即可将 PyTorch 程序编译为 MPK 巨型核函数。我们在 5 个广泛使用的模型、3 代 NVIDIA GPU（A100、H100、B200）上对 MPK 进行了评估。即使对于已被现有“每个算子对应一个核函数”系统（如用于 LLM 服务的 SGLang 和 vLLM）广泛部署和深度优化的工作负载，MPK 在单 GPU 和多 GPU 部署场景下仍能超越现有系统 1.0-1.7 倍 ，将 LLM 推理性能推向硬件极限。

二、背景知识

本节首先回顾面向核函数的 GPU 编程模型并讨论其局限性（第 2.1 节），然后介绍核函数融合与巨型核函数技术（第 2.2 节）——这些内容为 MPK 的设计提供了动机。

2.1 GPU 编程模型

在 GPU 中，计算以核函数为单位组织，每个核函数表示一段以单程序多数据（SPMD，Single-Program, Multiple-Data，即所有线程执行相同代码但处理不同数据）方式并行执行的函数。一个核函数由线程块（thread block）网格（grid）组成：每个线程块被分配给一个流式多处理器（SM），且包含多个线程，这些线程可对单个数据元素进行操作。

每个线程拥有私有寄存器文件（register file，GPU 中速度最快的存储单元，供单个线程独占使用）；
同一线程块内的线程可通过低延迟共享内存（shared memory，SM 内的高速共享存储区域，供线程块内所有线程访问）进行数据交换和集合操作。
所有核函数的输入和输出均存储在 GPU 设备内存（device memory，GPU 的全局内存，容量较大但访问延迟较高）中。

CUDA 编程模型不支持核函数内部【线程块】之间的直接同步—— 现代 GPU 架构缺乏此类协调所需的硬件机制 。因此，跨算子依赖关系（例如“矩阵乘法必须在另一个算子之前完成”）需通过核函数屏障来强制实现，GPU 运行时会在同一流上连续启动的核函数之间自动插入这种屏障。

实践中，全收集（AllGather，一种通信算子，需将所有设备上的数据汇总到每个设备）与矩阵乘法（MatMul）之间的数据依赖关系粒度要细得多： 由于全收集执行的是按元素操作，其每个线程块仅依赖于矩阵乘法的一个线程块的输出 。这种依赖结构支持细粒度核函数重叠——只要保留细粒度数据依赖关系，不同 SM 可并行执行矩阵乘法和全收集任务。如先前研究[41]所指出的，这种重叠能让系统同时利用现代 GPU 上的计算带宽和通信带宽。

然而，要实现这种重叠，需要在 SM 之间进行核函数级以下粒度的同步（如下图 3b 所示），而传统的核函数屏障无法支持这一点。

图 3：MPK 与现有方法在任务间细粒度核重叠支持上的对比。该图对比了 MPK 与现有方法在支持任务间细粒度核重叠方面的差异，黑色箭头（图 3b 中）代表数据依赖关系，确保执行正确性，直观体现 MPK 在并行利用计算与通信带宽上的优势。在 LLM 推理中，MatMul（矩阵乘法）后接 AllGather（全收集）是常见计算模式。现有方法将二者作为独立核执行，需等待所有 MatMul 线程块完成才能启动 AllGather，浪费资源（图 3a）。而 MPK 基于 SM 级任务依赖（图 3b），让不同 SM 可并行执行 MatMul 与 AllGather，只要满足细粒度数据依赖。这种优化契合现代 GPU 计算与通信资源分离的特点，能同时利用两类带宽，在多 GPU 张量并行场景中尤为关键，是 MPK 提升多 GPU 推理性能的核心技术之一

尽管核函数屏障简化了依赖关系管理，但它也阻碍了关键的 GPU 优化，例如跨核函数软件流水线和细粒度核函数重叠。

软件流水线

GPU 架构的异构性日益增强，集成了张量核心（tensor core，用于加速矩阵乘法等张量运算的专用硬件单元）和张量内存加速器（TMA，Tensor Memory Accelerator，用于高效执行张量数据传输的专用硬件）等专用加速器。由于 TMA 的加载和存储指令是异步执行的，数据移动可与张量核心、CUDA 核心（GPU 的通用计算核心）的计算操作并行进行 。要充分利用这些加速器，需要采用软件流水线技术—— 通过在任务的多个迭代中交错执行计算和数据移动的独立阶段，最大化硬件利用率。

图 2：MPK 与现有方法在任务内（intra-task）和跨任务（cross-task）流水线支持上的差异，清晰展现出 MPK 在消除流水线气泡、提升硬件利用率上的优势。现有系统受限于核屏障，仅能实现任务内流水线（如图 2a），不同任务间存在流水线气泡，导致 GPU 异构资源（如张量核、TMA）无法充分利用。而 MPK 通过 SM 级任务调度，打破任务边界，实现跨任务流水线，让数据预加载与计算并行。这种优化对 GPU 架构日益异构的趋势至关重要，能充分发挥 TMA 异步数据传输与张量核计算的协同作用，显著提升 LLM 推理等场景的硬件利用率，这也是 MPK 降低延迟的重要技术支撑

现有系统实现的是“任务内流水线” （intra-task pipelining），如图 2a 所示：将单个任务分解为多个迭代，在该模型中，TMA、张量核心和 CUDA 核心可分别为流水线中不同迭代执行数据传输、矩阵计算和辅助操作。但核函数屏障将流水线限制在单个任务内部，无法实现“任务间流水线”（inter-task pipelining），并会引入“流水线气泡” （pipeline bubble，指流水线中因等待依赖而出现的空闲周期），导致硬件资源未被充分利用。

细粒度核函数重叠

核函数屏障还会错失利用不同硬件资源（如计算与通信）的核函数之间的重叠机会，因为它仅在整个核函数粒度上强制依赖关系，而非单个数据单元粒度。

图 3a 展示了大型语言模型（LLM）中的一个常见模式：矩阵乘法（MatMul）算子之后紧跟全收集（AllGather）算子。现有系统通常将这两个算子作为独立核函数启动，这意味着全收集核函数的所有线程块必须等待前一个矩阵乘法 核函数的所有线程块执行完毕才能启动。

2.2 核函数融合

核函数融合（kernel fusion）通过将在同一数据上按顺序执行的多个 GPU 核函数合并为一个语义等价的核函数，消除核函数屏障。这种技术可通过避免实例化中间结果、减少设备内存访问、消除核函数启动开销来提升性能。

核函数融合已在张量程序编译器中得到广泛应用。

PyTorch、JAX、TVM 等框架采用基于规则的启发式方法融合相邻核函数[14,15,25]；
而 Mirage、TASO 等系统则通过编译器超优化自动发现融合规则[22,33]。

然而，现有编译器仅能融合小规模的局部算子组——生成一个能准确实现复杂张量程序的单个核函数，在计算上难度极大，且往往不可行。

巨型核函数范式将核函数融合推向极致：通过设备内存同步原语（用于实现多 SM 间同步的底层指令）协调 SM 间的执行，将张量程序的所有计算和通信融合为一个持久核函数。尽管具备性能优势，PyTorch、Triton、JAX、TVM 等现有机器学习编译器均不支持巨型核函数编译。

目前的巨型核函数均由 GPU 专家针对特定模型手工设计：例如，
* FlashDMoE 将混合专家（MoE，一种通过多个“专家”子网络并行处理不同数据的模型架构）计算与 GPU 间通信融合为单个核函数；
* 而 Spector 等人则为 LLAMA-1B 模型手工设计并实现了低延迟巨型核函数。

这些手工方法需要大量工程投入和深厚的 GPU 专业知识才能将张量程序“巨型核函数化”。相比之下，MPK 采用基于编译器的方法，能自动将张量程序转换为优化的巨型核函数，无需人工干预。

三、SM 级图表示

本节介绍 tGraph——一种以单个流式多处理器（SM）为粒度表示张量程序计算的图结构。这种细粒度表示暴露了更多并行性，支持跨算子软件流水线、细粒度核函数重叠等现有“每个算子对应一个核函数”执行模型无法实现的优化。

图 4 展示了 tGraph 的一个示例，图中每个节点要么表示一个“任务”（task），要么表示一个“事件”（event）：
* 任务：以蓝色（或橙色）矩形表示，指在单个 SM 上执行的计算或通信单元；
* 事件：以绿色圆形表示，指任务间的同步点。

在 tGraph 中，任务和事件交替排列：每个任务的出边仅指向“触发事件”（任务完成后会激活的事件），入边仅来自“依赖事件”（任务启动前必须已激活的事件）。
* 当一个任务的所有依赖事件均已激活时，该任务即可启动；任务执行完成后，会通知其触发事件。
* 当一个事件收到所有关联任务的通知后，该事件即被激活。

这种结构以远细于传统计算图的粒度捕获依赖关系。例如，在多 GPU LLM 服务中，常出现“矩阵乘法（MatMul）算子后接全归约（AllReduce）算子”的场景（图 4a）。现有系统由于依赖粗粒度的核函数屏障，通常只能让这两个算子按顺序执行。相比之下，SM 级任务图可表示精确的任务级依赖关系：由于全归约执行按元素通信，其每个任务仅依赖于矩阵乘法的一个对应任务。通过在存在依赖关系的任务对之间插入细粒度事件，MPK 可让计算密集型的矩阵乘法任务与通信密集型的全归约任务重叠执行，最大化 GPU 整体利用率。

同一计算图可对应多个 tGraph。图 4c 展示了一个备选但性能较差的 tGraph——其中事件仅捕获算子级依赖关系，与传统核函数屏障类似。第 4 节将介绍 MPK 如何通过推断精确的数据依赖关系来生成高性能任务图，以最大化并行性并最小化同步开销。

3.1 与 CUDA 图的对比

tGraph 可视为 CUDA 图的底层扩展，两者在结构上有若干相似之处：与 CUDA 图一样，tGraph 也是静态实例化的，并对操作间的显式依赖关系进行编码。然而，CUDA 图仅在核函数粒度捕获依赖关系，而 tGraph 则以单个 SM 任务和核函数级以下事件为粒度；CUDA 图主要描述核函数启动顺序，并依赖流语义进行同步，这将重叠和融合限制在核函数边界内。相比之下，tGraph 显式建模算子内部和算子间的依赖关系，支持 SM 间的细粒度同步以及单个核函数内计算与通信的重叠。这种设计使 MPK 能够利用 CUDA 图或核函数级执行模型无法访问的并行性。

四、MPK 编译器

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

4.1 tGraph 生成

算子分解

MPK 编译器通过对算子的输出张量进行分区，将输入计算图的每个算子分解为一组任务——确保所有任务计算输出的不相交子集，从而可在 SM 间并行执行。

大多数张量代数算子可跨多个输出维度分区：例如，矩阵乘法的输出张量可沿行和列两个维度切分（tiling），以暴露并行化机会。

分区策略的性能取决于问题形状（如张量维度、数据类型）和目标 GPU 架构。为找到高效的分区策略，MPK 会选择能最小化从设备内存到共享内存数据加载量的策略——因为访问设备内存的开销远高于访问共享内存或在 CUDA 核心/张量核心上执行计算的开销。默认情况下，MPK 生成的任务数量与 SM 数量成比例，以促进执行过程中 SM 间的负载均衡。此外，MPK 还提供接口，允许用户通过设置每个输出维度的期望并行度，轻松指定自定义分区策略。

依赖关系分析

MPK 利用“事件”捕获任务间的依赖关系。对于任何两个共享张量的算子，MPK 会枚举这两个算子的所有任务对，并仅当任务（来自前一个算子）生成的输出区域与任务（来自后一个算子）消耗的输入区域重叠时，才为该任务对引入一个事件。该事件作为同步点，表示必须在生成所需数据后才能启动。相应地，MPK 会在生成的 tGraph 中插入两条边：（完成后激活）和（需等待激活后启动）。这种细粒度依赖分析在保留所有生产者-消费者依赖关系的同时，最大化了独立任务间的并行性。

事件融合

MPK 采用两种互补的事件融合方式——“后继集融合”（successor-set fusion）和“前驱集融合”（predecessor-set fusion）——以消除冗余同步点并简化构建的 tGraph。对于事件，我们定义两个函数：

：触发的事件集合（即执行完成后会通知的任务）；
：依赖的事件集合（即需等待激活后才能启动的任务）。

通过这两个函数，可判断多个事件是否具有相同的依赖结构，进而将其融合。

1. 后继集融合

后继集融合合并那些作为同一组消费者任务（即相同）前置条件的事件。由于这些消费者任务必须等待所有此类事件激活后才能启动，将它们单独表示无法获得额外的调度灵活性。

定义 4.1（后继集融合）：对于 tGraph 中的任意两个事件和，当且仅当时，可应用后继集融合。MPK 从 tGraph 中移除和，并引入一个融合事件，其中（的触发任务为和的触发任务的并集），（的依赖任务与相同）。

例如，在图 5(b)中，事件和均为任务的前置条件（即），通过后继集融合可将它们合并为新事件（图 5(c)）。

2. 前驱集融合

前驱集融合合并那些依赖于同一组生产者任务（即相同）的事件。由于这些事件会被同时激活，将它们作为独立同步节点保留会增加不必要的图复杂度。

定义 4.2（前驱集融合）：对于 tGraph 中的任意两个事件和，当且仅当时，可应用前驱集融合。MPK 从 tGraph 中移除和，并引入一个融合事件，其中（的触发任务与相同），（的依赖任务为和的依赖任务的并集）。

例如，在图 5(c)中，事件、、、均依赖于任务和（即），通过前驱集融合可将它们合并为单个新事件（新 tGraph 中）。

MPK 需要解决的一个核心挑战是如何表示任务与事件之间的依赖关系。由于 MPK 在 SM 间并行执行任务和更新事件，运行时需要一种统一且低成本的表示方式，避免高昂的间接索引开销。这带来两个挑战：

一个任务可能依赖于多个事件，也可能触发多个事件。直接为每个任务预留最大数量的依赖事件和触发事件存储空间，会导致显著的内存开销；
事件融合后，一个事件可能需要触发大量任务。为每个事件预留最大“扇出”（fan-out，即事件触发的任务数量）存储空间同样代价高昂。

MPK 通过两种技术解决这些挑战：tGraph 规范化（tGraph normalization）和 tGraph 线性化（tGraph linearization）。

tGraph 规范化

MPK 通过 tGraph 规范化解决第一个挑战：将输入 tGraph 转换为功能等价的形式，确保每个任务最多有一个依赖事件和一个触发事件。

这一目标通过两种转换实现，将每个任务的最大“扇入”（fan-in，即任务依赖的事件数量）和“扇出”（任务触发的事件数量）均降至 1。

转换 1：减少任务的触发事件数量

当输入 tGraph 中的一个任务触发多个事件时，MPK 会修改 tGraph：引入一个新事件和个空的新任务（这些新任务不执行任何计算，且均依赖于）。转换后，任务仅触发，每个新任务则恰好触发一个原始事件（如图 6a 所示）。这种转换确保每个任务仅有一个触发事件。

下图图 5(e)展示了 MPK 如何应用此转换，将任务和的触发事件数量减少到 1。

图 5：MPK 编译器工作流程。该图展示了 MPK 编译器从输入计算图到生成 GPU 可执行结构的端到端流程。图（b）中，Q、K、V、A、O、R 分别代表分解查询投影、键投影、值投影、注意力、输出投影、RMSNorm 算子后生成的任务集。图（e）中，D1、D2 是 tGraph 规范化过程中插入的虚拟任务，用于确保每个任务仅有一个触发事件。图（f）展示了 MPK 如何将 tGraph 线性化并存储到 GPU 内存，任务与事件均采用统一标准表示。

流程涵盖算子分解、依赖分析、事件融合、tGraph 规范化与线性化等关键步骤。算子分解将大算子拆分为 SM 级任务；依赖分析捕捉细粒度依赖；事件融合（后继集/前驱集融合）减少冗余同步；规范化确保任务仅有一个依赖/触发事件；线性化则通过 BFS 算法让同一事件触发的任务索引连续，便于高效存储与调度。这些步骤环环相扣，既保证了 tGraph 的正确性与高效性，又为核内运行时的低开销调度奠定基础，是 MPK 实现自动优化的核心流程。

转换 2：减少任务的依赖事件数量

当输入 tGraph 中的一个任务依赖于多个事件时，MPK 会修改 tGraph：引入一个新事件和若干个空的新任务（这些新任务不执行任何计算，且均触发该新事件）。转换后，原任务和每个新任务均仅依赖于一个事件（如图 6b 所示）。

图 6：MPK 通过图变换规范化任意 tGraph。该图展示了 MPK 通过两种关键图变换对任意 tGraph 进行规范化的过程，确保每个任务最多仅有一个依赖事件和一个触发事件。针对任务多扇出/扇入事件的场景：当任务触发多个事件（图 6a），MPK 引入新事件与虚拟任务，让原任务仅触发新事件，虚拟任务触发原事件；当任务依赖多个事件（图 6b），则引入新事件与虚拟任务，让虚拟任务触发新事件，原任务依赖新事件。此设计虽会增加少量虚拟任务，但能统一任务的依赖/触发结构，简化核内运行时的任务管理与调度逻辑。实践中，由于真实模型多为“深度”结构（包含许多顺序执行的算子），并行算子较少，因此规范化引入的开销通常低于 1%，是平衡正确性与效率的关键设计。

tGraph 规范化通过添加额外任务和事件来“规范化”存在多扇入/多扇出事件的 tGraph，这会引入一定开销。但这种情况通常仅在原始计算图包含可并行执行的算子时才会发生。例如，图 5(d)中的任务有两个扇出事件，因为原始计算图中的 RMSNorm（一种归一化算子）和输出投影算子均依赖于注意力算子，且可并行运行。

在实践中，我们观察到规范化开销可忽略不计（评估中始终低于 1%），因为真实世界的模型通常是“深”的（包含许多顺序执行的算子），而非“宽”的（包含许多可并行执行的算子）。

tGraph 线性化

仅靠 tGraph 规范化无法解决第二个挑战：融合后，一个事件仍可能需要触发大量任务（例如图 5(e)中的事件触发 4 个任务），需额外存储来记录这些任务的索引。

MPK 通过基于广度优先搜索（BFS）的算法（见算法 1）对 tGraph 进行线性化，解决这一问题。

算法 1：MPK 的 tGraph 线性化算法。该算法确保每个任务仅入队到任务列表 T 一次，每个事件仅入队到事件队列 E 一次。第 5-7 行确保所有依赖于同一事件的任务在任务列表 T 中是连续排列的。在规范化 tGraph 后，事件可能触发大量任务，直接存储任务索引列表会产生高额内存开销。算法通过广度优先搜索（BFS），以事件为驱动遍历任务图，确保同一事件触发的任务在列表中连续排列。这种设计让事件的扇出信息可通过“首任务索引 + 尾任务索引”紧凑编码，无需存储完整任务列表，大幅降低 GPU 设备内存占用。同时，连续的任务排列也为核内运行时快速调度任务提供便利，是连接 tGraph 规范化与 GPU 内存存储的核心桥梁。

线性化确保由同一事件触发的所有任务在最终任务顺序中被分配连续的索引。因此，事件的扇出可仅用“第一个任务索引”和“最后一个任务索引”紧凑编码，无需存储依赖任务的显式列表，同时保留所有依赖语义。

图 5：MPK 编译器工作流程。

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

4.2 任务实现生成

除构建 tGraph 外，MPK 还需为每个任务生成在 GPU SM 上执行的设备函数。MPK 借鉴现有研究成果，采用 编译器超优化（superoptimization，一种通过穷举或智能搜索寻找最优代码序列的优化技术）方法，为每个任务自动生成高性能实现。

具体而言，MPK 在“线程块级”而非“核函数级”执行超优化： 每个计算任务都关联一个参考 PyTorch 实现，MPK 利用 Mirage 超优化器[33]自动搜索最优线程块图（thread block graph，描述线程块内计算和数据流动的结构）， 并将该图发送给 Mirage 编译器以生成 CUDA 实现。生成的 CUDA 实现包含所有 SM 内优化（如软件流水线、寄存器复用、布局优化），以最小化共享内存 bank 冲突（bank conflict，指多个线程同时访问共享内存同一 bank 导致的性能损耗）。

五、核内并行运行时（In-Kernel Parallel Runtime）

MPK 采用核内并行运行时，可在单个巨型内核（mega-kernel）中跨所有流式多处理器（SM）执行 tGraph。该设计消除了内核启动开销，并能对调度、同步和执行顺序进行细粒度控制。 巨型内核启动后，会持续管理计算和通信过程，直至推理任务完成。

为支持这种执行模型，MPK 将 GPU 的流式多处理器（SM）划分为工作器（worker）和调度器（scheduler）两类角色。
* 每个工作器运行在一个物理 SM 上，并维护一个独立的任务队列。工作器会执行一个轻量级循环：不断从队列中取出任务、执行相关计算或通信操作，再通过通知任务的触发事件来标记任务完成。这种设计既能确保工作器被充分利用，又能支持跨算子的异步执行。
* 调度器以线程束（warp）为粒度进行组织，每个 SM 承载 4 个调度器线程束。每个调度器维护一个事件队列，持续轮询新激活的事件，并将对应的任务分配给工作器。

工作器和调度器的分配在核启动时就已固定，且与 GPU 的物理 SM 数量相匹配，避免了在内核内部动态切换角色带来的开销。本节剩余部分将详细介绍核内运行时的架构。
* 5.1 节阐述 MPK 的事件驱动执行模型；
* 5.2 节介绍两种互补的任务启动机制，分析它们的权衡关系，并说明 MPK 如何结合这两种机制实现低延迟、负载均衡的执行；
* 5.3 节描述 MPK 的额外系统优化措施，这些措施可进一步降低开销并提升吞吐量。

5.1 事件驱动执行（Event-Driven Execution）

MPK 采用事件驱动模型执行 tGraph。每个 tGraph 以一个无前置条件的指定起始事件（例如图 7 中的 e0）开始，该事件会被初始加入调度器的事件队列。调度器从队列中取出事件后，会启动所有依赖于此事件的任务（例如 AT）。每个被启动的任务会分配给一个工作器，工作器执行任务后，会触发与该任务关联的后续事件，以标记任务完成。

当一个事件的所有前置任务均已完成，且这些任务已共同触发该事件达到所需次数时，该事件即被激活。激活的事件会被重新加入调度器的事件队列，从而持续推动 tGraph 的执行流程。通过这种方式，事件成为驱动任务执行的核心机制，实现了细粒度的异步调度与执行。

5.2 混合任务启动

MPK 支持两种任务启动方式：即时启动与提前启动。

即时启动：调度器仅在任务的依赖事件完全激活后，才将任务分配给工作器。任务分配后可立即开始执行。
提前启动：运行时会在任务的前置事件激活前，就将任务提前加入到预先分配好的工作器队列中。工作器需等待其依赖事件完全激活后才能执行该任务。

JIT 和 AOT 两种方式优势互补，共同构成 MPK 的混合任务启动策略。

JIT 启动有助于应对工作负载不均衡。例如，LLM 中注意力算子的执行时间因输入序列长度不同而差异显著。在 JIT 模式下，MPK 仅在注意力任务完成后才启动后续任务（如矩阵乘法或全归约）。较早完成注意力任务的工作器可以执行更多后续任务，从而缩短端到端延迟，并实现流式多处理器间的负载均衡。

AOT 启动则能降低任务分配延迟。在 JIT 模式下，从任务完成到后续任务启动，需要工作器与调度器之间进行两次同步通信。而在 AOT 模式下，任务已提前分配至工作器队列，工作器仅需等待依赖事件激活即可执行，仅需一次同步操作，从而降低了单次任务启动的延迟。

在 tGraph 构建阶段，编译器会根据算子特性将其分类：
* JIT 类型：执行时间依赖数据的算子（如注意力算子）及其下游任务，直到遇到全局屏障（消除负载不均衡）。
* AOT 类型：其余所有算子，以最小化任务分配开销。

工作器维护独立的 JIT 和 AOT 任务队列，并优先执行 JIT 队列中已就绪的任务。调度器仅处理 JIT 相关的事件，AOT 任务则通过轮询方式提前静态分配至各工作器队列。这种混合策略既通过 AOT 分摊了分配开销、减轻了调度器负载，又通过 JIT 实现了动态的负载均衡。

5.3 运行时优化

分页共享内存抽象

在 CUDA 等传统 GPU 编程模型中，共享内存的生命周期与内核执行绑定，内核结束后内存即被释放，且每个内核默认独占全部共享内存。这种设计阻碍了跨任务的软件流水线实现，因为流水线要求将后续任务的数据加载与当前任务的计算重叠执行，而这两个任务都需要访问共享内存。

为支持此类流水线，MPK 引入了分页共享内存抽象。共享内存被划分为多个固定大小的页面，任务实现被修改为基于这些页面运行，而非假设独占整块共享内存。一个任务可根据其共享内存占用量申请一个或多个页面，并在不再需要时释放。任务释放任一页面后，将不再允许申请更多共享内存，这种单调使用模式简化了调度。当任务释放页面后，MPK 可立即为后续任务预分配可用页面并开始数据预取。这种设计支持在巨型内核执行模型中对共享内存资源进行细粒度的按需分配。

跨任务软件流水线（Cross-Task Software Pipelining）

为支持同一工作器上执行的跨任务软件流水线（见 2.1 节），MPK 将每个任务分解为预加载阶段（pre-loading phase）和计算阶段（compute phase）。
* 预加载阶段仅执行数据传输指令，将所需张量的一部分从设备内存加载到共享内存，初始化任务内软件流水线，但不执行任何计算操作。
* 计算阶段执行实际的计算任务。当满足以下两个条件时，MPK 会伺机将当前任务 T_i 的计算阶段与后续任务 T_{i+1} 的预加载阶段重叠执行：（1）T_i 已发出所有自身的数据传输指令；（2）有足够的共享内存页面可供 T_{i+1} 的预加载阶段使用。需注意的是，这种流水线不会干扰 T_i 的执行，因为 MPK 会插入适当的 SM 内同步屏障，确保 T_{i+1} 的内存传输不会与 T_i 正在进行的数据传输产生冲突。

任务描述预取（Pre-fetching Task Descriptions）

每个工作器在设备内存中维护 JIT 和 AOT 两个任务队列。每个任务都关联一个任务描述，其中包含输入张量、输出张量和配置参数的信息；在当前实现中，每个任务描述占用 352 字节（见 6.1 节）。为降低任务入队/出队延迟并隐藏设备内存访问开销，MPK 采用轻量级预取机制，在任务需要执行前，将后续任务的描述加载到共享内存中。

六、评估

本节通过回答三个关键问题对 MPK 进行评估：
* 首先，6.3 节将 MPK 的巨型内核执行模型与当前最先进的“每个算子一个内核”系统进行对比；
* 其次，6.5 节考察 MPK 在多 GPU 执行大规模深度神经网络（DNN）任务时的可扩展性和效率；
* 最后，6.6 节分析 MPK 中的各项优化措施对整体性能提升的贡献。

本评估聚焦于大型语言模型（LLM）服务场景，原因有二：
* 第一，LLM 服务领域已有多个经过高度优化的“每个算子一个内核”基准系统（包括 SGLang 和 vLLM），与这些系统对比可构成严格的性能基准，凸显 MPK 巨型内核方案的优势；
* 第二，LLM 服务天然具有动态执行特性——每次服务迭代的批大小、序列长度以及预填充与解码阶段的比例都可能存在显著差异，形成异构工作负载，对编译器和运行时均提出挑战。

这种可变性使 LLM 服务成为评估 MPK 的代表性且具有挑战性的场景。需说明的是，MPK 编译器和运行时具有模型无关性，可轻松支持任意 DNN 架构。

6.1 实现细节

我们将 MPK 实现为 PyTorch 的内核后端。通过 PyTorch 的编译接口，指定 MPK 为后端（即 torch.compile(backend=MPK)），即可将 PyTorch 程序编译为 MPK 巨型内核。该调用会触发 MPK 编译器，生成巨型内核并将其作为可调用的 PyTorch 函数返回。
* 执行该函数时，会启动生成的巨型内核一次。
* 当前 MPK 实现包含约 40,000 行 C++代码、84,000 行 CUDA 代码和 10,000 行 Python 代码。
* 核内并行运行时采用 CUDA 编写，通过设备内存中的信号量实现工作器与调度器之间的协调。
* MPK 编译器采用 C++和 Python 实现，可自动将输入的张量程序转换为针对特定 GPU 类型优化的 tGraph。
* 对于计算任务，编译器集成了 Mirage 超级优化器以自动生成优化的 CUDA 实现，并使用 NVSHMEM 支持核内跨 GPU 通信。

我们的实现还包含多项关键优化，以最小化运行时开销并高效支持动态工作负载。

任务启动开销（Task-Launch Overhead）

由于 MPK 将计算分解为比传统 GPU 内核粒度更细的任务，因此最小化单次任务启动开销对性能至关重要。MPK 采用多种技术降低任务启动成本：
* 轻量级数据结构：运行时使用极轻量级的工作器和调度器——事件队列和任务队列均实现为 GPU 设备内存中的循环缓冲区，入队/出队操作仅依赖低成本的 atomicAdd 指令。
* 去中心化调度：MPK 采用去中心化调度策略，每个调度器仅使用本地状态分配任务，避免了全局协调，消除了全局协调调度中固有的通信和同步开销。

尽管 MPK 目前采用去中心化调度，但运行时的设计支持通过少量代码修改切换到其他调度策略（包括全局协调调度）。探索这些设计及其性能权衡是未来工作的一个重要方向。

支持运行时动态性

为证明 MPK 支持高度动态的工作负载，我们扩展了系统以支持 LLM 服务所需的机制，包括连续批处理和分页注意力。当处理任务图的起始事件时，调度器会通过以下步骤准备新的解码迭代：
1. 移除上一迭代中已完成的请求；
2. 接纳新到达的请求；
3. 更新每个请求的键值缓存元数据。

所有这些逻辑都在单个任务中执行，且 KV-cache 元数据存储在设备内存中，可供分页注意力任务直接使用。

为处理 LLM 服务中固有的动态批大小问题，MPK 为代表性批大小（2 的幂，最大为最大批大小）生成多个专用 tGraph。运行时会根据当前批大小选择合适的 tGraph。这种方式既允许编译器为特定批大小生成优化的 tGraph，又能为高度动态的工作负载保留灵活性。

6.2 实验设置

我们在 5 个广泛部署的 LLM（涵盖不同参数规模和架构类型）以及 3 代 NVIDIA GPU（A100、H100、B200）上对 MPK 进行评估。

表 1：评估中的 MPK 配置。该表展示了评估中 MPK 在不同 GPU 上的配置参数，包括 GPU 型号、流多处理器（SM）数量、工作器（Worker）数量、调度器（Scheduler）数量，为实验可重复性与性能分析提供关键参数参考。MPK 的核内运行时将 GPU 的 SM 划分为工作器（执行任务）与调度器（管理事件与任务分配），配置需平衡二者资源。实验中，所有 GPU 均预留 4 个 SM（16 个调度器 warp，每个 SM 支持 4 个并发 warp）作为调度器，剩余 SM 作为工作器（A100：108-4=104，H100：132-4=128，B200：148-4=144）。这种配置既保证调度器有足够资源高效处理事件与任务分配，又最大化工作器数量以提升计算吞吐量。同时，共享内存页大小固定为 32KB，各 GPU 的 SM 共享内存页数（A100：5，H100：7，B200：7）适配硬件特性，为页面化共享内存抽象（支持跨任务流水线）提供基础，是 MPK 在不同 GPU 上实现高效运行的重要参数保障。

表 1 总结了每种 GPU 类型对应的 MPK 配置：所有实验中，我们预留 4 个 SM 用于调度器（共分配 16 个调度器线程束，每个 SM 最多可容纳 4 个并发活跃的线程束），剩余 SM 分配给工作器。所有 GPU 的共享内存页面大小均固定为 32KB，该配置使 A100、H100、B200 每个 SM 分别拥有 5、7、7 个共享内存页面。

图 9：MPK 与现有系统在 A100、H100、B200 GPU 上对五种模型的性能对比。该图对比了 MPK 与 SGLang、vLLM 在 A100、H100、B200 三种 GPU 上对五种模型的性能，所有性能数据均以 MPK 为基准进行归一化（数值越高越好）。MPK 柱状图上方的数值表示其相对于性能最佳现有系统的加速比。由于模型超出单 A100 内存容量，故省略 Qwen3-30B-A3B 在 A100 上的结果。实验覆盖稠密模型与 MoE 模型，跨三代 GPU。结果显示，MPK 单批次推理性能提升 1.0-1.7 倍，小模型与新 GPU 上提升更显著。例如，Qwen3-8B 在 A100 上，MPK 将解码延迟从 14.5ms（vLLM/SGLang）降至 12.5ms，接近硬件理论下限（约 10ms）。

图 9 展示了评估中使用的 LLM，包括稠密模型和混合专家模型，覆盖多种模型规模。为控制请求到达模式的可变性，所有实验均在离线批处理推理设置下进行，最大批大小各不相同。所有请求的提示长度固定为 64，解码 1024 个 token。这种方法消除了因请求并发不足导致的服务器端停滞，使性能对比能清晰反映 MPK 的系统级性能优势。

6.3 端到端结果

首先，将 MPK 的端到端服务性能与当前最先进的两个 LLM 服务系统 SGLang 和 vLLM 进行对比。SGLang 和 vLLM 均采用“每个算子一个内核”模式，并依赖多种专用内核库：注意力算子使用 FlashInfer 和 FlashAttention 优化实现，矩阵乘法使用 cuBLAS 和 cuTLASS，其余算子使用 CUDA 或 Triton 实现。所有系统均从 Hugging Face Transformers 加载模型架构，采用 bfloat16 精度格式，并启用分页注意力和连续批处理。

关键架构差异在于：MPK 将页面分配和请求调度直接集成到巨型内核内部；而 SGLang 和 vLLM 在 CPU 上执行这些操作，会产生额外的主机-设备同步和分配开销。

对于每个模型，在 B200、H100、A100 GPU 上，针对 1 到 16 的不同最大批大小，评估了三个系统的性能，并报告服务吞吐量。由于 Qwen3-30B-A3B 模型超出了单个 A100 GPU 的内存容量，因此未在 A100 上评估该模型。

图 2：MPK 与现有方法在任务内和跨任务流水线支持上的对比。该图对比了 MPK 与现有方法在支持任务内和跨任务流水线方面的差异。现有系统受限于核屏障，仅能实现任务内流水线（如图 2a），不同任务间存在流水线气泡，导致 GPU 异构资源（如张量核、TMA）无法充分利用。而 MPK 通过 SM 级任务调度，打破任务边界，实现跨任务流水线，让数据预加载与计算并行。这种优化对 GPU 架构日益异构的趋势至关重要，能充分发挥 TMA 异步数据传输与张量核计算的协同作用。

图 9 展示了端到端吞吐量结果。在单批推理场景下，MPK 在不同模型和硬件上的服务性能提升了 1.0~1.7 倍。对于较小模型和较新一代 GPU，性能提升更为显著，这源于三个因素：
1. 即使使用 CUDA Graphs，“每个算子一个内核”模式仍存在较高的内核启动开销；
2. 由于内核抽象阻碍了跨任务流水线（见图 2），这些系统存在流水线空转问题；
3. SGLang 和 vLLM 在 CPU 上执行页面分配和请求调度，增加了 CPU-GPU 同步延迟。

随着模型规模减小和硬件性能提升，这些开销在总执行时间中的占比显著增加。6.6 节将通过消融实验评估这些优化措施的影响。

结果表明，MPK 非常适合低延迟优化场景（如单批、低延迟模型服务），能将 LLM 推理延迟推向硬件极限。例如，在 A100 GPU 上运行 Qwen3-8B 模型时，MPK 将每个 token 的解码延迟从 vLLM 和 SGLang 等高度优化系统已实现的 14.5 毫秒，进一步降低至 12.5 毫秒，接近理论下限（约 10 毫秒，根据 1.6 TB/s 内存带宽加载 16 GB 模型参数估算）。

除性能提升外，MPK 相比 vLLM 和 SGLang 的另一关键优势是易用性。vLLM 和 SGLang 需要大量工程工作来手动优化新模型并集成专用内核，而 MPK 采用编译器驱动的方案，只需修改少量代码即可自动将 PyTorch 模型转换为巨型内核。这种设计使 MPK 在保持 PyTorch 熟悉开发流程的同时，实现了高性能——相比原生 PyTorch，性能提升超过 10 倍。

6.4 案例研究：混合专家模型

为高效服务混合专家模型等动态工作负载，MPK 包含两项针对 MoE 的专用优化：
1. 混合工作负载均衡器；
2. 融合聚集-矩阵乘法实现。

混合工作负载均衡器

由于每个专家接收的 token 数量仅在运行时可知，因此预先确定有效的工作负载划分方案具有挑战性。朴素的静态方案会将固定组的 SM 分配给预先指定的部分专家，但当 token 分配高度不均衡时，这种策略会导致严重的负载失衡。与之相反，基于持久化分组通用矩阵乘法的全动态方案虽能实现 SM 间的负载均衡，但会引入显著的细粒度同步开销。

MPK 引入了一种混合策略，结合两种方案的优势：
* 编译阶段，编译器将工作静态划分为专家专用任务；
* 运行时，每个任务会接收由 topk-softmax 内核生成的元张量，其中包含全局 MoE 信息（具体为激活的专家数量和每个专家分配到的 token 数量）。

任务利用这些全局元数据动态调整工作负载分配，在均匀拆分工作的同时，避免全动态调度带来的同步开销。如图 10 所示，在所有批大小下，混合策略的性能均持续优于纯静态划分。

图 10：MPK 与现有系统在 B200 上对 Qwen3-30B-A3B 模型的性能对比。该图对比了 MPK 与现有系统在 B200 GPU 上对 Qwen3-30B-A3B 模型的性能，每个数值代表各方法实际的 MoE 运行时间（单位：微秒，数值越低越好）；柱状图上方的数值表示 MPK-Hybrid-MoE 相对于 SGLang-MoE 的加速比。

融合聚集-矩阵乘法

为利用 Hopper 和 Blackwell 架构 GPU 上的张量内存加速器（TMA），当前 MoE 实现会先执行聚集（gather）步骤，将分配给同一专家的 token 打包为连续内存布局。以 SGLang 中批大小为 1 的 Qwen3-30B-A3B 模型为例，该预处理步骤的执行时间占 MoE 总执行时间的比例最高可达 11%。此外，在 MPK 中引入额外预处理任务会进一步增加调度开销。

为解决这些问题，MPK 将基于 TMA 的聚集步骤替换为异步 token 级复制，并将其直接集成到矩阵乘法（GEMM）任务的数据加载阶段。这种融合设计消除了独立的聚集内核，避免了额外的调度节点，同时仍能实现高效的数据移动。因此，采用融合聚集-矩阵乘法的 MPK 相比 SGLang 的实现，实现了持续的性能提升。

6.5 多 GPU 结果

我们在 NVIDIA H100 DGX 实例上进一步评估了 MPK 在多 GPU 场景下的可扩展性。与对比基准系统一致，我们采用 Megatron-LM[27]中提出的张量模型并行（tensor model parallelism）策略。

用户只需在注意力和门控 MLP 块后插入全归约（AllReduce）层，即可指定张量并行执行；MPK 会自动将这些集合通信算子（collective operator）编译为两类任务：

跨 GPU 数据传输任务（使用 NVSHMEM 的nvshmem_signal_wait_until实现）；
本地归约任务。

这种分解将同步集合通信转换为完全异步的任务，使其能与 MPK 基于任务的事件驱动运行时无缝集成。

图 11：MPK 与现有系统在多 H100 GPU 上对 Qwen3-1.7B 模型张量并行推理的性能对比。所有系统的性能均以 MPK 为基准进行归一化（数值越高越好）。

图 11 展示了性能结果。PyTorch 结合了手动优化内核、CUDA Graphs 和torch.compile，而 MPK 的巨型内核执行模式相比 PyTorch，吞吐量提升最高可达 10 倍；相比 SGLang 和 vLLM 等高度优化的服务系统，MPK 在扩展到 8 个 H100 GPU 时，仍能实现 1.1~1.4 倍的性能提升。

这些性能增益源于“每个算子一个内核”基准系统缺失的三项关键优化：

MPK 将页面分配和请求调度直接集成到巨型内核内部，消除了 CPU 端的分配开销；
MPK 的异步执行模型实现了 计算任务与集合通信的重叠执行；
MPK 消除了内核屏障，支持跨任务软件流水线。

6.6 节将详细分析后两项优化的影响。

6.6 消融实验

本节评估 MPK 中两项关键优化措施的影响：跨任务流水线（cross-task pipelining）和计算-通信重叠（compute-communication overlap）。

跨任务流水线（Cross-Task Pipelining）

如 5.3 节所述，MPK 通过在当前任务计算的同时，预加载下一任务的输入张量块，实现跨任务流水线。

图 12：跨任务流水线消融实验。该图是跨任务流水线的消融实验结果，实验在 NVIDIA B200 GPU 上测量 Qwen3-8B 模型最终线性层的性能，报告实际运行时间（单位：微秒，数值越低越好）。

图 12 展示了针对 Qwen3-8B 模型最后一个线性层的消融实验结果：跨任务流水线将任务运行时间 缩短了 1.2~1.3 倍，甚至优于 cuBLAS 编译的内核。

计算-通信重叠（Compute-Communication Overlap）

MPK 捕捉任务间的细粒度依赖（见图 5），使运行时能伺机重叠执行计算任务和通信任务。

图 5：MPK 编译器工作流程。

图 13：计算 – 通信重叠消融实验。该图是计算 – 通信重叠的消融实验结果，实验在 4 块 H100 GPU 上测量 Qwen3-1.7B 模型张量并行推理的每迭代运行时间（单位：微秒，数值越低越好）。

图 13 展示了评估该重叠优化影响的消融实验：通过仅使用单个事件捕捉集合算子（如全归约）与前置任务间的粗粒度依赖（如图 5c 所示），即可禁用该重叠优化。结果显示，启用计算-通信重叠后，每次迭代的延迟 降低了 1.1 倍。

七、相关工作

7.1 手动设计内核（Manually Designed Kernels）

7.1 现有机器学习框架与专用内核

TensorFlow XLA[1,12]、PyTorch[25]、TensorRT[28]等主流机器学习框架普遍采用“每个算子一个内核”的模式，高度依赖GPU领域专家为单个算子手动设计和实现内核。以注意力算子为例，业界已开发出FlashAttention[5,17,20]、FasterTransformer[3]、FlashInfer[35]等多种专用内核，每种内核都针对特定的硬件架构特性或应用场景进行了深度优化。然而，当前系统依赖于这种分散的专用内核库生态， 难以将整个推理流水线统一到一个整体的巨型内核中。

7.2 机器学习编译器（ML Compilers）

大量研究工作聚焦于通过编译器自动生成高性能的张量程序内核。
* TVM[15,16]、Ansor[37]、Triton[29]等系统，以及其他相关工作[18,19,21,40]，均基于Halide[24,26]提出的“算法-调度分离”思想构建。
* 另一类研究则采用超级优化技术，从高层规范出发自动搜索高效的内核实现[22,31,32,34,39]。

然而，这些编译器本质上仍围绕算子级优化进行设计，既不支持生成统一的巨型内核，也无法有效协调跨算子的执行过程。

7.3 巨型内核（Mega-Kernels）

以往关于巨型内核的研究大多依赖于繁重的手动设计。
* 例如，FlashDMoE将混合专家模型的计算与跨GPU通信融合为一个手动设计的巨型内核[13]。
* 又如，Spector等人针对LLaMA-1B模型手动开发了低延迟巨型内核[9,30]。

这些方法需要投入大量的工程努力和深厚的GPU专业知识，并且难以跨不同模型或GPU平台进行推广。相比之下，MPK提供了一种基于编译器的自动化解决方案，能够将张量程序自动转换为高度优化的巨型内核，无需进行手动融合或编写内核实现。

八、结论

本文提出了MPK（Mirage Persistent Kernel）——首个能够自动将多GPU模型推理转换为完全融合巨型内核的编译器和运行时系统。

通过引入流式多处理器级的任务图和核内并行运行时，MPK克服了“每个算子一个内核”模式的根本局限，实现了跨算子的软件流水线、计算与通信的细粒度重叠执行，并消除了内核启动和CPU端调度的开销。

评估结果表明，MPK将大型语言模型的服务延迟推向了硬件极限，并在不同模型和GPU代际上显著提升了吞吐量。MPK通过在单个巨型内核中统一执行流程，同时保留了现有机器学习框架的易用性，为构建高性能、编译器驱动的推理系统开辟了新的路径。

参考文献

关注“鲸栖”小程序，掌握最新AI资讯

本文来自网络搜集，不代表鲸林向海立场，如有侵权，联系删除。转载请注明出处：http://www.itsolotime.com/archives/17087