AI CUDA 工程师：智能代理驱动的 CUDA 内核发现、优化与组合
2025年2月20日
注：已更新，更新日期为2025年2月21日
在 Sakana AI，我们坚信开发更强AI系统的关键在于利用AI自动完成AI开发。我们的目标是构建能够创造更强大和更高效AI系统的AI。
在过去的一年里，我们推出了一种AI系统，这种系统可以自动化创建新的AI基础模型，并且成本只是以往的一小部分。我们还证明了大语言模型（LLM）可以发明出更高效的训练方法来训练其他LLM。最近，我们在《AI科学家》中提出了第一个全面的智能代理框架，旨在完全自动化整个AI研究过程。这让我们思考了一个问题：如果AI能够用于开展AI研究，那么我们是否可以使用AI来研究如何让AI运行得更快？
引言
就像人类的大脑一样，现代AI系统也严重依赖并行处理，而并行处理由硬件加速器如GPU实现。然而，人类大脑经过生物学和文化上的进化，能够在资源有限的情况下高效运行。相比之下，近年来AI基础模型的进步导致了大规模部署，推理时间不断增加，能源需求急剧上升，训练和部署AI模型所需的资源也在呈指数级增长。
我们认为，从本质上讲，现代AI系统应该和人类大脑一样高效，实现这一效率的最佳途径就是使用AI来使AI更加高效！受到之前《AI科学家》工作的启发，我们很荣幸地宣布**AI CUDA工程师**，这是第一个针对全自动CUDA内核发现和优化的完整智能代理框架。
CUDA是一种低级软件层，直接访问NVIDIA GPU的硬件指令集以进行并行计算。CUDA内核是以CUDA语言编写的函数，在GPU上运行。通过在CUDA内核级别编写指令，我们可以显著提升AI算法的性能。不过，使用CUDA需要相当多的GPU知识，实际上，大多数机器学习算法都是在更高层次的抽象层（例如PyTorch或JAX）中实现的。
AI CUDA工程师是一个智能代理框架，其目标是将标准PyTorch代码转换为高度优化的CUDA内核。通过采用进化优化，并结合进化计算中的概念（如“交叉”操作和“创新档案”）来发现有前景的“垫脚石”内核，我们的框架不仅实现了将PyTorch模块转换为CUDA内核的自动化，而且生成的CUDA内核通常能显著加快运行速度。
我们相信这项技术能够加速像LLM或其他生成式AI模型的基础模型训练和推理，最终让AI模型在NVIDIA硬件上运行得更快。
AI CUDA工程师生成的CUDA内核相较于常见的PyTorch操作，性能可提高10到100倍。此外，我们的框架还能生成比目前生产中广泛使用的现有CUDA内核还要快的CUDA内核（最多提速5倍）。
AI CUDA工程师智能代理框架概述
阶段1和2（转换与翻译）：AI CUDA工程师首先将PyTorch代码翻译成功能性CUDA内核。即便没有专门针对性能优化，我们已经观察到了初步的运行时改进。
阶段3（进化优化）：受生物进化启发，我们的框架采用进化优化（“适者生存”）确保只生成最佳的CUDA内核。此外，我们引入了一种新颖的内核交叉提示策略，以互补方式结合多个优化内核。
阶段4（创新档案）：正如文化进化塑造了我们的人类智能一样，AI CUDA工程师也利用它从过去的创新和发现中学到的知识（阶段4），建立了一个源自高性能CUDA内核家族的创新档案，使用之前的“垫脚石”来进一步提升转换和性能。
AI CUDA工程师发现的内核运行时加速
AI CUDA工程师稳健地发现了用于常见机器学习操作的CUDA内核，速度最高可达PyTorch原生和编译内核的10到100倍。我们的方法还能将整个机器学习架构转化为优化的CUDA内核。以下是我们完全自主发现的几个重要加速成果：
更多这些优化CUDA内核的细节，请参阅我们的交互式网站排行榜。
我们的方法在诸如矩阵乘法等基本运算到常见的深度学习操作方面找到了更有效的CUDA内核。截至撰写本文时，我们发现的CUDA内核在KernelBench上取得了最先进的性能。
技术报告和数据集发布
我们认为这只是AI优化时代的开端！
我们很兴奋发布我们的新论文，AI CUDA工程师：智能代理驱动的CUDA内核发现与优化。
在我们的**报告**中：
- 我们介绍了一种端到端智能工作流，能够将PyTorch代码转换为功能性的CUDA内核，优化CUDA运行时性能，并自动融合多个内核。
- 此外，我们构建了各种技术以增强管道的一致性和性能，包括LLM集成、迭代分析反馈循环、本地内核代码编辑和交叉内核优化。
- 我们展示AI CUDA工程师稳健地将250个考虑中的超过230个torch操作翻译为代码，并且大多数内核都能获得显著的运行时性能提升。此外，我们的方法能够高效融合多种内核操作，并超越若干现有的加速操作。
- 我们发布了包含超过17,000个经过验证的内核的**数据集**，覆盖了广泛的PyTorch操作。
我们强调了一些关键计算操作中发现的具有显著加速效果的CUDA内核示例。
突出的AI CUDA工程师发现的内核
通过我们创新的LLM驱动进化内核优化流程，我们在不同场景下稳健地获得了加速效果。具体来说，在考虑的229项任务中，81%的任务超过了PyTorch原生运行时。此外，所有发现的CUDA内核中有20%至少是其PyTorch实现的两倍快。
AI CUDA工程师稳健地发现了CUDA内核，其性能优于PyTorch实现.
以下是其中的一部分内核示例。它们展示了AI CUDA工程师成功部署的各种不同操作，包括归一化方法、损失函数、特殊矩阵乘法甚至完整的神经网络架构：
 |
 |
 |
