Ada-MK：基于自动化 DAG 搜索的 LLM 推理自适应 MegaKernel 优化

# Ada-MK：通过基于 DAG 的自动化搜索实现面向 LLM 推理的自适应 MegaKernel 优化

来源：https://arxiv.org/html/2605.11581
Wenxin Dong, Mingqing Hu, Guanghui Yu, Qiang Fu, Peng Xu, Hui Xu, Yue Xing, Xuewu Jiao†, Shuanglong Li, Lin Liu

###### 摘要

当大语言模型（LLMs）在商业在线广告系统中提供实时推理服务时，端到端延迟必须严格控制在毫秒级范围内。然而，解码阶段生成的每个 Token 都会触发数千次内核启动，仅内核启动开销就可能占端到端推理时间的 14.6%。MegaKernel 通过将多个算子融合为单个持久化内核，消除了启动开销和算子间的高带宽存储器（HBM）往返延迟。然而，现有的 MegaKernel 实现在 NVIDIA Ada 等资源受限的 GPU 上面临着可移植性与效率之间的根本矛盾：手工调优的方案与特定架构紧密耦合，缺乏可移植性；而自动编译方法引入了运行时动态调度，其分支惩罚在延迟敏感的场景中是不可接受的。

我们观察到，在固定的部署配置下，MegaKernel 的最佳执行路径是唯一确定的，运行时的动态决策可以完全提升到编译时进行。基于这一洞察，我们提出了 Ada-MK：（1）结合 K 维分割的三维共享内存约束模型，将峰值共享内存使用量降低了 50%；（2）基于 MLIR 的细粒度 DAG 离线搜索，固化最佳执行路径，完全消除运行时分支；以及（3）一种异构混合推理引擎，将 MegaKernel 作为插件嵌入 TensorRT-LLM，结合高吞吐量的 Prefill（预填充）和低延迟的 Decode（解码）。在 NVIDIA L20 上，Ada-MK 相比原生 TensorRT-LLM 的单批吞吐量提高了高达 23.6%，相比 vLLM 提高了 50.2%，在所有测试场景中均取得了正向增益——这是 MegaKernel 首次在商业在线广告系统中实现工业级部署。

MegaKernel，基于 DAG 的优化，自适应优化，LLM 推理，共享内存复用，商业广告

## 1\.引言

