斯坦福意外用AI生成超强CUDA内核 性能好得出奇！华人主创

斯坦福意外用AI生成超强CUDA内核，性能好得出奇！华人主创

这结果确实有点出人意料：AI意外生成的计算内核（kernel），在性能上竟然能超越人类专家精心优化的版本。

免费影视、动漫、音乐、游戏、小说资源长期稳定更新！ 👉 点此立即查看 👈

斯坦福大学近期披露的一组发现，数据相当亮眼。由AI优化的内核，在多项深度学习基础操作上，性能不仅翻倍超越原生PyTorch实现，部分项目提升幅度甚至接近400%。

具体来看几个核心操作的性能对比：

矩阵乘法（Matmul，FP32）：性能达到PyTorch `torch.matmul` 的101.3%。

二维卷积（Conv2D）：性能达到`torch.nn.Conv2D`的179.9%。

Softmax：性能达到`torch.softmax`的111.8%。

层归一化（LayerNorm）：性能达到`torch.nn.LayerNorm`的484.4%。

Conv2D+ReLU+MaxPool组合操作：性能达到PyTorch参考实现的290.1%，以及`torch.compile()`参考实现的189.0%。

（基准测试环境为NVIDIA L40S GPU，性能百分比定义为参考运行时间除以生成内核的运行时间）

更惊人的是，这一突破性成果的实现，本身竟是一个“美丽的意外”。研究团队最初的计划，本是生成合成数据用于训练内核生成模型。没想到，在测试阶段，这些原本作为“训练材料”的合成数据，直接生成了性能极其优秀的内核。

消息一出，围观的技术社区立刻热闹起来。有网友调侃道：没想到，AI连内核工程师的饭碗也要开始“惦记”了。

除了性能的大幅跃升，研究团队采用的方法论也颇有趣味。他们并未采用简单的、类似“爬坡”算法的渐进式优化，而是在每次迭代之间，巧妙地加入了一个语言推理步骤。此举旨在鼓励搜索过程产生更多样化的解决思路。简单来说，就是让系统在每次试图改进时，先进行一番“思考”，发散出更多可能性，从而更有机会跳出局部最优，找到更佳的解决方案。

究竟是如何实现的？我们来一探究竟。

改代码前，先生成自然语言优化思想

根据斯坦福团队博客的描述，这种内核生成的基本思路其实很清晰：给定一段PyTorch代码，然后让模型尝试编写自定义CUDA内核来替换其中的算子。这些内核使用纯CUDA-C编写，不依赖CUTLASS、Triton这类专用库或领域特定语言。

关键在于其与众不同的流程：模型并非直接上手修改代码，而是分两步走。首先，用自然语言生成一系列的优化思想和策略；然后，再将每一个思想转化为具体的代码变体。团队解释说，传统“按顺序修改”的优化方式往往缺乏多样性，容易陷入局部最优，要么不断重复同类转换，要么在一条没有前途的路径上耗时费力。

为了进一步激发思路的多样性，团队还采用了多分支探索模式。具体而言，他们的方法不是每一步只盯着一个候选方案优化，而是让每个有价值的想法“开枝散叶”，衍生出多个实现版本，并从中选择性能最高的内核作为下一轮优化的种子。

团队使用OpenAI o3和Gemini 2.5 Pro模型，在自建的KernelBench基准（1级难度）中挑战了10个问题。经过多轮运行后，性能最佳的内核开始显现。有趣的是，大多数最佳结果出现在后续轮次（总5轮中，主要是第4或第5轮）。

这里提到的KernelBench，是斯坦福团队提出的一个用于评估AI生成内核能力的测试基准。其中的任务分为3个级别，1级任务特指单一的原始操作，涵盖了AI计算的基础构建块，例如卷积、矩阵乘法、损失函数、激活函数以及层归一化等。

这一发现，结合此前DeepMind的AlphaEvolve、o3发现Linux零日漏洞等事件，让许多业内人士认为，像Gemini 2.5 Pro和o3这类模型，其复杂推理和解决问题的能力已达到了一个新的层级。

回到斯坦福的项目，在观察模型的生成过程中，可以发现其优化思路开始呈现出与人类工程师经验高度相似的模式：

内存访问优化：专注于提升数据在不同内存层次（全局内存、共享内存、寄存器）间移动的效率，确保访问方式能最大化带宽、最小化冲突。

异步操作与延迟隐藏：通过将低速操作（如全局内存访问）与计算或其他内存传输重叠执行，来“隐藏”这些操作的延迟。

