斯坦福意外用AI生成超强CUDA内核，性能比人类专家优化得还要好！翻倍碾压原生PyTorch，华人主创

369 0 0

文章摘要

斯坦福团队意外发现AI生成的内核性能超越人类专家优化成果，展现出AI在内核生成领域的巨大潜力。

AI生成内核性能卓越：由AI优化的内核在常见深度学习操作上表现亮眼，在NVIDIA L40S GPU上进行基准测试，矩阵乘法性能达PyTorch torch.matmul的101.3%，二维卷积达torch.nn.Conv2D的179.9%，Softmax达torch.softmax的111.8%，层归一化达torch.nn.LayerNorm的484.4%，Conv2D + ReLU + MaxPool组合操作达PyTorch参考实现的290.1%和torch.compile()参考实现的189.0%。这一成果是意外实现的，原本目标是生成合成数据训练内核生成模型，测试阶段生成的合成数据却生成了高性能内核。

独特的内核生成思路：该内核生成思路是给定torch代码，编写自定义内核替换torch算子，用纯CUDA – C编写，无需特定库和DSL。模型先用自然语言生成优化思想，再转化为新代码变体，以增强思路多样性。团队使用多分支探索模式，分散想法衍生多个实现，用性能最高的内核作下一轮种子。使用OpenAI o3和Gemini 2.5 Pro挑战KernelBench 1级问题，多数最佳结果出现在后续轮次。

生成思路与人类经验相似：模型生成思路与人类经验相似，包括内存访问优化、异步操作和延迟隐藏、数据类型和精度优化、计算和指令优化、并行性和占用率增强、控制流和循环优化等。生成过程虽非每步都提升速度，但多步骤组合可大幅提升内核速度，超越PyTorch。

华人主创团队及未来展望：研究由Anne Ouyang、Azalia Mirhoseini和Percy Liang三位华人主创。此次发现呼应了无需大规模再训练的趋势，聪明的搜索和分支策略可解锁创新。研究还有优化空间，如FP16 Matmul和FP16 Flash Attention性能有待提升，但团队对未来前景乐观。此外，不止该团队在开发内核大模型，5月Cognition开源了通过强化学习编写CUDA内核的大模型Kevin – 32B，性能优于o3、o4 – mini。