AI论文精选
⏶22

CUDA-L1：通过对比强化学习改进CUDA优化

发表
由 Xiaoya LiXiaoya Li 提交
作者: Xiaoya LiXiaoya Li, Xiaofei Sun, Albert Wang, Jiwei Li, Chris Shum

摘要

由大型语言模型（LLMs）的快速发展所驱动的GPU计算资源需求的指数级增长，对自动化CUDA优化策略产生了迫切的需求。尽管LLMs的最新进展在代码生成方面展现出潜力，但当前最先进的模型（例如R1、o1）在提高CUDA速度方面的成功率较低。在本文中，我们介绍了CUDA-L1，一个用于CUDA优化的自动化强化学习框架。 CUDA-L1在CUDA优化任务上取得了性能提升：在NVIDIA A100上训练后，它在KernelBench的所有250个CUDA核（kernel）上实现了平均17.7倍的加速，峰值加速达到449倍。此外，该模型还展示了卓越的跨GPU架构可移植性，尽管专门针对A100进行了优化，但在H100上实现了平均17.8倍的加速，在RTX 3090上为19.0倍，在L40上为16.5倍，在H800上为14.7倍，在H20上为13.9倍。除了这些基准测试结果，CUDA-L1还展示了几个显著特性：1）发现多种CUDA优化技术，并学习策略性地组合它们以实现最佳性能；2）揭示CUDA优化的基本原理；3）识别非显而易见的性能瓶颈，并拒绝看似有益但实际有害于性能的优化。 CUDA-L1的能力表明，强化学习仅通过基于加速比的奖励信号，无需人工专业知识或领域知识，即可将最初性能不佳的LLM转化为有效的CUDA优化器。更重要的是，经过训练的强化学习模型能够将其习得的推理能力扩展到新的核。这种范式为自动化优化CUDA操作开辟了可能性，并有望大幅提升GPU效率，缓解GPU计算资源日益增长的压力。
查看 arXiv 页面查看 PDF
CUDA-L1：通过对比强化学习改进CUDA优化

评论

Xiaoya LiXiaoya Li
论文作者
论文提交者

我们想介绍一下 CUDA-L1，这是一个通过对比强化学习训练的LLM，旨在生成优化的CUDA内核代码。

📈 性能表现：

  • CUDA-L1 在 KernelBench 中的所有 250 个 CUDA 内核上实现了平均 3.12 倍（中位数 1.42 倍）的加速。

  • CUDA-L1 可以在不同的 GPU 架构上生成内核加速。例如，在 L40 上实现 3.12 倍加速，在 H100 上实现 2.39 倍加速，在 H20 上实现 2.37 倍加速等。

🌟 CUDA-L1 亮点：

  • 它发现了多种 CUDA 优化技术，并学会了如何策略性地组合它们。

  • 它揭示了 CUDA 优化的基本原理，例如优化的乘法特性。

  • 它识别出不明显的性能瓶颈，并拒绝了那些看似有益但实际上会损害性能的优化。

感谢您的时间和对社区的贡献。

Aryan V SAryan V S

你好 @xxiaoyali！几天/几周前我听说这篇论文时就对它很感兴趣了（它在 X 上很热门）。贵团队尝试使用强化学习（RL）进行内核生成，这是一项很棒的工作，然而，许多结果让我们这些爱好者感到怀疑。

包括我在内的一群人进行了独立测试，发现我们查看的所有问题都存在严重的基准测试错误和糟糕的基线。我们对论文中呈现的结果的主要抱怨可以总结为：

  • 基线表现不佳

  • 使用非默认流时，计时未进行适当同步

  • 每次重复运行基准测试时未清除设备缓存

  • 热身期间的缓存结果通过了正确性测试，但热身后的内核实际结果不正确（由于未处理同步）

  • 使用 CUDA 图来展示加速。现有系统的一项功能使得基准测试结果偏向您的论文。至少在我看来，正确的比较方式应该是比较 PyTorch eager 模式 + CUDA 图与您优化的内核 + CUDA 图，或者在两种情况下都不应用它。

  • 使用自定义的 PyTorch CUDA/cuDNN 后端标志来展示加速。同样，PyTorch 框架中可用的功能应该在基线和您优化的内核中都使用，或者都不使用，否则所呈现的基线会更差。

  • 1.K 倍到约 3 倍的加速通常可以归因于 PyTorch 框架中已有的可定制标志和功能。一些问题的解决方案似乎没有使用新颖的想法，但由于这个原因仍然显示出加速。即使生成的内核对于特定的问题形状更好，但在一般情况下它们仍不如专家优化的内核。

  • 一些问题没有自定义生成的 CUDA 代码，而只是调用 PyTorch API。

  • （还报告了其他一些问题，但我目前想不起来了）

  • 此外，将生成的内核与 Inductor（torch.compile）生成的 Triton 程序进行比较，将是评估通过应用强化学习进行内核生成所发现的真正创新的良好实践。

就我个人而言，我随机测试了总共 5 个问题，在修正了基准测试问题、移除了 CUDA 图的使用、PyTorch 后端标志以及其他框架提供的优化后，发现几乎没有或根本没有加速（< 1.5 倍）。这并不是说这项工作不令人惊叹，但其方法论确实引人深思。

尽管如此，我希望再次尝试这个项目，以了解通过强化学习发现并应用的重要优化。请问，GitHub 代码和论文的最新版本是否已用正确的基准测试方法进行了修正？

祝贺您的工作，感谢您将其开源！我期待贵团队在该研究领域取得更多令人惊叹的成果。

Xiaoya LiXiaoya Li
此评论已隐藏。
Jiwei LiJiwei Li

您好，@a-r-r-o-w

非常感谢您提出的宝贵意见，我们对此深表感谢。

同步不当的计时问题

我们已在最新版本中纠正了这个问题，并更新了 KernelBench 评估中的计时方法，以确保正确的同步。实际上，所有当前代码都不再包含新构建的 CUDA 流。

缓存结果

我认为它们现在不应该影响加速比，因为它应该无法通过正确性检查。由于评估是使用不同的输入进行的，这意味着其影响并未反映在性能中（如果我错了请指正）。

但我同意这确实令人困惑，我们将通过删除所有缓存来更新代码。

有些问题没有自定义生成的 CUDA 代码，只是调用 PyTorch API

很抱歉造成困惑。这是因为我们只是为非加速任务复制了原始代码。我们将纠正这个问题。

使用 CUDA 图来展示加速。现有系统的一个特性使基准测试偏向于您的论文。在我看来，正确的比较方法应该是比较 PyTorch Eager + CUDA 图 与您的优化内核 + CUDA 图，或者在两种情况下都不应用。

这篇论文之所以这样组织，是因为我们使用 KernelBench 参考代码作为强化学习训练的起点。因此，我们将我们的结果与原始的 KernelBench 参考代码进行了比较。

但我完全同意您的观点！我们将 a) 将当前版本与 torch.compile 版本进行比较；b) 与 PyTorch Eager + CUDA 图进行比较；c) 使用 torch.compile 版本作为起点重新训练模型。

再次非常感谢您抽出时间撰写评论。我们将很快更新论文和代码，并随时向您通报最新进展。