剪辑：剪辑部

斯坦福和普林斯顿商酌者发现，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引诱漫衍式系统体育游戏app平台，并在NVIDIA DRIVE部门从事自动驾驶汽车的引诱责任。