技术摘要：MPK —— 一种用于超大规模内核化张量程序的编译器与运行时

1. 问题陈述

当前用于大语言模型（LLM）和其他张量程序的高性能 GPU 推理系统主要依赖于逐算子内核执行模型（kernel-per-operator execution model）。在这种范式下，每个算子（例如矩阵乘法、注意力机制、All-Reduce）都作为一个独立的 GPU 内核进行执行。虽然这种方法简化了编程，但它引入了三个关键限制，阻碍了性能表现：

1. 隐式内核屏障（Implicit Kernel Barriers）： 传统的 GPU 执行在同一流上的连续内核启动之间强制执行屏障。这阻止了跨算子软件流水线化（cross-operator software pipelining），迫使依赖算子必须严格顺序执行，从而产生了“流水线气泡（pipeline bubbles）”，导致硬件资源（如 Tensor Cores、TMA）处于闲置状态。
2. 粗粒度重叠（Coarse-Grained Overlap）： 依赖关系是在整个算子的粒度上强制执行的。因此，计算与通信之间的细粒度重叠是不可能的。例如，一个 AllReduce 操作必须等待前序 MatMul 的全部完成后才能开始，尽管该归约操作中的特定片段仅依赖于乘法输出的特定片段。
3. 启动开销与静态灵活性（Launch Overhead and Static Flexibility）： 每次推理迭代执行数百或数千个内核会产生显著的启动开销。虽然 CUDA Graphs 可以缓解这一问题，但它们在很大程度上是静态的；动态变化的控制流、张量形状或数据依赖关系通常需要重新实例化图，这限制了处理如 LLM 服务等动态工作负载时的灵活性。

现有的解决方案（如手动编写的 mega-kernels，例如 FlashDMoE, Spector 等）虽然提供了性能优势，但需要大量的工程投入和深厚的 GPU 专业知识，无法在不同模型或硬件之间实现泛化。

2. 方法论

作者提出了 Mirage 持续内核（Mirage Persistent Kernel, MPK），这是第一个能够自动将多 GPU 模型推理转换为单个高性能 mega-kernel 的编译器与运行时系统。MPK 通过以下两个主要组件实现：

A. SM 级图表示 (tGraph)

MPK 引入了一种名为 tGraph 的新型图表示，它在**流式多处理器（Streaming Multiprocessors, SM）**的粒度上捕捉计算与通信，而非整个 GPU 级别。

• 节点（Nodes）： 代表在单个 SM 上运行的任务（计算或通信）。
• 边（Edges）： 通过**事件（events）**对任务之间的细粒度依赖关系进行编码。
• 优势： 这种粒度暴露了在内核级模型中无法实现的并行性，从而实现了跨算子的软件流水线化以及计算与通信的细粒度重叠。

编译器优化：
MPK 编译器通过多个阶段将标准的张量程序转换为优化的 tGraph：

1. 算子分解（Operator Decomposition）： 根据输出张量的划分方式（例如，对矩阵乘法的行/列进行分块）将算子拆分为不相交的任务集，以匹配可用的 SM 数量。
2. 依赖分析（Dependency Analysis）： 枚举任务之间的生产者-消费者对，以便仅在存在数据重叠时插入同步事件。
3. 事件融合（Event Fusion）： 使用后继集融合（successor-set fusion）（合并具有相同消费者集的事件）和前驱集融合（predecessor-set fusion）（合并具有相同生产者集的事件）来合并冗余的同步点，以降低图的复杂度。
4. 归一化（Normalization）： 转换图结构，使得每个任务最多只有一个依赖事件和一个触发事件。这通过插入虚拟任务和事件来处理扇入/扇出（fan-in/fan-out），确保实现统一且低开销的表示。
5. 线性化（Linearization）： 使用基于 BFS 的算法，为由同一事件触发的任务分配连续索引。这允许运行时将事件的扇出编码为简单的索引范围（首个/最后一个任务 ID），而不是显式的列表，从而显著降低了设备内存占用。
6. 任务实现生成（Task Implementation Generation）： 利用编译器超优化技术（Mirage）自动生成每个独立任务的高性能 CUDA 实现，包括寄存器复用和共享内存布局等 SM 内部优化。

B. 内核内并行运行时 (In-Kernel Parallel Runtime)

MPK 在一个单一的持续性 mega-kernel 中执行 tGraph，采用去中心化的事件驱动架构：

