[博客翻译]AI CUDA工程师：代理CUDA内核发现、优化和组合

由 发布于 cuda优化人工智能智能代理并行计算
 收藏

原文地址：https://sakana.ai/ai-cuda-engineer/


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工程师成功部署的各种不同操作，包括归一化方法、损失函数、特殊矩阵乘法甚至完整的神经网络架构：
[](https://sakana.ai/ai-cuda-engineer/<https:/pub.sakana.ai/ai-cuda-engineer/kernel/1/97/optimize-b10-s4-e0-sweep/9/2/0/blocks