大语言模型（LLMs）越来越多地部署在搜索、推荐和对话应用程序中，参数量从数十亿增长到数千亿。因此，推理延迟已成为工业应用的关键瓶颈，尤其是在对延迟敏感的在线服务中，在有限的硬件资源下进行高效推理至关重要。对于极端延迟敏感的场景，GPU 上频繁的全球内存访问和内核启动开销已成为端到端性能的主要瓶颈（Kwon et al., 2023 (https://arxiv.org/html/2605.11581#bib.bib14); NVIDIA, 2024 (https://arxiv.org/html/2605.11581#bib.bib22)）。在 TensorRT-LLM 下使用 Nsight Systems 对 Qwen2.5-1.5B 进行性能剖析显示，内核启动开销约占端到端推理时间的 14.6%（1,655,550 次启动消耗约 3.3 秒）。在传统流水线中，相邻内核必须通过 HBM 交换中间结果；而 MegaKernel 则利用共享内存和寄存器实现无缝的算子链式处理和深度并行，从根本上消除了算子切换开销和 HBM 往返延迟。MegaKernel 通过持久化内核实现单次启动的持久计算，通过 Warp 专业化构建生产者-消费者流水线，并通过 TMA 和异步 I/O 重叠计算与内存访问（HazyResearch, 2025 (https://arxiv.org/html/2605.11581#bib.bib10)）。

我们的目标部署环境是运行在 NVIDIA Ada (L20) GPU 上的商业在线广告系统（NVIDIA, 2022a (https://arxiv.org/html/2605.11581#bib.bib19)）。在 Ada 上部署 MegaKernel 面临严重的双重限制。首先，业务要求端到端延迟确定性地在 1-5 毫秒以内。其次，Ada 架构原生缺乏 TMA 硬件支持，需要使用 PTX 汇编和手工制作的软件流水线来模拟异步数据移动；其片上共享内存仅为 H100 的一半（128 KB 对比 227 KB）（NVIDIA, 2022b (https://arxiv.org/html/2605.11581#bib.bib20)），严重压缩了流水线阶段和 Tile 大小的优化空间。相比之下，Hopper/Blackwell 架构提供 TMA 硬件支持和更大的共享内存（227 KB），使得 MegaKernel 的部署更加直接（HazyResearch, 2025 (https://arxiv.org/html/2605.11581#bib.bib10); NVIDIA, 2022b (https://arxiv.org/html/2605.11581#bib.bib20)）。在实践中，Ada 上的共享内存限制使得可达到的流水线阶段从理论最优的 4 阶段减少到仅 2 阶段，导致流水线占空比损失超过 30%。

现有解决方案无法满足这些要求。斯坦福的 MegaKernel（HazyResearch, 2025 (https://arxiv.org/html/2605.11581#bib.bib10)）提供了强大的性能，但其代码库深度依赖于 Hopper/Blackwell 特定的汇编优化，且仅支持少数模型架构（如 Llama-1B），缺乏对 Qwen 等其他广泛使用模型的支持（Qwen Team, 2024 (https://arxiv.org/html/2605.11581#bib.bib24), 2025 (https://arxiv.org/html/2605.11581#bib.bib25)）；它也不支持长序列和大批量的 Prefill 阶段。Mirage MPK（Cheng et al., 2025 (https://arxiv.org/html/2605.11581#bib.bib2); Wu et al., 2024 (https://arxiv.org/html/2605.11581#bib.bib28)）通过自动调优提高了易用性，但其托管指针机制引入了基于共享内存页面状态的运行时 if-else 分支，降低了超低延迟场景下的指令发射效率，且无法达到手工调优算子的性能。

为了解决这些挑战，我们提出了 Ada-MK，具有以下核心贡献：

- **自适应共享内存管理**。我们从硬件规格、模型架构和动态工作负载三个维度对共享内存分配进行建模，通过 K 维细粒度分割将峰值共享内存需求降低 50%，同时启用跨算子页面复用，在 Ada 受限的资源上重建高效流水线。
- **基于细粒度 DAG 的自动搜索**。我们利用 MLIR 降阶（Lattner et al., 2021 (https://arxiv.org/html/2605.11581#bib.bib16)）构建 PTX 级别的依赖 DAG，并通过离线性能剖析固化最佳执行轨迹，完全消除运行时动态决策开销。与 Ansor（Zheng et al., 2020 (https://arxiv.org/html/2605.11581#bib.bib32)）和其他传统自动调优框架不同，Ada-MK 的 DAG 级搜索能够捕捉更细粒度的并行机会。
- **异构混合推理引擎**。我们将 MegaKernel 作为插件嵌入 TensorRT-LLM（NVIDIA, 2024 (https://arxiv.org/html/2605.11581#bib.bib22)），复用 TensorRT-LLM 的原生算子进行 Prefill，并切换到 MegaKernel 引擎进行 Decode，实现高吞吐量和低延迟——这是 MegaKernel 的首次工业级部署。

实验表明，Ada-MK 在 Qwen 模型系列上相比 vLLM（Kwon et al., 2023 (https://arxiv.org/html/2605.11581#bib.bib14)）、SGLang（Zheng et al., 2025b (https://arxiv.org/html/2605.11581#bib.bib34)）和原生 TensorRT-LLM（NVIDIA, 2024 (https://arxiv.org/html/2605.11581#bib.bib22)）将端到端延迟降低了 10%-50%，并已在百度的商业在线广告系统中投入生产使用。本文其余部分组织如下：§2 (https://arxiv.org/html/2605.11581#S2) 介绍背景和相关工作；§3 (https://arxiv.org/html/2605.11581#S3) 描述 Ada-MK 的整体架构；§4 (https://arxiv.org/html/2605.11581#S4) 详细介绍三项核心优化；§5 (https://arxiv.org/html/2605.11581#S5) 展示实验评估；§6 (https://arxiv.org/html/2605.11581#S6) 总结并展望未来方向。

## 2\.背景和相关工作

参见标题 Figure 1. Ada-MK 整体架构。第一阶段（离线 MegaKernel 合成）：Transformer Decoder 和 LM Head 被解析为细粒度 DAG，通过硬件资源约束进行剪枝，并通过性能剖析选择最佳执行轨迹，然后进行序列化。第二阶段（在线 Warp 专业化执行）：序列化的 MegaKernel 作为 TensorRT-LLM 插件嵌入，Warp 专业化角色（Loader、Consumer、Storer、Controller、Launcher）通过共享内存协作，形成高效的软件流水线。### 2.1\.LLM GPU 内核技术：从局部到全局融合

GPU 算子优化已从局部融合发展到全局融合，形成了几个不同的研究方向：

传统内核融合将相邻内核合并以减少启动次数，但仍依赖全局内存进行中间数据交换，在极端模型规模下难以规避内存带宽和数据依赖瓶颈。CUDAGraph 通过捕获和重放内核序列来减少启动开销，但无法实现深层的算子间融合。

编译器级别的算子融合。Triton（Tillet et al., 2019 (https://arxiv.org/html/2605.11581#bib.bib26)）通过块级编程抽象实现编译时算子融合，而 Hidet（Ding et al., 2023 (https://arxiv.org/html/2605.11581#bib.bib7)）通过任务映射范式提供更细粒度的调度控制。这些方法提高了融合深度，但仍基于“每算子一内核”的执行模型，无法消除算子切换开销。此外，MCFuser（Zhang et al., 2024 (https://arxiv.org/html/2605.11581#bib.bib30)）实现了内存受限计算密集型算子链的高性能融合，Neptune（Zhao et al., 2025 (https://arxiv.org/html/2605.11581#bib.bib31)）通过打破循环依赖实现了归约算子的高级融合，Magneto（Di et al., 2025 (https://arxiv.org/html/2605.11581#bib.bib5)）专注于并行算子结构的协调优化。

单算子极致融合。FlashAttention（Dao et al., 2022 (https://arxiv.org/html/2605.11581#bib.bib4)）通过 IO 感知的分块技术将注意力的多个步骤融合为单个内核，实现了突破性的内存效率；FlashAttention-2（Dao, 2024 (https://arxiv.org/html/2605.11581#bib.bib3)）进一步优化了并行性和工作分区。然而，这些方法仅限于单个算子。

MegaKernel 全局融合（HazyResearch, 2025 (https://arxiv.org/html/2605.11581#bib.bib10)）将整个计算块集成到一个持久化内核中，通过 Warp 专业化将 SM 内的 Warp 划分为生产者-消费者角色，并通过多级存储和异步复制实现流水线并行，彻底消除了算子切换开销并突破了全局内存访问限制。MegaKernel 代表了大内核方向的最前沿进展，实现了全链路融合和持久计算。

### 2.2\.主流方法分析

斯坦福的 MegaKernel 建立在 Tiny-Llama 和 FlashAttention（Dao et al., 2022 (https://arxiv.org/html/2605.11581#bib.bib4)）原语之上。它将 SM 资源分解为五种异构角色——Loader（异步预取）、Consumer（张量计算）、Storer（异步写回）、Controller（指令分发协调）和 Launcher（资源释放管理），利用异步协作和信号量同步实现计算、内存访问和指令分发的深度流水线解耦。其实现大量借鉴了 NVIDIA CUTLASS（NVIDIA, 2023 (https://arxiv.org/html/2605.11581#bib.bib21)）的 Warp 专业化 GEMM 设计。该框架专注于 Hopper/Blackwell 硬件，并为特定模型（如 Llama-1B）硬编码了 Tile 大小，缺乏对 Qwen 变体的支持。代码库包含大量的架构特定汇编优化，移植到 Ada 的成本极高。初始版本仅支持低负载的 Decoder 阶段，不支持长序列和大批量的 Prefill，在工业部署方面存在显著差距。

Mirage MPK（Cheng et al., 2025 (https://arxiv.org/html/2605.11581#bib.bib2)）是算子编译领域的代表性近期工作，提出了基于 tGraph 的多级流水线抽象，自动化搜索分块策略和存储分配。Mirage（Wu et al., 2024 (https://arxiv.org/html/2605.11581#bib.bib28)）进一步通过 μGraph 统一实现了跨层级优化，发现了结合代数变换、调度变换和自定义内核生成的新颖优化。OLLIE（Zheng et al., 2022 (https://arxiv.org/html/2605.11581#bib.bib33)）通过基于推导的变换扩展了张量代数表达式的搜索空间，Korch（Hu et al., 2024 (https://arxiv.org/html/2605.11581#bib.bib12)）通过算子裂变和约束优化实现了最佳内核编排，PET（Wang et al., 2023 (https://arxiv.org/html/2605.11581#bib.bib27)）通过部分等价变换发现了传统方法不可见的优化机会，TASO（Jia et al., 2019 (https://arxiv.org/html/2605.11581#bib.bib13)）实现了计算图替换生成和验证的自动化。这些工作从不同角度推进了张量程序优化，但均未解决 MegaKernel 级别的全局算子融合问题。MPK 的托管指针机制在运行时动态确定执行路径，由此产生的 if-else 分支降低了指令发射效率，使其无法达到硬件峰值性能。其自动化搜索也不容易在合理时间内收敛到不规则 Tile 或复杂算子链的全局最优解。

## 3\.Ada-MK 整体架构

如图 1 (https://arxiv.org/html/2605.11581#S2.F1) 所示，Ada-MK 系统架构由两个紧密耦合的阶段组成：*离线 MegaKernel 合成*（第一阶段）和*在线 Warp 专业化执行*（第二阶段）。

第一阶段：离线 MegaKernel 合成。此阶段负责计算图生成和受硬件约束的自动搜索。系统首先将 Transformer Decoder 和 LM Head 的逻辑结构解析为细粒度 DAG。硬件资源过滤器随后评估共享内存容量和寄存器限制等约束，直接剪除超出资源预算的无效分支。在由此产生的可行搜索空间内，系统进行性能剖析以识别目标 GPU 的最佳 DAG 执行轨迹，并将该轨迹序列化以供运行时调用。

第二阶段：在线 Warp 专业化执行。此阶段将序列化的 MegaKernel 作为插件嵌入 TensorRT-LLM 推理框架，实现 Prefill（使用 TensorRT-LLM 的原生算子）和 Decode（使用 MegaKernel 引擎）之间的无缝切换。在每个流多处理器（SM）中，MegaKernel 构建了一个高效的*warp 专业化软件流水线*：计算资源在空间上划分为多个 Warp 组——Loader、Consumer、Storer、Controller 和 Launcher，它们共享 SM 资源并通过共享内存协作，最大限度地重叠内存访问延迟与计算。

## 4\.核心优化

### 4.1\.资源约束下的自适应共享内存管理

#### 4.1.1\.多维度参数感知资源建模

我们从硬件规格、模型架构和动态工作负载三个维度对共享内存分配进行建模，实现细粒度共享内存管