世界首个「AI CUDA工程师」诞生！AI自己写代码优化CUDA内核，性能狂飙100倍

Transformer论文八位作者之一Llion Jones创立的Sakana AI发布重磅成果——全球首个「AI CUDA工程师」！它能将PyTorch代码自动转换为高度优化的CUDA内核，速度比PyTorch原生实现快10-100倍。

就在刚刚，由Transformer八大金刚之一Llion Jones创办的Sakana AI，推出了世界上首个「AI CUDA工程师」——The AI CUDA Engineer。

同时，还发布了一个由「AI CUDA工程师」生成的，包含17,000多个已验证CUDA内核的数据集。

论文地址：https://pub.sakana.ai/static/paper.pdf

项目主页：https://pub.sakana.ai/ai-cuda-engineer

自从去年提出了世界上首个「AI科学家」之后，一个疑问一直萦绕在团队心头：如果可以使用AI来进行AI研究，是不是也能用AI来研究如何让AI运行得更快？

团队认为，AI系统也应该像人脑一样高效地利用资源。而实现的最佳途径，就是利用AI来让AI更高效！

于是，「AI CUDA工程师」就这样诞生了。

这位AI工程师，能自动生成高度优化的CUDA内核，不仅比PyTorch中常见的机器学习操作快10-100倍，而且也比生产中广泛使用的CUDA内核快最高5倍。

通过LLM驱动的进化代码优化技术，「AI CUDA工程师」能将PyTorch代码转换为CUDA内核，并通过进化算法优化 CUDA内核的执行效率，实现多个运算操作的融合。

更重要的是，它能够通过学习历史创新成果，发现编写高效CUDA运算的全新方法。

最终，开启一个AI驱动加速AI的新纪元，实现自动化推理时优化。

英伟达高级研究科学家Jim Fan表示：「干得漂亮，Sakana AI！这是我最近看到的最酷的自主编程AI智能体。AutoML王者归来！」

对计算资源最有价值的使用方式，就是提高相同计算资源的未来生产力。它具备了所有超越人类表现的关键要素：

1. CUDA调试异常繁琐，需要大量领域专业知识，有时候就像是炼金术；
2. 奖励函数定义得非常清晰：数值正确性和执行时间；
3. 进化周期很快，因为正确性检查很轻量，奖励评估也很简单；
4. 可以进行完全并行化的搜索。

而且，「创新档案」这个概念也很有趣，跟英伟达「我的世界」Voyage智能体中使用的「技能库」非常相似。

本质上，这是一个完全由LLM在探索CUDA设计空间时编写的代码仓库，还保存了最佳代码片段。这是进化搜索过程中产生的学习成果，无需更新参数权重。

Jim Fan顺便推荐了一下，英伟达最近也在展开的一项相关研究，同样也是用强大了的LLM推理来改进CUDA内核。

「AI CUDA工程师」：基于智能体的CUDA内核自动发现、优化、组合

CUDA是一个底层软件层，能够直接访问英伟达GPU的硬件指令集来进行并行计算。CUDA内核，则是使用CUDA语言编写的、在GPU上运行的函数。

通过在CUDA内核层面直接编写指令，我们能够显著提升AI算法的性能。

然而，使用CUDA需要扎实的GPU专业知识，而在实践中，大多数机器学习算法都是在PyTorch或JAX等更高层次的抽象层中编写的。

「AI CUDA工程师」是一个基于前沿LLM的AI智能体框架，目的是将标准PyTorch代码自动转换为高度优化的CUDA内核。

通过运用进化优化技术，并借鉴进化计算中的「交叉」操作和「创新档案」等概念来发现具有潜力的过渡性内核，研究人员提出的框架不仅能够自动将PyTorch模块转换为CUDA内核，而且他们的高度优化的CUDA内核通常能够实现显著的性能提升。

研究团队相信这项技术能够带来性能加速，从而加快LLM或其他生成式AI模型等基础大模型的训练和运行（推理）速度，最终让AI模型在英伟达GPU上运行得更快。

• 第1和第2阶段（转换和翻译）：「AI CUDA工程师」首先将PyTorch代码转换为可运行的CUDA内核。即便在没有专门优化的情况下，它们就已经观察到初步的性能提升。
• 第3阶段（进化优化）：借鉴生物进化中的适者生存法则，框架采用进化的方法来确保只保留最优秀的CUDA内核。同时，团队创新性地提出了内核交叉提示策略，能够将多个优化后的内核进行有效组合。
• 第4阶段（创新档案）：正如文化通过千年的文明传承来塑造人类智慧一样，「AI CUDA工程师」也会利用在过往创新和发现中积累的经验，构建一个由高性能CUDA内核构成的创新档案库，并以这些积累的优化经验为基础，实现更进一步的转换和性能突破。

