Ada-MK：通过基于有向无环图的自动搜索实现大语言模型推理的自适应 MegaKernel 优化
计算与语言
2026-05-13 v1
摘要
当大语言模型 (LLM) 在商业在线广告系统中提供实时推理服务时，端到端延迟必须严格限制在毫秒级。然而，解码阶段生成的每个 token 都会触发数千次内核启动，仅内核启动开销就可能占端到端推理时间的 14.6%。MegaKernel 通过将多个算子融合到一个持久内核中，消除了启动开销和算子间的高带宽内存 (HBM) 往返。然而，现有的 MegaKernel 实现在资源受限的 GPU（如 NVIDIA Ada）上面临可移植性与效率之间的根本矛盾：手工调优的方案与特定架构紧密耦合，缺乏可移植性；而自动编译的方法引入了运行时动态调度，其分支惩罚在延迟关键场景中是不可接受的。我们观察到，在固定的部署配置下，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 在商业在线广告系统中的首次工业部署。
引用
@article{arxiv.2605.11581,
title = {Ada-MK: Adaptive MegaKernel Optimization via Automated DAG-based Search for LLM Inference},
author = {Wenxin Dong and Mingqing Hu and Guanghui Yu and Qiang Fu and Peng Xu and Hui Xu and Yue Xing and Xuewu Jiao and Shuanglong Li and Lin Liu},
journal= {arXiv preprint arXiv:2605.11581},
year = {2026}
}
备注
10 pages, 8 figures