数据类型与精度优化：在可能的情况下，优先使用低精度数据类型（如FP16或BF16），以降低内存带宽需求，提升缓存效率。

计算与指令优化：致力于提升算术计算本身的效率，减少指令数量，或利用特定的硬件指令。

并行性与占用率增强：最大化流多处理器（SM）上的活跃线程数量，以更好地隐藏延迟，提高整体吞吐量。

控制流与循环优化：减少与循环、条件分支以及索引计算相关的开销。

团队还展示了一段具体的优化轨迹。从轨迹中可以看出，并非每一步优化都必然带来速度提升，中间甚至可能出现性能回退。但经过多个步骤的组合与迭代，内核的最终性能实现了大幅飞跃，成功超越PyTorch原生实现。

关于实现细节，有人提出了一个关键问题：AI生成的优化建议，是真的被转化为了对应的代码实现，还是仅仅触发了随机的代码探索？对此，作者回应称，尽管尚未进行更严谨的系统性验证，但在手动检查的多个案例中，生成的CUDA代码与其提出的优化建议大致匹配。这表明，AI并非在盲目地随机尝试，而是在有意识地实践它自己构思的策略。

华人主创团队的意外发现

这项研究由三位作者共同完成：Anne Ouyang、Azalia Mirhoseini以及Percy Liang。

Anne Ouyang目前是斯坦福大学扩展智能实验室的博士生，本科与硕士均毕业于麻省理工学院，并曾任职于英伟达的cuDNN团队。

Percy Liang是斯坦福大学计算机科学副教授兼统计学助理教授，目前担任斯坦福基础模型研究中心（CRFM）主任。他此前曾与李飞飞教授等多位学者合作，共同推进了多项前沿研究工作。

Azalia Mirhoseini是斯坦福大学计算机科学助理教授，也是斯坦福扩展实验室的创始人。她曾在DeepMind、Google Brain以及Anthropic等顶尖机构工作，其研究履历包括混合专家模型（MoE）、芯片设计算法AlphaChip等知名项目。

回到研究本身，团队的初衷本是生成高质量数据，用以训练更强大的内核生成模型。然而，实验过程中间出现了计划外的惊喜：仅在测试阶段生成的合成数据，其本身就足以产生性能卓越的内核。由于这些内核实现了一些此前被认为较难达成的高级优化和硬件特性利用，团队决定以博客形式先行分享这一意外成果。

不过，关于数据生成的具体方法论，研究团队暂未对外公开，仅提及整体设计理念其实相当简洁。最关键的是，现有结果已经充分展示了该方向的巨大潜力。

此外，团队也指出，此次发现与近期AI领域的某些趋势不谋而合——大规模重新训练模型并非总是必需品。有时候，精巧的搜索策略与多分支探索，本身就足以解锁科学创新并解决复杂问题。结合验证器进行广泛搜索，往往能带来额外收获。将强大的推理能力与同时探索多种假设结合起来，已被证明能产生更优的结果，正如AlphaEvolve、AlphaEvolution以及Gemini 2.5 Pro的“深度思考”模式所展示的那样。

当然，团队也坦诚表示，这项研究仍有广阔的优化空间。他们手头就在继续攻关两个维度：FP16精度的矩阵乘法（目前性能为`torch.matmul`的52%）以及FP16精度的Flash Attention（目前性能为`torch.nn.functional.scaled_dot_product_attention`的9%）。相比FP16或BF16，FP32在新硬件上的官方优化通常较少，这或许解释了为何在FP32内核上更容易实现相对于PyTorch的性能提升。

尽管面临不少现有限制，团队对未来前景持乐观态度。毕竟在项目初期，他们连一个能正常运行的内核都难以生成。而通过持续优化搜索方法，如今Flash Attention的性能已被提升到一个相当不错的水平。值得一提的是，整个搜索过程消耗的计算资源并不夸张，大约只使用了300万Token的输入和400万Token的输出。

One More Thing

事实上，探索用大模型生成高性能计算内核的团队并非只有斯坦福一家。

就在今年5月，开发了AI程序员Devin的Cognition公司，开源了首个仅通过强化学习即可编写CUDA内核的大模型Kevin-32B。该模型基于QwQ-32B，在KernelBench数据集上使用GRPO（Group Relative Policy Optimization）算法进行了多轮强化学习，其性能据称优于o3和o4-mini模型。

参考链接：

[1]https://crfm.stanford.edu/2025/05/28/fast-kernels.html

[2]https://x.com/anneouyang/status/1928124885567467768

[3]https://x.com/cognition_labs/status/1919835720493236295