令人惊叹！AI意外生成的内核，性能竟然超越了人类专家精心优化的结果！
斯坦福大学近期公布了一项令人瞩目的研究成果。
由AI算法优化的内核，在常见的深度学习操作中表现出惊人的性能提升，相较于原生PyTorch，性能提升甚至接近400%。具体表现如下：
矩阵乘法（Matmul，FP32）：性能达到PyTorch torch.matmul的101.3%。这一表现与原生实现相当，但AI自动生成的潜力不容小觑。
二维卷积（Conv2D）：性能达到 torch.nn.Conv2D的179.9%。接近两倍的性能提升，体现了AI在优化卷积操作方面的强大能力。
Softmax：性能达到 torch.softmax的111.8%。虽然提升幅度不大，但也显示了AI对softmax操作的优化能力。
层归一化（LayerNorm）：性能达到torch.nn.LayerNorm的484.4%。这一数据格外亮眼，表明AI在层归一化操作的优化上有着显著优势。
Conv2D+ReLU+MaxPool组合操作：性能达到PyTorch参考实现的290.1%，以及torch.compile()参考实现的189.0%。在实际应用场景中，AI优化的内核也表现出了卓越的性能。
（以上基准测试在NVIDIA L40S GPU上进行，性能百分比定义为参考时间除以生成的kernel_size时间）
更令人惊讶的是，这一突破性的成果并非研究团队的最初目标。
最初，他们的目的是生成用于训练内核生成模型的合成数据。
然而， 他们意外地发现， 在测试阶段生成的合成数据本身， 竟然就可以生成性能非常优异的内核。
一些网友戏称：难道AI也要取代内核工程师了吗？ 这一评论也反映了人们对AI在软件开发领域潜力的关注。
除了性能的显著提升，研究团队所采用的方法也颇具创新性：
他们并没有采用传统的操作上逐步优化的方法（类似于爬坡算法），而是在每次迭代之间加入了一个语言推理的步骤。 这种方法鼓励搜索过程更加多样化，更有效地探索解空间。
换句话说，研究人员让AI在每次改进时通过类似“思考”的方式产生更多想法，从而寻找更优的解决方案。 这种引入“思考”的优化方式，模仿了人类工程师在优化代码时的思维过程。
让我们深入了解一下具体是如何实现的。
自然语言引导代码优化
根据斯坦福团队博客的描述，他们采用了一种相对简洁的内核生成思路：给定一段PyTorch代码，然后利用AI编写自定义CUDA内核来替换PyTorch算子。
这些内核完全使用CUDA-C编写，避免了使用CUTLASS和Triton等库及DSL（Domain-Specific Language，领域专用语言）。 这种方式更加灵活，也便于AI进行更深层次的优化。
与传统方法不同，该模型并非直接修改代码，而是先用自然语言生成优化思路，再将这些思路转化为新的代码变体。 这种“先思考，后行动”的方法，借鉴了人类解决问题的模式。
研究团队认为，传统的“按顺序修改”式的优化思路缺乏多样性，容易陷入局部最优解， 导致重复访问同一类转换或无休止地优化没有前景的轨迹。 这一点与遗传算法中避免早熟收敛的思想类似。
为了进一步增强思路的多样性，斯坦福团队还采用了多分支探索模式。
具体而言，他们的方法并非每一步只优化一个候选方案，而是将每个想法分散开来，使其衍生出多个实现，并使用性能最高的内核作为下一轮的种子。 这种并行探索的策略，显著提高了寻找全局最优解的概率。
研究团队使用OpenAI的o3和Gemini 2.5 Pro来挑战KernelBench 1级中的10个问题。经过多轮运行，最佳内核开始显现。
大多数最佳结果出现在后续轮次（总共5轮），特别是第4轮或第5轮，说明优化过程需要一定的迭代累积。
KernelBench 是斯坦福团队提出的AI生成内核测试基准。 该基准中的任务分为3个级别，其中1级是指单一原始操作（Single primitive operation），包括AI的基础构建块，例如卷积、矩阵-向量与矩阵-矩阵乘法、损失函数、激活函数以及层归一化。
这一发现，以及之前DeepMind的AplhaEvolve，以及o3发现Linux的0day漏洞等事件，都暗示着Gemini 2.5Pro和o3的能力水平达到了新的高度。 AI不再仅仅是执行命令的工具，而是具备了解决复杂问题的能力。
回到斯坦福的这个项目，在内核生成过程中，模型的生成思路开始呈现出与人类工程师的经验相似之处：
内存访问优化： 提高不同内存层次结构(全局内存、共享内存、寄存器)之间数据移动的效率，并确保以最大化带宽和最小化冲突的方式访问数据; 这与缓存一致性等优化策略异曲同工。
异步操作和延迟隐藏： 通过将慢速操作(如全局内存访问)与计算或其他内存传输重叠，“隐藏”慢速操作的延迟。 这种技术类似于流水线操作，旨在充分利用硬件资源。
数据类型和精度优化： 尽可能使用低精度数据类型(如 FP16 或 BF16)以减少内存带宽要求并提高缓存效率。 降低精度是提高计算效率的常见手段，但需要在精度和性能之间做出权衡。
计算和指令优化： 提高算术计算本身的效率，减少指令数量，或利用专门的硬件指令。 例如，利用SIMD指令可以显著提高计算效率。
并行性和占用率增强： 最大化流多处理器(SM)上的活动线程数量，以更好地隐藏延迟并提高整体吞吐量。 提高并行度一直是GPU优化的核心目标。
控制流和循环优化： 减少与循环、分支和索引计算相关的开销。 例如，循环展开可以减少循环带来的额外开销。
斯坦福团队还展示了一组具体的优化轨迹。 结果表明， 并非每一步优化都能立即提升速度，但经过多个步骤的组合， 内核的速度能够得到大幅提升，最终超越PyTorch的原生实现。 这印证了优化是一个迭代和累积的过程。
在具体实现上，有人询问AI生成CUDA内核时，其提出的优化建议是否被转化为了实际代码，或者仅仅是触发了随机探索？
作者回应说，虽然尚未进行严格的系统验证，但在手动检查的案例中， 生成的CUDA视线与提出的优化建议大致匹配。 这表明AI并非完全随机地进行优化，而是尝试实现其提出的策略。
华人主创团队的意外发现
这项研究由三位作者共同完成：Anne Ouyang、Azalia Mirhoseini 和 Percy Liang。
Anne Ouyang 目前是斯坦福大学扩展智能实验室的博士生， 她本硕毕业于麻省理工学院，并在英伟达的 cuDNN 团队工作过。 她的经历为该项目带来了深厚的理论基础和实践经验。
Percy Liang 是斯坦福大学计算机科学副教授兼统计学助理教授， 目前担任斯坦福基础模型研究中心主任。 他曾与李飞飞教授合作发布并推进了多项研究工作。
Azalia Mirhoseini 是斯坦福大学计算机科学助理教授， 也是斯坦福扩展实验室的创始人。 她曾在DeepMind、Google Brain 和 Anthropic 工作过。 其此前参与的研究包括 MoE、芯片设计算法 AlphaChip 等， 拥有丰富的AI研究经验。
这项研究最初的目的是生成用于训练内核生成模型的合成数据。
然而，在研究过程中， 团队却意外发现仅在测试阶段生成的合成数据本身， 竟然可以生成性能非常优秀的内核。 这一意外发现为AI在内核优化领域的应用打开了新的思路。
由于这些内核利用了先前被认为难以实现的高级优化和硬件特性， 研究团队决定以博客形式分享此次成果， 希望能与社区共同探讨。
不过，研究团队暂时没有公布具体的数据生成方法， 只是提到其设计理念非常简单，并且已经展示出了巨大的潜力。 这可能涉及到一些尚未公开的技术细节， 或许会在后续的研究中揭晓。
此外，研究团队认为此次发现与最近的趋势相符： 大规模的预训练并非总是必需的。 相反，聪明的搜索和分支策略可以解锁科学创新并解决复杂问题，通过验证器进行广泛搜索还能有更多收获。
将强大的推理能力与同时探索多个假设相结合，能够带来更好的结果，正如AlphaEvolve、AlphaEvolution 和 Gemini 2.5 Pro 的深度思考模式一样。 这一观点强调了AI在解决复杂问题时，进行推理和探索的重要性。
最后，团队表示这项研究还有很多可优化的空间。 例如，他们正在优化以下两个方面：
FP16 Matmul：性能达到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的输出， 表明该方法具有较高的效率。
更多探索
实际上，尝试开发内核大模型的团队不止斯坦福一个。
早在今年5月，开发了Devin的Cognition公司就开源了首个通过强化学习编写CUDA内核的大模型Kevin-32B。
该模型基于QwQ-32B，在KernelBench数据集上使用GRPO进行多轮强化学习，性能优于o3和o4-mini， 表明强化学习在内核优化方面也具有广阔的应用前景。
总而言之，斯坦福大学的这项研究为AI在内核优化领域的应用提供了新的视角和方法。 这一意外的发现，不仅展示了AI在解决复杂问题上的强大能力，也为未来的软件开发模式带来了新的可能性。 随着AI技术的不断发展，我们有理由期待AI在内核优化，乃至整个软件工程领域发挥更大的作用。
[1]
[2]
[3]
