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

当大语言模型（LLM）在商业在线广告系统中提供实时推理服务时，端到端延迟必须严格限制在毫秒范围内。然而，解码阶段生成的每个词元都会触发数千次内核启动，仅内核启动开销就可以占端到端推理时间的 14.6%。 MegaKernel 通过将多个运算符融合到单个持久内核中，消除了启动开销和运算符间 HBM 往返。然而，现有的 MegaKernel 实现在 NVIDIA Ada 等资源受限的 GPU 上面临可移植性和效率之间的根本矛盾：手动调整的解决方案与特定架构紧密耦合，缺乏可移植性，而自动编译方法引入了运行时动态调度，其分支惩罚在延迟关键设置中是不可接受的。我们观察到，在固定的部署配置下，MegaKernel 的最佳执行路径是唯一确定的，并且运行时动态决策可以完全提升到编译时。基于这一见解，我们提出了 Ada-MK：（1）结合 K 维分割的三维共享内存约束模型，可将峰值共享内存使用量减少 50%； （2）基于MLIR的细粒度DAG离线搜索，固化最优执行路径，完全消除运行时分支； (3) 异构混合推理引擎，将 MegaKernel 作为插件嵌入到 TensorRT-LLM 中，将高吞吐量预填充与低延迟解码相结合。在 NVIDIA L20 上，Ada-MK 的单批次吞吐量比普通 TensorRT-LLM 提高了 23.6%，比 vLLM 提高了 50.2%，在所有测试场景中都取得了积极的成果——这是 MegaKernel 在商业在线广告系统中的首次工业部署。