近日，来自斯坦福和普林斯顿的研究者发现，DeepSeek-R1 已经能生成自定义 CUDA 内核了，而且还在一众推理模型中，直接拿下了 TOP 1！
紧随其后，OpenAI o1 和 Claude 3.5 Sonnet 分别排第二和第三。
具体过程，就是给定一个 PyTorch 程序，让模型对其优化，然后生成一个包含自定义 CUDA 内核的 PyTorch 版本。
在此期间中，模型可以自由决定优化哪些操作，以提高计算效率。
引导模型生成 GPU 内核，潜力巨大
如今，传统人工优化内核的方式，在效率上已经不足以应对大量涌现的 AI 模型架构和硬件平台。
既然是为了 LLM 进行优化，那么，我们是否也能够借助 LLM 来模拟 AI 工程师的工作流程，凭借编译器反馈、硬件规格等丰富的信息，自动编写出准确且经过优化的内核代码呢？
为此，研究团队提出了一种全新的 KernelBench 框架，用于生成和评估不同 AI 任务（单个操作、操作序列、端到端架构）的内核，并模拟了工程师迭代优化的过程。
GPU 的本质，是硬件依赖的。因此，研究者们希望尝试，看是否能通过以下方式，引导模型生成 GPU 内核。
首先，向模型提供硬件信息（如内存带宽、TFLOPS），以针对特定 GPU（A100、H100 等）进行优化。
然后，要让模型在上下文中展示代表性的内核优化技巧，例如矩阵乘法中的分块（tiling）或 Flash Attention 中的在线 softmax。
研究者们发现，只有更强大的模型，会偶尔表现出利用这些优化的能力。
上下滑动查看
比如，DeepSeek-R1 有时会使用特定于硬件的指令（如 Tensor Core 的 wmma），但往往无法正确编译或使用它们，从而限制了最终性能。
总的来说，研究发现，前沿模型在 KernelBench 上的开箱即用性较差，OpenAI o1 和 DeepSeek-R1 在不到 20% 的任务上超过 PyTorch Eager 基线。
这些模型生成的内核存在大量执行错误、功能正确性问题，并且无法进行特定平台的优化。
具体来说，研究者发现：
-
对模型而言，编写功能正确的内核仍然具有挑战性；
-
模型通过优化展示了生成高性能内核的潜力；
-
利用反馈对于减少执行错误和发现更快的方案很重要。
当然，KernelBench 目前还只是让 GPU 加速奔跑的起点，但也是让整个 GPU 编程自动化的起始催化剂。
令人兴奋的是，现在已经有了许多新的工作，专注于解决 KernelBench 中涉及的问题。
比如 2 月 12 日，英伟达就发出博客文章，探讨如何使用 DeepSeek-R1 进行 GPU 内核自动生成与推理时 scaling。
随后在 2 月 12 日，Meta 也发文测试了前沿模型编写 GPU 内核方面的性能，他们发现，最佳模型可以在 KernelBench 上提供平均 1.8 倍的加速。
Sakana AI 更是推出「AI CUDA 工程师」，让 AI 自己写代码优化 CUDA 内核，速度声称比 PyTorch 原生实现快了 10-100 倍。
如雨后春笋般出现的研究表明，如今，我们已经进入了 AI 驱动加速 AI 的新纪元！
在未来，KernelBench 还将持续演进。它不会仅限于当前收集的 250 个问题，还可以扩展到新的 AI 任务。与此同时，评测指标 fast_p 也可以随着时间的推移进行调整，提高加速门槛，以推动更高效的优化方案
KernelBench：AI 内核生成框架
KernelBench 是一个开源框架，旨在评估 LLM 在编写 GPU 内核方面的能力。
任务格式
KernelBench 包含 250 个任务，涵盖了各种 AI 工作负载，并且易于扩展到新的工作负载。
下图 1 展示了 KernelBench 评估语言模型（LM）生成高性能 GPU 内核的能力。KernelBench 要求语言模型为给定的目标 PyTorch 模型架构生成优化的 CUDA 内核，并进行自动化评估。
· 任务输入
给定一个 AI 工作负载，任务的输入是用 PyTorch 编写的参考实现。模仿研究人员的工作流程，PyTorch 代码包含一个继承自 torch.nn.Module () 的名为 Model 的类，其中标准的__init__和 forward () 函数（以及任何辅助函数）被填充为 AI 工作负载的 PyTorch 操作。
AI 算法通常在大型张量数据上进行操作。工作负载的最优内核取决于张量的大小和数据类型（如 BF16、FP8）。因此，每个任务还包含 get_inputs () 和 get_init_inputs () 函数，用于指定内核需要处理的精确输入张量。
· 任务输出
给定输入，LLM 需要输出一个继承自 torch.nn.Module () 的名为 ModelNew 的新类，其中包含自定义优化。例如，LLM 可以在 forward () 函数中使用 PyTorch 的 CUDA-C 扩展来集成内联内核调用。
为了成功完成任务，模型需要确定（1）Model 类中的哪些操作最能从优化中受益；（2）如何优化这些操作。LLM 可以使用任何硬件高效技术（如融合和分块）或专用指令（如张量核心）以及任何编程库（如 PTX、CUDA、CUTLASS、Triton、ThunderKittens）。
上下滑动查看
任务选择
这些任务根据包含的基本操作或 PyTorch 库函数的数量分为三个级别。
Level 1 包含 100 个单个基本操作，如卷积、矩阵乘法等 AI 基础构建块。虽然 PyTorch 调用了经过优化的闭源内核，让 LLM 超越基线具有挑战性，但如果能生成开源内核，将有重要价值。
Level 2 包含 100 个操作序列，如卷积、ReLU 和 Bias 的组合，这些操作可以融合成一个内核以提高性能。
由于基于编译器的工具（如 PyTorch 编译器）在融合方面非常有效，LLM 要在这方面超越它们也具有挑战性。然而，LLM 可能会提出更复杂的算法。
Level 3 包含 50 个完整的机器学习架构，如 AlexNet 和 MiniGPT 等，这些架构在运行训练和推理时对内核的性能要求极高。
评估指标
KernelBench 引入了一个新的评估指标 fast_p，衡量生成的内核中功能正确且加速大于阈值 p 的任务比例。
通过调整阈值参数 p，研究者可以评估不同加速阈值下的内核性能，并捕捉加速分布。
fast_0 相当于 LLM 的正确率，它衡量代码功能正确的任务比例，而不考虑其速度。在实际评估中，通常以 p=1 作为起点。
LLM 在 KernelBench 上的表现
研究人员对一系列 LLM 在 KernelBench 上进行了评估，结果显示，目前的 LLM 在生成正确且优于 PyTorch 基线速度的内核方面仍有困难。
在一次性基线评估中，LLM 生成的内核平均在不到 20% 的任务中比 PyTorch Eager 更快。这表明，仅靠简单提示，LLM 很难在性能上超越传统的 PyTorch 内核。
LLM 生成的内核存在大量的执行错误和功能正确性问题，经常由于简单的编译器和运行时错误而失败。
执行错误包括 CUDA/nvcc/Python 编译时错误、CUDA 内存违规和运行时错误等；正确性错误则主要表现为输出张量形状和值不匹配。
推理模型（o1，R1）生成的错误解决方案（<55%）比其他模型（>70%）少。然而，这主要是由于执行失败的情况较少。在功能正确性方面，所有 LLM 都面临类似程度的困难。
在性能方面，模型生成功能正确的内核在多数情况下也未能优于 PyTorch 基线。
随着 p 的提高，模型生成的内核中能达到要求的比例越来越低。在 p=1 时，在所有 KernelBench 级别中，不到 15% 的 LLM 生成内核优于 PyTorch。
推理模型通常在提供加速方面优于其他 LLM，但总体仍有不足。
模型生成的内核在不同硬件平台上的通用性不佳。DeepSeek-R1 生成的内核在 NVIDIA L40S 上实现了 36% 的加速，而在 NVIDIA A10G 上则为 47%。
这表明 LLM 在生成特定目标硬件的高效内核方面还存在很大的提升空间。
模型能力分析
测试时利用 KernelBench 环境反馈
正如上面观察到的，执行失败是 LM 生成的内核中最常见的失败模式。
KernelBench 提供的环境允许收集丰富的信号，包括编译器错误、正确性检查和运行时性能分析指标，所有这些都可以反馈给 LM 以帮助它解决内核故障。
为了探索 LM 如何利用这些反馈，研究团队评估和比较了两个基线：第一个令 LM 为每个 KernelBench 任务生成多个并行样本，另一个通过允许 LM 利用执行反馈逐步改进，依次为每个 KernelBench 任务生成内核。
· 重复采样
KernelBench 环境支持对 LM 生成的内核进行程序化验证，允许研究团队收集和评估每个任务的多个 LM 生成。他们使用 fastp@k 评估这种重复采样方法。重复采样有助于 LM 发现更多快速且正确的解决方案。
如下图 4 所示，随着 k 值的增加，在 DeepSeek-V3 和 Llama 3.1 70B 的三个级别上，通过高温度参数重复采样可以提升 fast1 的性能。
值得注意的是，在 Level 2 上，DeepSeek-V3 在 k=100 个样本时达到了 37% 的 fast1，而在单次提示基线中仅为 4%。
通过检查样本，我们发现高温采样有助于探索解决方案空间，增加了生成具有更好优化的无错误内核的机会。然而，如果一个模型解决任务的固有概率非常低，仅仅增加采样预算的影响有限。
例如，即使尝试了 100 个样本，DeepSeek-V3 也从未能够为 Level 1 中的一组 34 个卷积变体生成任何正确的解决方案。
· 生成结果的迭代优化
KernelBench 环境非常适合收集编译器反馈、执行错误和使用 PyTorch 分析器等工具进行的时间分析，作为真实信号（ground-truth signals）。
研究人员研究了利用这些反馈是否能帮助语言模型（LMs）迭代优化其生成结果。
下图 5 显示，KernelBench 框架使模型能够在迭代优化过程中接收并利用反馈。这些真实信号包括 NVCC 编译器错误信息、执行统计数据（例如正确性检查和挂钟时间），以及 PyTorch 分析器（操作时间分解）。
他们在多轮过程中为模型提供每次生成的反馈：在初始生成后，向模型提供其之前的生成结果 G，以及当前生成对应的编译器 / 执行反馈 E 和 / 或分析器输出 P。
然后将每次生成及其后续反馈定义为一轮（turn），并在 N 轮内运行这一迭代优化过程。利用执行反馈有助于减少错误，并随时间提升整体加速效果。
研究人员在下表 2 中检查了第 N=10 轮时的 fast1 行为，发现迭代优化在不同模型和 KernelBench 的各个级别上均持续提升了性能。
DeepSeek-R1 在 Level 2 上的改进最为显著，其中执行反馈 E 和分析器反馈 P 的组合将 fast1 从 36% 提升至 72%（如下图 6 所示）。
此外，通过分析迭代优化轨迹，他们发现模型在执行反馈 E 的帮助下能更有效地自我纠正，尤其是在修复与执行错误相关的问题上。
DeepSeek-R1 在 Level 1 和 Level 2 上，经过 10 轮优化后，能在超过 90% 的任务中生成功能正确的内核（下表 9）。
然而，剩余的错误内核几乎总是由于功能不正确而失败，这可能是因为正确性反馈的颗粒度不如执行失败信息细致。
· 比较重复采样与迭代优化
在上表 2 中，研究人员比较了在固定 10 次推理调用预算下重复采样和迭代优化的效果。两种方法相较于单次基线（one-shot baseline）均取得了显著改进，其中迭代优化在 6 个案例中的 5 个中表现更优。
然而，他们最终发现，测试时方法的效果本质上依赖于基础模型的质量。
例如，在重复采样中，DeepSeek-V3 在所有三个级别上始终优于 Llama-3.1 70B。类似地，在迭代优化中，DeepSeek-R1 通过反馈 E 和 P 持续改进，而 DeepSeek-V3 和 Llama-3.1 70B 并非总能从这些信息中获益。
提供硬件知识生成硬件高效内核
显然，语言模型在生成硬件高效内核方面表现有限。
这可能是由于训练数据中内核代码的稀缺性，以及最佳内核可能需要根据硬件平台的特定属性而变化。
在本案例研究中，研究团队探索了提供以下内容的效果：（1）提供内核工程最佳实践的示例，并将其置于（语言模型的）上下文之中；（2）提供详细的硬件规格说明，并将其置于（语言模型的）上下文之中。
· 硬件感知的上下文示例
编写良好的内核通常使用融合（fusion）、分块（tiling）、重计算（recompute）和异步（asynchrony）等技术来最大化性能。
具体来说，研究人员纳入了三个上下文示例：使用操作融合的 GeLU、使用分块的矩阵乘法，以及展示共享内存 I/O 管理的最小 Flash-Attention 内核。
结果则显示，上下文示例降低了语言模型的整体 fast1 分数，因为模型尝试了更激进的优化策略，但导致更多执行失败。与上面基线生成的代码相比，OpenAI o1 在使用少样本示例时生成的代码平均长度增加了 25%。
然而，在正确的解决方案中，语言模型应用了一些有趣的优化：他们发现，在 KernelBench Level 1 的 77% 的 GEMM 变体中，o1 应用了分块并提升了速度，优于单次基线。在 Level 2，o1 在 11 个问题上应用了激进的共享内存 I/O 管理，并能够超越 PyTorch Eager。
· 指定硬件信息
正如上面讨论过的，内核性能因硬件平台而异。
例如，FlashAttention-2 从 NVIDIA A100 迁移到 H100 GPU 时硬件利用率下降了 47%。FlashAttention-3 是一个完全不同的算法，专为 H100 编写。
在本研究中，研究团队探讨语言模型是否能利用上下文中的以下信息生成改进的内核：（1）硬件规格，例如 GPU 类型（H100、A100 等）、内存大小、带宽、TFLOPS；（2）硬件知识（例如线程、线程束、线程块、流多处理器的定义）。
结果显示，模型很少生成针对底层硬件优化的内核，这表明未来模型仍有改进空间。
某些新一代 GPU（例如 H100）引入了与前代不同的新硬件单元和指令。提供硬件信息对 Llama 3.1 70B 或 DeepSeek-V3 的输出影响不大。
有趣的是，他们发现 OpenAI o1 和 DeepSeek-R1 生成的部分内核使用了特定于硬件的指令和优化。
R1 在大约 50% 的 Level 1 矩阵乘法问题中尝试生成 warp 矩阵乘加（wmma）指令（下图 10），尽管大多数未能编译。
在功能正确的生成中，R1 和 o1 在每个级别产生了 1-3 个异常值，比 Level 4 基线快 2 倍以上。
总体而言，研究团队发现，与提供硬件信息相比，语言模型通过少样本示例调整策略时表现更好。
结论
研究人员在本论文中提出了 KernelBench，一个为语言模型驱动的内核优化奠定基础的框架；他们评估了多种模型和方法，分析了它们的优势和局限性，并提供了改进机会的见解。
总的来说，尽管大多数基准测试最终会达到饱和，但 KernelBench 被设计为随着新的 AI 工作负载的出现而动态演进。
他们的 fastp 指标可以随时间调整，以测量相对于日益先进的基线（即超出工作中使用的 PyTorch 基线）的加速阈值（p）。
由于 PyTorch 具备跨硬件平台兼容性，KernelBench 中基于 PyTorch 的任务可以在每个新硬件平台发布时进行评估。最后，与许多基准测试不同，在 KernelBench 上的成功直接映射到生产价值和现实世界的影响（降低成本并大规模减少能耗）。
这些特性确保了 KernelBench 在不断演变的 AI 领域中将保持其价值。
下一步工作
研究团队表示在当前可用模型下，KernelBench 仍有显著的改进空间。
首先，未来的工作可以探索开发先进的微调和推理技术，包括智能体工作流（agentic workflows）。由于 CUDA 是一种低资源语言，未来工作开源更多高质量数据将具有重要价值。
其次，在他们的实验中，语言模型生成的是原始 CUDA 代码。然而，未来的工作可以探索使用其他编程抽象（例如 ThunderKittens、CUTLASS、Triton 等）生成代码是否能简化生成问题，例如使语言模型更容易利用张量核心指令。
最后，研究团队的评估至今仅限于 GPU，未来的工作可以扩展到其他硬件加速器。
作者介绍
Anne Ouyang
Anne Ouyang 目前是斯坦福大学计算机科学（CS）博士生，在 Scaling Intelligence Lab（规模化智能实验室）进行研究。
她的研究兴趣主要集中在可扩展的自我改进机器学习系统，同时也广泛关注实证机器学习（empirical ML）和性能工程（performance engineering）。
她曾获得了 MIT 学士和硕士学位，并曾在 NVIDIA cuDNN 团队工作，负责编写 CUDA 内核，用于加速 GPU 上的深度学习工作负载。
Simon Guo
Simon Guo 是斯坦福大学计算机科学专业的一年级博士生，目前正在扩展智能实验室（Scaling Intelligence Lab）跟随 Azalia Mirhoseini 教授进行轮转研究。
他曾获得了 UC 伯克利电气工程和计算机科学学士学位。他的研究兴趣在计算机系统和机器学习。
最近，他在 Cohere 从事语言模型预训练工作。在此之前，他曾在苹果公司设计 GPU，在 Anyscale 开发分布式系统，并在 NVIDIA DRIVE 部门从事自动驾驶汽车的开发工作。
参考资料：HNYZs