像人类一样优化CUDA：微剖析工具作为基于LLM的GPU内核优化的专家替代

# 像人类一样优化 CUDA：微性能剖析工具作为基于 LLM 的 GPU 内核优化的专家代理 来源：https://arxiv.org/html/2606.26453 Jiading Gai\*jiadingg@amazon\.com AmazonShuai Zhang\*shuaizs@amazon\.com AmazonKaj Bostrombostromk@amazon\.com AmazonJin Huangjinhun@amazon\.com AmazonVihang Patilpvihang@amazon\.com AmazonHaoyang Fanghaoyfang@amazon\.com AmazonBernie Wangyuyawang@amazon\.com AmazonHuzefa Rangwala†rhuzefa@gmail\.com SiemensGeorge Karypis†karypis@umn\.edu University of Minnesota ###### 摘要 我们提出了 KernelPro，一个闭环多智能体系统，它通过将大语言模型 (LLM) 代码生成与硬件性能分析器反馈以及可插拔的瓶颈检测工具集成，自动生成、性能分析并迭代优化 GPU 内核代码。KernelPro 引入四项贡献：(1) 一个*语义反馈算子*，将专家启发式方法编码为可插拔的微性能剖析工具，将原始硬件指标转化为可操作的自然语言指导；(2) 一个*两阶段工具调用架构*，其中基于屋顶线的瓶颈分类筛选出要执行的特定分析工具，结合了内核级别 (ncu)、指令级别 (SASS) 和系统级别 (nsys) 性能分析；(3) 一个*领域自适应的 MCTS*，具有渐进式扩展、非对称分支、对数奖励校准、死胡同剪枝以及用于跨迭代学习的搜索记忆；(4) 通过*直接生成 CuTe 源码级代码*，即在 CUTLASS/CuTe 代码库上进行自主代码搜索，模拟专家工程师使用原始 CUDA+CuTe 编写高性能 GPU 内核的方式。在 KernelBench 上，KernelPro 在 Level 1/2/3 上分别实现了 2.42×/4.69×/5.30× 的几何平均加速，在所有难度级别上均达到了最先进 (SOTA) 性能。在 VeOmni 专家优化的 MoE 训练内核上，KernelPro 通过从头生成原始 CUDA+CuTe Hopper WGMMA 内核，实现了相对于手工调优 Triton 的 1.23× 加速。消融研究表明，每个设计组件都独立且显著地提高了优化质量：微性能剖析工具 (与原始指标相比 p<0.0001)、MCTS 搜索 (与贪心相比几何平均高 26%，p=0.004) 以及主动工具编排 (改进 23%，p=0.035)。最后，KernelPro 是首个*优化能效*的 CUDA 内核编码智能体，超越了先前系统仅关注速度的局限，在速度匹配的情况下实现了 11.6% 的测量能耗降低。 11footnotetext:\*同等贡献。代码将在发表后发布。 ## 1 引言 优化 GPU 内核需要横跨硬件架构、内存层次结构、指令调度和供应商特定编程模型的专业知识。经验丰富的工程师遵循一个严格的周期：对内核进行分析，识别指标中的瓶颈模式，诊断根本原因，应用有针对性的转换，然后重新分析以验证。每次成功优化后瓶颈都会转移，这个过程会重复进行，直至达到硬件极限。专家之所以有效，并非因为他们能读取性能分析器的数字，而是因为他们对数字运用了*模式识别*和*诊断推理*——这些启发式方法经过多年积累，常常被编纂到个人脚本和工具中 (NVIDIA Corporation, 2024 (https://arxiv.org/html/2606.26453#bib.bib24); Yang 等人, 2020 (https://arxiv.org/html/2606.26453#bib.bib33))。 最近的基于 LLM 的 CUDA 优化系统 (Zhang 等人, 2025 (https://arxiv.org/html/2606.26453#bib.bib37); Dong 等人, 2026 (https://arxiv.org/html/2606.26453#bib.bib8); Li 等人, 2026 (https://arxiv.org/html/2606.26453#bib.bib20); Chen 等人, 2025 (https://arxiv.org/html/2606.26453#bib.bib7)) 已经表明，当给予性能分析数据时，语言模型可以生成有竞争力的内核代码。然而，这些系统将原始或轻微总结的指标直接呈现给 LLM，并依赖模型隐式地复制专家推理。这混淆了两种不同的能力：*解释硬件遥测数据*（一项结构化的、遵循规则的任务）和*生成优化代码*（一项创造性的、依赖于上下文的的任务）。通过要求 LLM 同时完成这两项任务，现有智能体放弃了使人类专家可靠的那种系统性中间分析——它们会错过优化机会，产生不一致的诊断，并且无法在不重新训练的情况下轻易地融入新的分析模式。 我们提出了 KernelPro，一个闭环多智能体系统，它通过将专家启发式方法编码为*可插拔的微性能分析工具*——即在代码生成 LLM 看到之前，将原始硬件指标转换为可操作的自然语言指导的可执行分析函数——来分离这些关注点。每个工具实现一个特定的诊断模式：一个触发条件 (例如，张量核心利用率低于 10%)、分析逻辑 (例如，检查内核是否执行可约的矩阵操作) 以及一个指示性的推荐 (例如，“为获得主要的预期加速，使用 CUTLASS GEMM 重写”)。一个两阶段调用流水线首先通过屋顶线分析对内核的瓶颈类型进行分类，然后只调度相关的工具——从而减少提示噪音，并将 LLM 聚焦于相关的优化。 KernelPro 还贡献了一种领域自适应的蒙特卡洛树搜索 (MCTS) 策略，该策略在利用有希望的优化路径与探索替代方法之间取得平衡。与先前基于 LLM 的 CUDA 优化系统使用的平面迭代优化不同，采用上置信界 (UCT) 选择的 MCTS 避免了在复杂的多步优化问题上过早收敛到局部最优。结合一个在迭代间提炼发现的搜索记忆机制，KernelPro 在固定计算预算下实现了单调递增的性能曲线。 在 KernelBench 基准测试 (Ouyang 等人, 2025 (https://arxiv.org/html/2606.26453#bib.bib26)) 上，KernelPro 在 Level 1/2/3 上分别实现了 2.42×/4.69×/5.30× 的几何平均加速，在所有级别上都超越了当前 SOTA 的 KernelBlaster (Ding 等人, 2026 (https://arxiv.org/html/2606.26453#bib.bib8))。遵循设计选择应基于经验而非假设的原则，我们通过受控的消融实验（每项都使用配对 Wilcoxon 符号秩检验 (Hollander 等人, 2014 (https://arxiv.org/html/2606.26453#bib.bib10)) 评估其显著性）在固定的 42 任务 KernelBench 子集（Level 1–3，因与 transformer 相关性而被选中，并在本文中全程使用 (第 5.3 节)）上验证了 KernelPro 的每个组件。例如，微性能剖析工具提供的加速比原始指标高 125%（p<0.0001），MCTS 显著优于贪心搜索（Wilcoxon p=0.004，几何平均高 26%）。总体而言，KernelPro 贡献并综合了五个想法，以使基于 LLM 的优化在现代 GPU 架构上更加扎实、可解释和有效： 1. 1. 一个**语义反馈算子**，将专家启发式方法形式化为可插拔的、可被 LLM 调用的微性能剖析工具，将硬件指标转化为自然语言优化指导。 2. 2. 一个**两阶段工具调用架构**，其中基于屋顶线的瓶颈分类（第一阶段）筛选出要执行的特定分析工具（第二阶段），结合了内核级、指令级和系统级性能分析。 3. 3. 用于基于 LLM 的 CUDA 优化的**领域自适应 MCTS**，具有渐进式扩展、非对称分支、对数奖励校准和死胡同剪枝。搜索记忆实现了跨迭代学习。 4. 4. **直接 CuTe 源码级代码生成**，通过在 CUTLASS/CuTe 代码库 (NVIDIA, 2023 (https://arxiv.org/html/2606.26453#bib.bib23))（CUTLASS 是一个高性能内核的 CUDA 模板库；CuTe 是其底层的张量/布局代数层，用于从头组合内核）上进行自主代码搜索，模拟了专家工程师使用原始 CUDA+CuTe 编写高性能 GPU 内核的方式（完整生成源码见附录 H）。 5. 5. 首个**能效感知的 CUDA 内核编码智能体**：先前基于 LLM 的内核智能体仅针对速度进行优化，而 KernelPro 通过字典序能效奖励，将其奖励和性能分析工具扩展到*能效*作为次要目标。作为初步研究，一个匹配速度的 A/B 测试发现，通过选择低能耗指令，在相同速度下测量能耗降低了 11.6%（第 5.5 节）；全面的评估留待未来工作。 ## 2 相关工作 有几个近期系统将 LLM 应用于 CUDA 内核优化，主要区别在于硬件性能分析数据如何到达 LLM。一些系统仅依赖端到端计时，无需硬件性能分析 (Andrews & Witteveen, 2025 (https://arxiv.org/html/2606.26453#bib.bib1); Zhang 等人, 2026 (https://arxiv.org/html/2606.26453#bib.bib38); Wiedemann 等人, 2026 (https://arxiv.org/html/2606.26453#bib.bib31); Chen 等人, 2026 (https://arxiv.org/html/2606.26453#bib.bib6))。其他系统则直接将原始或经过统计过滤的 ncu 指标传递到 LLM 的提示中，并依赖 LLM 自身进行解释 (Zhang 等人, 2025 (https://arxiv.org/html/2606.26453#bib.bib37); Dong 等人, 2026 (https://arxiv.org/html/2606.26453#bib.bib8); Li 等人, 2026 (https://arxiv.org/html/2606.26453#bib.bib20); Team & Meta, 2025 (https://arxiv.org/html/2606.26453#bib.bib29); Han 等人, 2026 (https://arxiv.org/html/2606.26453#bib.bib9))。cuPilot (Chen 等人, 2025 (https://arxiv.org/html/2606.26453#bib.bib7)) 增加了基于屋顶线的瓶颈分类，但止于给出一个标签（内存受限 vs. 计算受限），而未能从底层计数器生成可操作的指令。在所有情况下，LLM 仍然是性能分析数据的唯一解释者——没有先前的基于 LLM 的系统提供对硬件计数器到优化指令的程序性转换。 这一挑战在 LLM 出现之前就已存在：经典的性能顾问如 GPA (Zhou 等人, 2021 (https://arxiv.org/html/2606.26453#bib.bib39)) 使用指令采样和数据流分析将停顿归因于其根本原因，并发出人类可读的优化建议，并且独立地观察到原始性能分析器输出“对如何改进代码提供的洞察很少”。KernelPro 采用了相同的“先诊断后开方”的理念，但*闭环*了：其微性能剖析工具在 LLM 介入*之前*就将计数器转化为自然语言的指令，然后 LLM 直接应用修复，而不是留给人类工程师。我们的消融实验（附录 C）证实，这种基于工具的解释显著优于无反馈和仅原始指标的反馈。 CUTLASS 和 CuTe 代码生成。有多个智能体针对 NVIDIA 的 CUTLASS 库，但处于不同的抽象级别。AVO (Chen 等人, 2026 (https://arxiv.org/html/2606.26453#bib.bib6)) 和 StitchCUDA (Li 等人, 2026 (https://arxiv.org/html/2606.26453#bib.bib20)) 生成高级别的 CUTLASS 模板实例化——从预定义的配置空间中选择 tile 大小、流水线阶段和 epilogue 融合。KernelPro 在更低的级别上操作：它通过在 CUTLASS/CuTe 代码库中自主搜索布局代数、复制原子和 MMA 原子，然后将它们组合成新颖的内核，生成原始的 CUDA+CuTe 源代码——这与专家内核工程师使用的流程相同。 用于 LLM 代码优化的树搜索和进化方法。越来越多的研究工作将蒙特卡洛树搜索与 LLM 集成。MCTSr (Zhang 等人, 2024 (https://arxiv.org/html/2606.26453#bib.bib36)) 将 UCT 应用于数学奥林匹克问题，其中每个节点代表一个通过自我评估改进的答案版本。DeepSearch (Wu 等人, 2025 (https://arxiv.org/html/2606.26453#bib.bib32)) 将 MCTS 嵌入到 RL 训练循环中，并采用全局前沿选择，目标是模型权重更新而非推理时优化。TreeRL (Hou 等人, 2025 (https://arxiv.org/html/2606.26453#bib.bib13)) 提出了熵引导的令牌级分支，从高不确定性的令牌处进行分支。Tree-GRPO (Ji 等人, 2025 (https://arxiv.org/html/2606.26453#bib.bib15)) 将完整的智能体步骤作为树节点，并从树内相对优势中推导出过程监督信号。这些系统将树搜索视为*训练*LLM 或生成*推理轨迹*的机制。AlphaEvolve (Novikov 等人, 2025 (https://arxiv.org/html/2606.26453#bib.bib22)) 采用进化方法，使用基于种群的 LLM 采样和标量适应度来优化包括 TPU 内核在内的算法，但需要数百万次采样和数天的计算。 KernelPro 引入了用于基于 LLM 的 GPU 内核优化的领域自适应 MCTS，其中每个节点代表一个完整的、已编译且已分析的 CUDA 内核。这种粗粒度需要不同的设计选择：渐进式扩展 (Auger 等人, 2013 (https://arxiv.org/html/2606.26453#bib.bib2)) 限制了扩展速度，因为每个节点都需要编译、执行和性能分析，确保探索随访问次数呈次线性增长；对数奖励校准消除了对学习价值函数的需求；非对称分支因子反映了实证观察，即有效优化多于有效修复；死胡同剪枝移除已耗尽的子树；搜索记忆实现了跨迭代学习。KernelPro 采用了 MCTSr 的根节点重新扩展策略（当现有子树停滞时注入新的种子），但用实测加速替换了自我奖励，消除了奖励操纵的风险。在撰写本文时，我们观察到 OptiML (Bhattacharjee 等人, 2026 (https://arxiv.org/html/2606.26453#bib.bib4)) 也独立地将 MCTS 与 UCT 应用于 CUDA 内核优化，我们认为这提供了一个趋同的证据，表明树搜索非常适合这个领域。 ## 3 KernelPro 的架构 ### 3.1 高级流水线 KernelPro 作为一个多智能体系统运行，由三个主要智能体通过一个搜索编排器协调。 **基准测试智能体**（第一阶段）分析参考实现，以建立基线性能指标，并通过屋顶线分析对内核的瓶颈类型（计算受限、内存受限、延迟受限或混合）进行分类。**程序员智能体**以三种模式生成优化的 CUDA 代码：初始生成、调试编译/正确性错误，以及基于性能分析反馈的迭代优化。**性能分析智能体**（第二阶段）执行由第一阶段瓶颈分类筛选过的微性能分析工具，生成指导下一次优化迭代的语义反馈。这种两阶段架构——先进行瓶颈分类，再执行筛选后的工具——减少了提示噪音，并将 LLM 聚焦于相关的优化。 第一阶段（一次性）迭代循环第二阶段微性能分析工具CUDA/Triton/PyTorch内核基准测试智能体（屋顶线分类）瓶颈类型搜索编排器MCTS / 贪心程序员智能体初始化 | 调试 | 优化N个候选最佳验证CUDA内核搜索记忆跨迭代学习编译并执行验证正确性性能分析智能体（语义反馈）ncu（8个工具）nsys（5个工具）SASS（2个工具）输入内核性能分析选择节点收敛CUDA内核代码瓶颈类型筛选注入记忆化记忆化语义反馈ncu指标 + .cubin控制/主流程语义反馈内存流工具调用瓶颈筛选 图1：KernelPro 智能体优化工作流。第一阶段（基准测试智能体）执行一次性的基于屋顶线的瓶颈分类。迭代循环包括搜索、编程、编译、正确性验证、性能分析和语义反馈。