Kernel Forge：一个基于LLM的CUDA内核生成与优化智能体框架

# 用于基于大语言模型生成和优化CUDA内核的Agent框架

来源：https://arxiv.org/html/2607.24762

---

\IEEEauthorblockNJoshua Brodsky\IEEEauthorrefmark1\IEEEauthorblockA密歇根大学 美国 [email protected]\IEEEauthorblockNDhravid Kumar\IEEEauthorrefmark1\IEEEauthorblockA密歇根大学 美国 [email protected]{\@IEEEauthorhalign}\IEEEauthorblockNSavini Kashmira\IEEEauthorblockA密歇根大学 美国 [email protected]\IEEEauthorblockNJayanaka Danatanarayana\IEEEauthorblockA密歇根大学 美国 [email protected]{\@IEEEauthorhalign}\IEEEauthorblockNJason Mars\IEEEauthorblockA密歇根大学 美国 [email protected]\IEEEauthorblockNLingjia Tang\IEEEauthorblockA密歇根大学 美国 [email protected]

###### 摘要

机器学习模型正越来越多地嵌入日常软件中，其大部分运行时间都花在一小组计算内核上，如矩阵乘法、卷积和归一化。优化这些内核是降低延迟和成本最直接的方法之一，但传统上需要专家工程师手动编写底层GPU代码。基于大语言模型（LLM）的智能体系统现在可以以少得多的人力生成和优化内核，然而现有工具大多在随机生成的张量和孤立内核上进行评估，输出需要开发者手动重新集成的独立CUDA代码，主要仅针对LLM PyTorch模型，并且对检查与调试结果的支持有限。我们提出**Kernel Forge**，一个**开源**的端到端智能体框架，它接受**任何**未经修改的原生PyTorch模型。Kernel Forge支持视觉、扩散和LLM工作负载，使用蒙特卡洛树搜索（MCTS）探索多种优化路径而非单一的线性精炼链，并配备图形用户界面用于监控进度、检查候选内核以及调试失败。我们在NVIDIA DGX Spark（配备GB10 GPU）上对涵盖视觉、扩散和LLM工作负载的四个PyTorch模型评估了Kernel Forge。每个内核仅经过50次优化迭代，它便优化了14个内核，使其性能超越PyTorch eager模式：在ResNet-50上达到adaptive\_avgpool2d的**1.52倍**，在Stable Diffusion 3.5 Medium上达到group\_norm的**1.70倍**，在Gemma 4 E2B上达到softmax的**2.83倍**，在Qwen 3.5 35B-A3B上达到softmax的**1.54倍**。代码地址：Kernel Forge (https://github.com/TheJoshBrod/KernelForge)

11注：同等贡献。

## 1 引言

机器学习（ML）模型正越来越多地应用于日常软件系统[sze2017efficient]。因此，提升其运行时性能对于降低延迟和计算成本至关重要。模型的大部分运行时间都花费在一小组计算内核上，例如矩阵乘法、卷积和归一化[jouppi2017tpu,llminfgpu2025]。因此，优化这些内核是使ML系统更快、更高效的最直接方法之一。传统上，这种优化依赖于专家工程师手动编写和调优底层GPU代码[tillet2019triton]。尽管这种方法能产生非常高效的实现[dao2022flashattention]，但它速度慢、成本高，并且需要大量的专业知识。

基于大语言模型（LLM）的智能体AI系统为内核优化提供了一种新途径。这些系统能够以比传统手动调优更少的人力直接参与来生成、修订和评估候选内核。近期系统[hong2025autocomp,wang2025geak]表明，这种方法可以为单个内核带来显著的性能提升。然而，一个在孤立环境下表现良好的内核，当嵌入到已部署的ML模型中时，可能无法提供同样的优势。许多基准测试使用随机生成的张量和孤立的内核调用（而非完整的模型执行）来评估正确性和延迟[ouyang2025kernelbench,wang2026kernelbenchx,li2025tritonbench,wen2025multikernelbench,zhu2026cudabench,backendbench2025]。相比之下，在实际工作负载中，内核在模型特定条件下运行，包括张量形状、激活分布、相邻算子和内存行为，这些在孤立内核基准测试中可能并不存在。因此，我们需要一种工具，能够弥合在孤立环境下优化内核与在实际模型执行和真实工作负载中评估其行为之间的差距。

有前途的工具（包括一些开源系统）已开始着手解决这个问题[flashinferbench2026,cudaforge2025,astra2025,lange2025robustkbench,hong2025autocomp,kernelagent2025,ouyang2025kernelbench]。然而，在将内核优化应用于真实PyTorch模型方面，它们仍然留下若干挑战未解。

首先，现有工具并未使缺乏底层CUDA专业知识的开发者能够轻松进行模型级优化。将内核优化应用于真实的PyTorch工作负载需要的不仅仅是生成一个快速的独立内核：优化后的代码必须被插入到模型中并进行正确性验证。而许多当前工具将这一集成步骤留给开发者处理[cudaforge2025,astra2025,lange2025robustkbench]。其次，一些工具主要关注LLM工作负载或基准内核，而真实的PyTorch应用包括扩散模型、视觉模型和音频模型，它们使用相同的算子，但具有不同的形状、布局和参数[flashinferbench2026,ouyang2025kernelbench,astra2025]。由于内核性能依赖于这些模型特定的细节，为一种工作负载调优的优化可能无法迁移到其他工作负载。第三，现有工具可能无法有效地探索替代优化路径。线性精炼会使结果严重依赖早期设计选择[cudaforge2025,astra2025]，而束搜索（beam-style search）可能会剪掉早期表现不佳的路径，导致在启发式搜索中得到次优解[zhang1998completebeam]。这具有挑战性，因为有用的优化路径可能会经过暂时较慢的中间候选内核。最后，许多系统主要作为研究原型、脚本或基准测试工作流程而非集成的图形化工具提供。这使得内核检查、进度跟踪、候选比较、故障诊断以及应用于不同PyTorch模型的过程更加手工化[cudaforge2025,astra2025,lange2025robustkbench,kernelagent2025]。总之，这些局限性促使我们开发一个开源的端到端系统，用于在真实PyTorch工作负载中应用、检查和评估内核优化。

参见图注图1：Kernel Forge图形界面。编号的截图展示了主要工作流程：(1) 项目启动器列出了保存的优化项目；(2) 项目创建允许用户直接从Hugging Face加载模型和权重、配置目标后端并提供验证输入，无需修改模型代码；(3) 算子工作台展示了捕获的内核、它们的运行时占比及优化状态，以便用户选择要优化的对象；(4) 结果视图总结了每个算子的计时、加速比、成本以及后端或回退状态；(5) 内核工作台展示了生成的CUDA修订版本、智能体推理、优化决策以及每次优化运行的MCTS修订树。为解决这些局限性，我们引入了Kernel Forge，一个开源的智能体AI框架，用于优化超越LLM的真实PyTorch工作负载。Kernel Forge接受任何标准PyTorch模型以及用户提供的示例输入，无需用户修改模型代码。不同于输出独立优化代码的工具，Kernel Forge根据PyTorch eager模式输出验证生成的内核，并自动将其集成回PyTorch模型执行路径。它支持多样化的PyTorch工作负载，包括语言、扩散、视觉和音频模型，而非局限于单一工作负载族。为了比线性精炼或束搜索更灵活地探索优化路径，Kernel Forge使用蒙特卡洛树搜索（MCTS）[kocsis2006uct]，使其能够重新审视可能不会立即带来收益的候选内核。如图1 (https://arxiv.org/html/2607.24762#S1.F1) 所示，为提高可观察性和可用性，Kernel Forge同时提供图形用户界面（GUI）和命令行界面（CLI）与优化过程交互。用户可以轻松监控优化进度、检查候选内核、调试故障，并在不同模型提供商之间进行选择以平衡成本、性能和部署需求。

Kernel Forge的工作原理是：在提供的示例输入上运行模型，并捕获推理期间执行的PyTorch算子。然后，它挂接到诸如conv2d、matmul、group\_norm和softmax等高级算子，并使用智能体循环来提出、生成、验证和分析专门的硬件感知CUDA内核。生成的内核在捕获的工作负载上根据PyTorch eager模式的输出进行数值验证，并在eager模式模型执行中透明地运行。

我们在代表视觉、扩散和LLM工作负载的四个PyTorch模型上评估了Kernel Forge：ResNet-50[he2016resnet]、Stable Diffusion 3.5 Medium[stability2024sd35]、Gemma 4 E2B[google\_gemma4\_model\_card]和Qwen 3.5 35B-A3B[qwen35\_model\_card]。所有实验均在配备GB10 GPU的NVIDIA DGX Spark[nvidia\_dgx\_spark\_hw]上运行。每个内核仅经过50次优化迭代，Kernel Forge就生成了在所有四个模型上均超越PyTorch eager模式性能的内核。它在ResNet-50上达到adaptive\_avgpool2d的eager模式性能的1.52倍；在Stable Diffusion 3.5 Medium上达到group\_norm的1.70倍；在Gemma 4 E2B上达到softmax的2.83倍；在Qwen 3.5 35B-A3B上达到softmax的1.54倍。这些结果表明，LLM驱动的内核优化不仅可以应用于孤立基准测试，还可以应用于从完整PyTorch模型和代表性工作负载中捕获的算子。

在本文中，我们做出以下贡献：

1. 1. 我们提出了Kernel Forge，一个开源的端到端工具，通过捕获执行的算子、生成专门的CUDA内核、探索多种优化路径以及将优化后的内核集成回模型执行路径，来优化任何真实PyTorch模型的执行。
2. 2. 我们提供GUI和CLI，用于监控优化进度、检查生成的内核、调试故障以及在不同模型提供商之间进行选择。
3. 3. 我们在代表视觉、扩散和LLM工作负载的四个PyTorch模型上评估了Kernel Forge，结果显示，在每内核仅50次优化迭代下，选定算子上的eager模式性能提升高达2.83倍。
4. 4. 我们将Kernel Forge作为开源工具发布，以支持可重复的研究和智能体内核优化在PyTorch工作负载中的实际应用。

参见图注图2：Kernel Forge系统概览。模型摄取流水线将工作负载特定的PyTorch算子变体捕获为算子卡片（operator cards），而锻造炉（Forge）使用这些卡片生成、验证、基准测试并通过MCTS控制器迭代优化CUDA候选内核。

## 2 背景

### 2.1 用于内核优化的PyTorch算子捕获

PyTorch模型使用高级算子编写，例如矩阵乘法、卷积、归一化和softmax[paszke2019pytorch]。在推理期间，这些算子会调用执行实际计算的底层GPU内核。每个内核的效率取决于运行时传递给算子的张量和参数，包括形状、布局和数据类型。

一个独立的基准测试可能会使用随机生成的张量或一小套固定形状来测试内核[ouyang2025kernelbench]。相比之下，一个真实的PyTorch模型可能会以不同的形状、布局、参数以及周围的内存行为来调用同一个算子。因此，观察模型执行了哪些算子，并使用推理过程中观察到的输入来优化内核是很有用的。一旦知道了相关的算子和输入，优化器就可以为这些情况生成专门的内核。然后，生成的根内核可以在捕获的工作负载上根据PyTorch eager模式的输出进行验证，并在模型执行路径中使用。

### 2.2 GPU内核优化

GPU内核实现了PyTorch算子背后的底层计算。对于NVIDIA GPU，这些内核通常用CUDA编写，这使开发者能够控制线程组织、内存访问、同步和数据移动[nvidia\_cuda\_programming\_guide]。优化内核通常需要硬件感知的选择，例如分块（tiling）、内存合并、共享内存重用、算子融合以及减少同步或启动开销[nvidia\_cuda\_best\_practices]。

这些选择会显著影响性能，但手工做出这些选择很困难。开发者必须保留原始算子的行为，同时推理GPU硬件细节和工作负载特定的输入。现有库、编译器和领域特定语言（DSL），如TVM和Triton，有助于生成高效的实现[chen2018tvm,chen2018autotvm,zheng2020ansor,tillet2019triton]，但当模型以特定形状、参数或内存行为调用算子时，专门的内核仍然很有用[dao2022flashattention]。

### 2.3 智能体内核生成

最近的一些系统使用LLM通过迭代反馈循环来生成和改进GPU内核[hong2025autocomp,wang2025geak]。在一个典型的循环中，系统会生成一个候选内核，编译并运行它，检查其输出是否与参考实现一致，测量性能，并利用结果指导后续尝试。这使得可以在多次迭代中改进内核，而不是依赖一次生成的实现。

一个关键挑战是决定哪些候选内核值得继续探索。现有工具通常使用线性精炼（一次改进一个候选内核，通过一系列编辑进行优化）[cudaforge2025,astra2025]，或束搜索（保持多个候选内核，但可能优先考虑即时性能提升）[hong2025autocomp]。在内核优化中，这可能会造成局限：一个在早期迭代中较慢的候选内核可能在后期更改后变得有用，而一个在某个测试用例上较快的候选内核可能无法泛化到模型工作负载。蒙特卡洛树搜索（MCTS）在此场景中很有用，因为它可以探索多个优化路径，同时利用正确性和性能反馈将搜索集中在有希望的方向上[kocsis2006uct]。实际的内核优化因此需要代码生成、正确性检查、性能分析以及用于检查优化过程的工具。

## 3 Kernel Forge系统设计

Kernel Forge是一个端到端框架，用于优化在真实模型推理期间执行的PyTorch算子。给定PyTorch模型、其权重和示例输入，Kernel Forge会捕获支持的算子调用，将它们分组为工作负载特定的变体，并使用生成的CUDA内核优化选定的变体。每个候选内核在PyTorch eager执行的基础上进行验证和基准测试，然后通过持久化的MCTS搜索进行精炼。最终输出是一个受保护的包，仅在捕获的测量显示有效时才会分发生成的CUDA内核……【翻译中断，原文至此】