斯坦福大学AI优化研究突破：深度学习助力生成超强CUDA内核，性能提升显著，华人主创团队成果惊人

斯坦福大学的一项新研究揭示了AI生成的内核在性能上超越了人类专家优化的结果。这些由AI优化的内核在常见的深度学习操作中表现卓越，最高提升了近400%的性能。

具体来说，在NVIDIA L40S GPU上的基准测试显示：
– 矩阵乘法（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%。

这一突破是研究团队在尝试生成合成数据以训练内核生成模型时意外发现的。他们原本的目标并非直接生成高性能内核，但结果却令人惊喜。

此外，研究团队采用了一种新颖的方法来增强搜索过程的多样性。他们在每次迭代之间加入了一个语言推理步骤，鼓励系统通过类似“思考”的方式产生更多想法，从而找到更好的解决方案。

斯坦福团队的具体方法包括先用自然语言生成优化思想，然后将这些思想转化为新的代码变体。这种方法避免了传统“按顺序修改”式的优化思路，防止陷入局部极小值。

为了进一步增强思路多样性，团队使用了多分支的探索模式。他们并未每一步只优化一个候选方案，而是让每个想法衍生出多个实现，并选择性能最高的内核作为下一轮的种子。

斯坦福团队使用OpenAI o3和Gemini 2.5 Pro挑战KernelBench 1级中的10个问题，经过多轮运行后，最佳内核开始出现。大多数最佳结果出现在后续轮次，主要是第4轮或第5轮。

这项研究由三位华人主创完成：Anne Ouyang、Azalia Mirhoseini和Percy Liang。Ouyang目前是斯坦福大学扩展智能实验室的博士生，曾就职于英伟达cuDNN团队。Percy Liang是斯坦福大学计算机科学副教授兼统计学助理教授，担任斯坦福基础模型研究中心主任。Azalia Mirhoseini是斯坦福大学计算机科学助理教授及斯坦福扩展实验室创始人。

研究团队认为此次发现与最近的一些趋势相呼应，即大规模再训练并非必需。有时聪明的搜索和分支策略可以解锁科学创新并解决复杂问题。通过结合强大推理能力和同时探索多个假设，能带来更优结果。

尽管取得了显著成果，团队仍认为还有许多可优化的空间。例如，他们正在优化两个维度：FP16 Matmul和FP16 Flash Attention。与FP16或BF16相比，FP32在新推出硬件上的优化程度通常较低，这也是为何使用FP32内核比PyTorch更容易实现性能提升的原因。

值得注意的是，不只是斯坦福一个团队在尝试开发内核大模型。今年5月，Cognition开源了首个通过强化学习即可编写CUDA内核的大模型Kevin-32B。它基于QwQ-32B在KernelBench数据集上使用GRPO，实现了多轮强化学习，性能优于o3、o4-mini。

本文来源：