• Worker/Scheduler 分区： 将 GPU 的 SM 静态划分为 Workers（执行任务）和 Schedulers（管理事件队列）。
• 事件驱动执行： Schedulers 轮询已激活的事件（即所有前提条件均已满足的事件），并将相应的任务分发给 Workers。这消除了对主机端同步或重复内核启动的需求。
• 混合任务启动（Hybrid Task Launch）： 为了平衡延迟与负载均衡，MPK 采用了两种模式：
  • 即时启动 (Just-In-Time, JIT)： 仅在依赖关系满足后才启动任务。这能适应动态工作负载（例如，Attention 中的可变序列长度）以防止负载不均衡。
  • 预先启动 (Ahead-Of-Time, AOT)： 在依赖关系满足之前就将任务预先入队到 Worker 队列中。这降低了可预测算子的调度延迟。
• 分页共享内存（Paged Shared Memory）： 为了实现跨任务的软件流水线化，MPK 引入了共享内存的页面抽象。任务获取并释放固定大小的页面，允许下一个任务在当前任务计算的同时，将数据预取到共享内存中，而不会产生资源冲突。

3. 核心贡献

1. 首个自动化 Mega-Kernel 系统： MPK 是第一个能够自动将任意多 GPU 张量程序转换为单个 mega-kernel 的系统，消除了对手动内核融合的需求。
2. SM 级抽象： 引入 tGraph 表示，将依赖跟踪从算子级转向 SM-任务级，实现了此前在内核逐算子模型下无法实现的优化（流水线化、细粒度重叠）。
3. 端到端编译器与运行时： 该系统集成了一个优化任务图并生成 CUDA 代码的编译器，以及一个在单个内核内异步执行这些任务的运行时，支持连续批处理（continuous batching）和分页注意力机制（paged attention）等动态工作负载。
4. 混合执行策略： 结合 JIT 和 AOT 任务启动，使系统既能处理 LLM 推理的动态特性（可变的 Batch/序列长度），又能保持静态算子的执行效率。

4. 评估结果

作者在五种广泛使用的 LLM（包括 Qwen3 和 Llama-3 变体）以及三代 NVIDIA GPU（A100, H100, B200）上对 MPK 进行了评估，并将其与最先进的逐算子内核系统（SGLang, vLLM）及原生 PyTorch 进行了对比。

• 性能提升： 与 SGLang 和 vLLM 相比，MPK 实现了 1.0× 至 1.7× 更低的端到端推理延迟。在较小的模型和更新的硬件世代（如 B200）中，这种提升最为显著，因为在这些场景下，内核启动开销和流水线气泡相对于计算时间变得更加突出。
• 吞吐量： 在多 GPU 设置下（高达 8 个 H100），MPK 比优化后的基准系统提升了 1.1× 至 1.4× 的吞吐量。
• 延迟极限： 在 A100 上，MPK 将 Qwen3-8B 的每 Token 解码延迟从 14.5 ms（vLLM/SGLang）降低至 12.5 ms，接近了约 10 ms 的估计硬件下限。
• 消融实验：
  • 跨任务流水线化（Cross-Task Pipelining）： 减少了 1.2–1.3× 的任务运行时间。
  • 计算-通信重叠（Compute-Communication Overlap）： 减少了 1.1× 的单次迭代延迟。
  • 减少内核启动（Kernel Launch Reduction）： 消除了基准系统中每 Token 对应 293 次内核启动所带来的约 1.1 ms 的开销。
• 可扩展性： MPK 在 2、4 和 8 个 GPU 上均表现出良好的扩展性，并始终保持相对于基准系统的加速。

5. 重要性与主张

论文声称 MPK 通过将执行流水线统一到单个 mega-kernel 中，代表了 GPU 推理范式的转变。其重要性在于：

• 打破内核屏障： 通过将同步移动到 SM 级别，MPK 解锁了目前在逐算子内核系统中无法实现的硬件利用潜力。
• 自动化 vs. 手动工程： 不同于以往需要专家手动设计的 mega-kernel 方法，MPK 实现了过程的自动化，使得标准 PyTorch 开发者只需极少的代码修改即可获得高性能的 mega-kernel 执行能力。
• 贴近硬件： 该系统通过最小化软件开销（启动、调度、同步）并最大化硬件重叠，将 LLM 推理性能推向了底层硬件的理论极限。

作者总结道，MPK 为构建高性能、编译器驱动的推理系统开辟了新路径，在保留现有编程模型灵活性的同时，提供了接近硬件极限的性能。