arXiv:2605.11581cs.CL2026-05被引 1

通过编译时搜索优化大模型推理,显著降低延迟并提升吞吐量。

Ada-MK: Adaptive MegaKernel Optimization via Automated DAG-based Search for LLM Inference

论文配图:Ada-MK: Adaptive MegaKernel Optimization via Automated DAG-based Search for LLM Inference
图 1 · 摘自论文原文
  • 基于MLIR的有向无环图搜索,将运行时分支提前到编译阶段。
  • 共享内存优化使峰值使用量减少50%,提升资源利用率。
  • 在广告系统中首次实现工业级大核优化,全场景性能提升超20%。

当大语言模型在商业在线广告系统中进行实时推理时,端到端延迟必须控制在毫秒级别。然而,解码阶段每个生成的令牌会触发数千次内核调用,仅内核启动开销就占端到端推理时间的14.6%。大核(MegaKernel)通过融合多个算子为单一持久内核,消除启动开销和算子间HBM往返。但现有实现面临可移植性与效率的矛盾:手工调优方案紧耦合特定架构,缺乏可移植性;自动编译方法引入运行时动态调度,分支惩罚在低延迟场景不可接受。我们发现,在固定部署配置下,大核的最佳执行路径是唯一的,运行时动态决策可完全移至编译时。基于此,提出Ada-MK:(1) 三维共享内存约束模型结合K维拆分,将峰值共享内存使用量降低50%;(2) 基于MLIR的细粒度有向无环图离线搜索,固化最优执行路径,彻底消除运行时分支;(3) 异构混合推理引擎,将大核作为插件嵌入TensorRT-LLM,融合高吞吐预填充与低延迟解码。在NVIDIA L20上,Ada-MK相比原生TensorRT-LLM单批吞吐提升最高23.6%,相比vLLM提升50.2%,所有测试场景均获正向收益——首次在商业在线广告系统中实现大核工业级部署。

原文摘要 · Abstract (English)

When large language models (LLMs) serve real-time inference in commercial online advertising systems, end-to-end latency must be strictly bounded to the millisecond range. Yet every token generated during the decode phase triggers thousands of kernel launches, and kernel launch overhead alone can account for 14.6% of end-to-end inference time. MegaKernel eliminates launch overhead and inter-operator HBM round-trips by fusing multiple operators into a single persistent kernel. However, existing MegaKernel implementations face a fundamental tension between portability and efficiency on resource-constrained GPUs such as NVIDIA Ada: hand-tuned solutions are tightly coupled to specific architectures and lack portability, while auto-compiled approaches introduce runtime dynamic scheduling whose branch penalties are unacceptable in latency-critical settings. We observe that under a fixed deployment configuration, the optimal execution path of a MegaKernel is uniquely determined, and runtime dynamic decision-making can be entirely hoisted to compile time. Building on this insight, we propose Ada-MK: (1) a three-dimensional shared-memory constraint model combined with K-dimension splitting that reduces peak shared memory usage by 50%; (2) MLIR-based fine-grained DAG offline search that solidifies the optimal execution path, completely eliminating runtime branching; and (3) a heterogeneous hybrid inference engine that embeds MegaKernel as a plugin into TensorRT-LLM, combining high-throughput Prefill with low-latency Decode. On an NVIDIA L20, Ada-MK improves single-batch throughput by up to 23.6% over vanilla TensorRT-LLM and 50.2% over vLLM, achieving positive gains across all tested scenarios--the first industrial deployment of MegaKernel in a commercial online advertising system.

大模型推理内核优化低延迟MLIR

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。