搜索
【导读】斯坦福和普林斯顿研究者发现，DeepSeek-R1生成的自定义CUDA内核，完爆了o1和Claude 3.5 Sonnet，拿下总排名第一。虽然目前只能在不到20%任务上超越PyTorch Eager基线，但GPU编程加速自动化的按钮，已经被按下！
近日，来自斯坦福和普林斯顿的研究者发现，DeepSeek-R1已经能生成自定义CUDA内核了，而且还在一众推理模型中，直接拿下了TOP 1！
紧随其后，OpenAI o1和Claude 3.5 Sonnet分别排第二和第三。
具体过程，就是给定一个PyTorch程序，让模型对其优化，然后生成一个包含自定义CUDA内核的PyTorch版本。
在此期间中，模型可以自由决定优化哪些操作，以提高计算效率。
引导模型生成GPU内核，潜力巨大
如今，传统人工优化内核的方式，在效率上已经不足以应对大量涌现的AI模型架构和硬件平台。
既然是为了LLM进行优化，那么，我们是否也能够借助LLM来模拟AI工程师的工作流程，凭借编译器反馈、硬件规格等丰富的信息，自动编写出准确且经过优化的内核代码呢？
为此，研究团队提出了一种全新的KernelBench框架，用于生成和评估不同AI任务（单个操作、操作序列、端到端架构）的内核，并模拟了工程师迭代优化的过程。
论文地址：https://arxiv.org/abs/2502.10517
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
https://arxiv.org/abs/2502.10517
https://x.com/simonguozirui/status/1894455889039692263