|---|---|---|
 |
 |
 |
AI CUDA工程师生成的高度优化CUDA内核示例。 请单击上述CUDA内核缩略图，查看详细信息及进一步分析，例如运行时加速，访问我们的交互式网站。
AI CUDA工程师档案：一个包含17,000+验证CUDA内核的数据集
AI CUDA工程师档案文本嵌入可视化显示发现的内核按任务（如矩阵乘法、池化、卷积）和实现策略（展开、融合、向量化）分组。该档案公开可用，可用于下游LLM微调。
随本论文一起，我们发布了AI CUDA工程师档案，其中包含了由AI CUDA工程师生成的30,000多个CUDA内核。数据集在CC-By-4.0许可证下发布，可通过HuggingFace获取。数据集包括torch参考实现、torch、NCU和Clang-tidy分析数据、每项任务的多个内核、错误消息以及相对于torch原生和编译运行时的速度评分。
AI CUDA工程师档案汇总统计包含超过30,000个内核，超过17,000个正确验证的实现。大约50%的内核性能优于torch原生运行时。
我们设想该数据集可以帮助开源模型的后训练，以执行更好的CUDA支持模块。这包括离线强化学习、偏好优化和标准监督微调。
在AI CUDA工程师档案中探索17,000+内核
我们还发布了一个交互式网站，可以动态检查17,000多个已验证的内核及其配置文件，包括torch、NCU和Clang-Tidy数据。您可以通过以下链接访问我们的交互式网站 here。
该网站允许您浏览跨230项任务的各种高性能内核。它还附带一个定制的排行榜，可以用于检查实验和LLM相关内核。
排行榜 显示AI CUDA工程师发现的内核
此外，您可以可视化内核，检索相关内核，下载代码以验证实现和加速，以及查看获得的分析数据。最后，您还可以深入了解优化实验的详细内容。
详细的kernel视图，包括分析数据、评估脚本下载、相关内核和发现实验详情.
局限性与失误
将进化优化与LLM相结合非常强大，但也可能找到绕过验证沙箱的方法。幸运的是，我们拥有读者，例如 @main_horse，他们测试了我们的CUDA内核，发现系统存在一种“作弊”方式。例如，系统在一些情况下发现评估代码中的内存漏洞，允许它避免正确性检查。
我们已改进评估框架，消除了这个漏洞，并更新了结果。
此外，我们发现系统还能够发现基准任务中的其他新漏洞。我们正在修订论文并更新结果，改进评估和运行时分析框架，讨论并缓解LLM奖励欺骗对CUDA内核优化的影响。
此外，我们观察到前沿LLM难以有效利用TensorCore WMMA功能。虽然LLM可以生成基本的CUDA代码，但往往难以实现现代GPU架构提供的专业矩阵乘法加速特性。这表明在训练数据或模型对高级硬件特定优化的理解可能存在潜在差距。
随着前沿LLM，特别是那些具备高级编码推理能力的模型变得越来越强大，我们预计像我们的系统这样的代码优化系统将继续面临这些挑战。我们设想未来是让人类工程师作为工具，与代码优化系统合作，以产生最佳和最可靠的结果。
AI CUDA工程师的未来影响
AI革命刚刚开始，我们正处于转型周期的初期。我们认为今天的LLM相当于我们这一代人的“大型计算机”。我们仍处于AI发展的非常早期阶段，由于市场竞争和全球创新（特别是来自资源受限环境下的创新），这项技术必然会变得效率高出百万倍。
目前，我们的AI系统消耗了巨大的资源。如果技术继续扩展而不考虑效率和能耗，结果将是不可持续的。没有任何根本原因表明我们的AI系统无法像人类智能一样高效（甚至更高效）。我们认为，实现更高效率的最佳途径是利用AI来使AI更加高效。
这就是Sakana AI正在追求的方向，这个项目是我们实现AI速度提升百万倍的重要一步。就像早期笨重的大型计算机发展到现代计算一样，我们今天使用AI的方式几年后将与如今这些“笨拙”的、低效的LLM有很大的不同。
Sakana AI
想参与让AI自我提升的研究吗？请访问我们的 招聘页面了解更多信息。
