KernelOPT：面向 GPU 内核优化的调度感知智能体搜索

arXiv:2609.30059v2 公告类型：替换并交叉列出 摘要：深度学习推理和训练的性能在很大程度上取决于 GPU 内核的效率。PyTorch Inductor 等现代编译器能够从高层模型代码自动生成 GPU 内核，但其性能往往远逊于专家编写的实现。近期基于大语言模型（LLM）辅助的内核优化器能够缩小独立内核上的这一差距，但它们将编译后的模型视为黑箱，通常仅优化单个独立内核，既不遵循编译器的结构性决策，也不对模型进行端到端验证。我们提出 KernelOPT，这是一个将编译后的模型视为结构化产物的多智能体系统。它保留供应商库调用（cuBLAS、cuDNN），并使用五个由性能剖析引导的 LLM 智能体，专门针对生成的 Triton 子内核进行优化。四道关卡组成的级联验证流程依次执行静态验证、多随机种子正确性检查、模型级 float64 回退验证以及性能门控（$\gamma{=}1.03$），以筛选候选方案，并对重新拼接的模型进行端到端验证。当候选方案未通过验证时，系统保留编译器基线实现。该系统支持输入 PyTorch nn Modules、独立 Triton 内核和 Helion 内核。在 NVIDIA H200 上对 250 道 KernelBench 问题（100 道 Level 1、100 道 Level 2 和 50 道 Level 3）进行评估时，KernelOPT 在所有内核（包括回退案例）上相较于 torch compile 实现的几何平均加速比分别为 1.40$\times$（L1）、1.15$\times$（L2）和 1.07$\times$（L3）。仅统计成功优化案例（不包括验证关卡决定保留编译器基线实现的案例）时，几何平均加速比显著更高：2.54$\times$（L1：36/100）、1.84$\times$（L2：23/100）和 1.37$\times$（L3：11/50），体现了优化器能够取得实质性收益的场景。