提升内核运行速度

研究团队相信，这仅仅是人工智能伟大优化的开始！

在论文中，他们引入了一种端到端的智能体工作流，能够将PyTorch代码转换为可工作的CUDA内核，优化CUDA运行时性能，并自动融合多个内核。

此外，他们构建了各种技术来增强流程的一致性和性能，包括LLM集成、迭代分析反馈循环、本地内核代码编辑以及交叉内核优化。

AI发现的部分CUDA内核在KernelBench上实现了SOTA的性能

研究人员表示，「AI CUDA工程师」能够稳健地转换250个考虑的Torch操作中的230多个，并且对于大多数内核实现了显著的运行时性能提升。此外，该方法还能够有效地融合各种内核操作，并且可以超越几种现有的加速操作。

具体而言，在229个考虑的任务中，他们在81%的任务上优于PyTorch原生运行。此外，所有发现的CUDA内核中有20%至少比其PyTorch实现快两倍。

下图展示了PyTorch原生内核与经过优化处理后的内核在不同级别（Level 1、Level 2、Level 3）下的加速对比。

从下图展示的内核可以看出，「AI CUDA工程师」能实现各种操作，包括归一化方法、损失函数、特殊矩阵乘法，甚至是完整的神经网络架构。

17,000+个已验证CUDA内核数据集

研究团队随论文一同发布的「资料库」（The AI CUDA Engineer Archive）是一个包含超过3万个CUDA内核的数据集，这些内核均由「AI CUDA工程师」生成。

数据集采用CC-By-4.0许可证发布，可通过HuggingFace平台访问，并可用于LLM的进阶微调。

其中包括，PyTorch参考实现、PyTorch/NCU/Clang-tidy性能分析数据、每个任务的多个内核版本、错误信息，以及与PyTorch原生和编译运行时间的性能对比数据。

团队希望这个数据集能够促进开源模型的后续训练，从而开发更出色的CUDA功能模块。包括离线强化学习、偏好优化和标准监督微调等技术。

数据集地址：https://huggingface.co/datasets/SakanaAI/AI-CUDA-Engineer-Archive

统计数据显示，「资料库」共包含超过3万个内核，其中已验证正确的实现超过1.7万个，约有50%的内核性能优于PyTorch原生运行时间。

基于文本嵌入进一步分析发现，这些内核可根据不同的任务类型（如矩阵乘法、池化操作、卷积运算）和实现策略（包括循环展开、运算融合、向量化处理）自然地聚类。

探索内核资料库

除此之外，研究人员还开发了一个交互式网站，供用户实时查看超过1.7万个已验证内核及其配置信息，包括PyTorch、NCU和Clang-Tidy的相关数据。

在网站中，用户可以探索涵盖230个任务的各类高性能内核。并且可以动手定制一个排行榜，查看不同实验和LLM间的相关内核。

内核性能排行榜

并且，还可以直观地查看内核可视化结果、搜索相关内核、下载验证代码和性能提升数据，以及查看详细的性能分析结果。并且还可以深入了解优化实验的具体细节。

性能分析数据、评估脚本下载、相关内核信息和实验发现细节的实例归一化内核

AI革命才刚刚起步

研究人员认为，今天的LLM就像是我们这一代的「大型机」。

我们仍然处在AI发展的早期阶段，在市场竞争和全球创新的推动下（尤其是那些在资源受限条件下进行创新的团队），这项技术必将实现百万量级的效率提升。

目前，我们的AI系统消耗着海量资源，如果这项技术继续扩张而不考虑效率和能源消耗问题，其发展将难以持续。

从根本上来说，没有任何理由表明我们的AI系统不能达到甚至超越人类智能的效率水平。我们相信，要实现这种效率提升，最佳途径就是利用AI来优化AI本身。

这正是Sakana AI努力的方向，而这个项目正是朝着实现AI百万倍性能提升目标迈出的关键一步。

就像从早期庞大臃肿的大型机演变到现代计算机一样，相比当今「臃肿」低效的LLM，未来几年我们使用AI的方式将会发生翻天覆地的变化。

参考资料：

https://sakana.ai/ai-cuda-engineer/

文章来自于“新智元”，作者“编辑部 HNZ”。