Kernel Forge: An Agent Harness for LLM-based Generation and Optimization of CUDA Kernels
Kernel Forge 是一个开源的端到端智能体框架，它利用大语言模型和蒙特卡洛树搜索来自动生成并优化针对多样化 PyTorch 工作负载的 CUDA 核函数，在实现比 PyTorch eager 模式显著加速的同时，还提供了用于检查和调试的图形化界面。
原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。 阅读完整免责声明
想象一下你正试图在电脑上运行一款电子游戏。有时，游戏运行得非常流畅，但有时却会出现卡顿或延迟。这是因为电脑必须每秒执行数百万次微小的数学计算，以绘制图像并处理物理效果。在人工智能（AI）的世界里，这些计算甚至更加密集。为了让 AI 进行“思考”，它依赖于被称为**内核（kernels）**的特殊、高速的指令。你可以把内核想象成厨师在切菜时使用的一种特定的、超高效的食谱。如果厨师使用钝刀和缓慢的方法，整顿晚餐就会变晚；如果他们使用锋利且专业的刀具，食物就能瞬间准备就绪。
多年来，编写这些“食谱”就像是用一种你几乎不精通的语言来写小说。这需要高度专业的专家来为计算机芯片（GPU）编写复杂的指令，以使其运行得更快。但现在，我们有了一种新的助手：大语言模型（LLMs）。你可能知道它们是那些可以写故事或解决数学问题的智能聊天机器人。科学家们正在教这些聊天机器人编写“食谱”（内核），希望它们能比人类做得更快、更好。然而，这里有一个陷阱：仅仅因为一个聊天机器人写出的食谱在纸面上看起来很好，并不意味着它在真实的厨房里真的能把饭做得更快。
这就是名为 Kernel Forge 的新项目发挥作用的地方。来自密歇根大学的研究人员构建了一个智能系统，它不仅仅是要求聊天机器人编写代码；它更像是一个全职教练、一名质量检查员以及一名项目经理的结合体。他们想看看 AI 是否真的能够接管这些 AI 模型的现实世界优化工作，而不仅仅是在虚构的测试案例中表现出色。
问题所在：“孤立”陷阱
想象一下你正在训练一个机器人跑步。如果你只在一条完美平坦、空旷的跑道上测试它，它看起来可能像是一名世界级的短跑选手。但一旦你把它放到一条崎岖、拥挤的城市街道上，它可能会摔倒。这正是以往尝试使用 AI 编写 GPU 内核时发生的情况。大多数工具通过给 AI 提供随机的、虚构的数学问题并在空旷的跑道上进行测试来测试它。它们生成一段代码，在隔离状态下进行测试，然后说：“看，这有多快！”
但在现实世界中，AI 模型就像繁忙的城市。这些“食谱”（内核）必须与特定的数据形状、特定数量的内存协同工作，并且必须与旁边的其他食谱和谐共处。一个在随机测试中很快的食谱，在用于像生成图像或与你聊天这样的真实 AI 模型内部时，可能会变慢甚至崩溃。以前的工具往往把如何将这些新的、由 AI 编写的食谱重新整合进复杂机器的任务留给了人类开发人员，这既繁琐又容易出错。
解决方案：Kernel Forge
研究人员创建了 Kernel Forge，这是一个作为 AI 编写者与现实世界之间桥梁的开源工具。Kernel Forge 不再只是让 AI 去猜测，它首先会观察一个真实的 AI 模型（例如识别面部或编写故事的模型）在计算机上运行的情况。它会记录下究竟正在使用哪些“食谱”、数据有多大以及运行时间有多长。
一旦掌握了这些现实世界的数据，它就会将任务交给一个 AI 智能体（一个由大语言模型驱动的智能机器人）。但这并不是一次性的工作。该系统使用了一种被称为**蒙特卡洛树搜索（Monte Carlo Tree Search）**的巧妙策略。想象一位侦探正在破解谜案：侦探不会只凭一个直觉，而是会尝试许多不同的路径。如果一条路径导致了死胡同，他们不会放弃，而是回到原点尝试另一个角度。同样地，Kernel Forge 不仅仅是编写一个版本的食谱然后听天由命。它会生成许多个版本并进行测试，如果其中一个很慢，它会尝试修复它，或者尝试一个完全不同的方法。它会保留所有尝试过程的“家族树”，即使最初并不完美，也会记住那些有潜力的尝试。
结果：速度提升，但并非到处如此
团队在四种不同类型的 AI 模型上测试了 Kernel Forge：一个用于识别图像的模型（ResNet-50），一个用于生成艺术的模型（Stable Diffusion 3.5 Medium），以及两个用于聊天和推理的模型（Gemma 4 和 Qwen 3.5）。他们在配备了 NVIDIA GB10 GPU 的高性能计算机上运行了这些测试。
以下是他们的发现：
- 它有效，但很有挑剔性： 该系统成功地为许多 AI 部分生成了新的、更快的食谱。在某些情况下，AI 编写的代码比计算机通常使用的标准代码快得多。例如，在图像生成器中，用于“组归一化”（group normalization）的新代码快了 1.70 倍。在聊天机器人 Gemma 4 中，用于“softmax”（一种用于做决策的数学步骤）的新代码更是实现了惊人的 2.83 倍加速。
- “受保护”的安全网： 该系统非常谨慎。它有一个“守卫”来检查每一个新食谱。如果新的 AI 编写的食谱变慢了或者出现了错误，系统会立即切换回原始的、可靠的代码。它绝不会把一个糟糕的食谱强加给用户。这意味着，即使 AI 尝试优化某段代码并失败了，计算机也不会崩溃或变慢，它只会使用旧的、安全的方法。
- “大玩家”难以超越： 研究人员注意到了一些有趣的事情。那些占据最多时间的 AI 部分（即“大玩家”）通常已经由像 NVIDIA 这样公司的成熟、专家级编写的代码处理了。AI 在挑战这些专家方面显得有些吃力。例如，在图像生成器中，“线性”（linear）操作（占据了超过 50% 的时间）在 AI 尝试重写它时实际上变得更慢了。系统明智地决定对这部分保持原有的代码。
- 小胜积聚成大胜： 最显著的速度提升发生在代码中较小、但不那么关键的部分。虽然 AI 在处理最大的任务时无法击败专家，但它找到了巧妙的方法，将这些较小的任务提速了 1.5 到 2.8 倍。
好奇心的代价
研究人员还调查了这种“思考”的成本。他们使用了一个强大的 AI 模型来生成代码，而每当 AI 编写一行代码或检查一个结果时，都会产生一点费用（基于 API 使用量）。他们发现，随着他们让 AI 尝试更多的变体（多达 50 轮尝试与修复），成本也会上升。然而，他们获得的额外速度并不总是能匹配所花费的额外资金。这表明，虽然 AI 可以找到极佳的捷径，但我们需要聪明地决定何时停止搜索，特别是当我们要修复的代码部分对整体速度影响不大时。
总结
Kernel Forge 表明，我们正在进入一个 AI 可以帮助编写使其他 AI 运行得更快的底层代码的新时代。它不是一个能瞬间解决一切的魔杖——它目前还不能轻易击败处理最大任务时的世界顶尖人类专家。但它是一个强大的工具，可以在保持系统安全和稳定的同时，发现隐藏在细节中的速度提升。
研究人员将他们的工具作为开源发布，这意味着任何人都可以使用它来尝试让自己的 AI 模型运行得更快。他们证明了，通过将 AI 的创造力与现实世界的测试安全性相结合，我们可以开始优化我们数字世界的引擎，而无需具备计算机科学博士学位。这是迈向一个未来的一步——在那个未来，我们的计算机不仅更聪明，而且效率也更高。
您所在领域的论文太多了？
获取与您研究关键词匹配的最新论文每日摘要——附技术摘要，使用您的语言。