摘要：本想练练手合成点数据，没想到却一不小心干翻了PyTorch专家内核！斯坦福华人团队用纯CUDA-C写出的AI生成内核，瞬间惊艳圈内并登上Hacker News热榜。团队甚至表示：本来不想发这个结果的。
编辑：Aeneas 好困
【新智元导读】本想练练手合成点数据，没想到却一不小心干翻了PyTorch专家内核！斯坦福华人团队用纯CUDA-C写出的AI生成内核，瞬间惊艳圈内并登上Hacker News热榜。团队甚至表示：本来不想发这个结果的。
就在刚刚，斯坦福HAI华人大神团队又出惊人神作了。
他们用纯CUDA-C语言编写的快速AI生成内核，竟然超越了PyTorch！
在这个过程中，完全不用借助CUTLASS和Triton等库和领域特定语言（DSL），就能让性能表现接近PyTorch内置的、经过专家优化的标准生产级内核，甚至在某些情况下还更胜一筹。
作者团队都是我们熟悉的名字——Anne Ouyang、Azalia Mirhoseini和Percy Liang，有趣的是，他们甚至直言，这个结果其实本不想拿出来发布。
一经发布，这个发现就引爆了技术圈，现在已经登顶Hacker News总榜第二。
说起来，这个发现还有很多意外的成分。
本来，他们的目标是生成合成数据，来训练更好的内核生成模型，合成数据生成的设计也十分简单。
然而，意想不到的事情发生了，仅用于测试的合成数据生成本身，竟开始生成非常优秀的内核，甚至超越了人类专家优化的PyTorch基线，而且还利用了高级优化和硬件特性。
而在此前，这是一项很艰难的挑战。
由此，研究者们决定提前撰写博文，把自己的发现分享出来。
总结来说，研究的亮点成果如下：
矩阵乘法（Matmul, FP32）：性能达到PyTorch FP32 torch.matmul的101.3%二维卷积（Conv2D, FP32）：性能达到PyTorch FP32 torch.nn.Conv2D的179.9%Softmax（FP32）：性能达到PyTorch FP32 torch.softmax的111.8%层归一化（LayerNorm, FP32）：性能达到PyTorch FP32 torch.nn.LayerNorm的484.4%二维卷积 + ReLU + 最大池化（Conv2D + ReLU + MaxPool, FP32）：性能达到PyTorch FP32参考实现的 290.1%，达到PyTorch FP32 torch.compile参考实现的189.0%以上结果在英伟达L40S GPU上进行了基准测试，性能百分比定义为参考时间除以生成的内核时间。
网友：强制LLM推理，实在太有趣了
在Hacker News上，网友们也对此展开了热烈讨论。
比如为什么使用FP32内核会比PyTorch更容易实现性能提升，理由就相当有趣。
如果AI真的能以更低成本，实现更优化的内核，的确潜力巨大。
最令人震撼的就是，无论是最近谷歌的AlphaEvolve，还是o3在Linux内核中发现了零日漏洞，都在提醒我们——
Gemini Pro 2.5和o3已经达到了一个全新的能力水平，那些曾经在其他模型上尝试失败的想法，现在突然奏效了。
可以说，我们已经到达了一个节点，LLM能比用人类快得多的速度进行迭代和测试，信息组合、进步和智能应用的蛮力，似乎正在成功！
接下来，我们来看看斯坦福研究者们博客中的具体内容。
博客全文
在博客中，研究者分享了具体方法、五个优化后的内核（包括4个基础机器学习算子和1个AlexNet模块的融合内核）、一个优化过程的实例，以及一些思考，关于这些发现对高性能内核生成可能意味着什么。
可以说，这些内容将是他们后续探索的第一步。
方法
研究者们采用了KernelBench的任务设置（这是他们在2024年12月发布的一款基于AI的内核生成基准测试）。
具体来说，给定一段torch代码，LLM会编写自定义内核来替换原有的torch算子，目标是实现加速。
依照KernelBench最初的设计，参考代码默认使用FP32精度；在给定的容差阈值（1e-02）下，采用较低精度的解决方案也是被允许的。
此外，由于存在大量针对特定规模的优化手段，KernelBench中的每个问题都设定了具体的输入大小。
因此，该基准测试旨在找出针对特定问题规模的最快内核，而非一个适用于任意问题规模的高速内核。
而且，研究者会同时运行torch参考代码和生成的代码，并通过在多种随机输入下比较两者输出的数值是否一致，来检验其正确性。
当前，在优化内核这个问题上，业界扩展测试时计算资源最常用的方法是顺序修订（sequential revision）。
这是一种多轮迭代的循环：模型首先对内核进行增量式修改，接着检查其正确性和性能，然后根据结果再次尝试。
也就是说，要么修复有问题的内核，要么进一步提升现有内核的性能。
这个循环过程非常直观，也容易实现。模型会修复失效的内核，微调可用的内核，一步步优化出性能更佳的版本。
这种方法的主要局限，在于优化思路缺乏多样性。
顺序循环往往容易陷入局部最优的困境，比如反复尝试同类型的转换，或是在缺乏潜力的优化路径上无休止地调整。
其结果便是测试时计算资源的低效利用，并且难以促使模型产生具有根本性创新的优化思路。
为解决这一问题，研究者引入了两项关键改变：
运用自然语言对优化思路进行推理他们不再于每一步直接生成新的内核，而是以先前尝试过的思路为条件，用自然语言生成优化思路，随后将这些思路具化为新的代码变体。
在每个优化步骤进行分支扩展他们不是每步只改进一个候选方案，而是进行分支扩展，让每个思路都能派生出多种实现版本，其中性能最佳的内核将作为下一轮优化的种子。
（研究者也会保留一个表现优异的现有内核库，用于提供种子）。
这种方式解锁了大规模的并行处理能力，使他们能够在每一轮探索截然不同的优化方向，避免陷入狭窄的优化路径。
其结果是，这种测试时循环不再像顺序修订那般，仅仅是与编译器「对话」，而是更接近一种结构化的探索性搜索。
这种搜索由明确的优化假设指导，并采用大规模并行评估的方式进行。
研究者运行了KernelBench第1级的10个问题，以进行测试。
他们调整了问题规模，以确保内核启动开销相对于问题的整体运行时间而言可以忽略不计。
然后，使用OpenAI o3和Gemini 2.5 Pro模型进行了5轮实验。
下图展示了首次发现性能最佳内核所在的轮次分布情况。
可以看到，大多数最优结果出现在靠后的轮次（总共5轮），其中绝大部分出现在第4轮或第5轮。
随着扩大搜索范围，研究者还发现：许多高性能内核的优化策略高度相似，集中在少数几种常见的模式上，这与他们手动编写内核的经验也是一致的。
主要的优化类别归纳如下——
内存访问优化：提升不同内存层级（全局内存、共享内存、寄存器）之间数据迁移的效率，并确保数据访问方式能够最大化带宽、最小化冲突。异步操作与延迟隐藏：通过将耗时较长的操作（例如全局内存访问）与计算或其他内存传输重叠执行，来隐藏其带来的延迟。数据类型与精度优化：在允许的条件下，尽可能使用较低精度的数据类型（如FP16或BF16），以降低内存带宽需求，提升缓存效率，并有望利用专门的硬件加速单元。计算与指令优化：提升算术运算本身的效率，削减指令数量，或利用专门的硬件指令。并行性与占用率增强：最大化流式多处理器（SM）上活跃线程束（warp）的数量，以便更好地隐藏延迟，提高整体吞吐率。控制流与循环优化：减少由循环、分支及索引计算等引入的额外开销。总结
这次研究者采用的方法，与AI研究中一个日益显著的趋势不谋而合——
将强大的推理能力与对多个假设的并行探索相结合，能够带来性能的提升。
正如一些近期研究（例如AlphaEvolve、Gemini 2.5 Pro Deep Think）所强调的，我们并不总是需要大规模的重新训练。
论文地址：https://storage.googleapis.com/deepmind-media/DeepMind.com/Blog/alphaevolve-a-gemini-powered-coding-agent-for-designing-advanced-algorithms/AlphaEvolve.pdf
有时，巧妙的搜索和分支策略便足以催生科学创新、攻克复杂难题，而借助验证器进行广泛搜索，则可能带来更大的收益。
然而，这并不意味着我们不需要进一步的训练。
恰恰相反，研究者的这种方法，也有助于生成更优质的合成数据，用以改进未来的模型训练（这需要更多的问题实例）。
因此，它既是一种强大的测试时扩展方法，也是我们迈向更智能、数据效率更高的模型开发之路的一步。
而且，这次研究者展现的仅仅是初步的成果。这些优化结果的质量看起来相当可观，但仍有广阔的提升空间，例如产生更优的优化思路、生成更高质量的最终代码，以及将此方法应用于日益复杂的内核。
目前，研究者仍在积极改进的两个具体例子包括：
FP16 Matmul：性能达到torch.matmul的52%FP16 Flash Attention：性能达到torch.nn.functional.scaled_dot_product_attention的9%在现代机器学习任务中，FP32的应用不如FP16或BF16普遍，并且在较新的硬件上，针对FP32的优化往往也更少。
这或许能部分解释，为何基于FP32的内核更容易在性能上超越PyTorch。
作者介绍
Anne Ouyang
Anne Ouyang目前是斯坦福大学计算机科学（CS）博士生，在Scaling Intelligence Lab（可扩展智能实验室）进行研究。
她的研究兴趣主要集中在可扩展的自我改进机器学习系统，同时也广泛关注实证机器学习（empirical ML）和性能工程（performance engineering）。
此前，她在MIT获得学士和硕士学位，并曾在NVIDIA cuDNN团队工作，负责编写CUDA内核，用于加速GPU上的深度学习工作负载。
Azalia Mirhoseini
Azalia Mirhoseini是斯坦福大学计算机科学助理教授，也是Scaling Intelligence Lab（可扩展智能实验室）的创始人，并在Google DeepMind兼任高级研究科学家。
她的实验室致力于开发可扩展的自主演进人工智能系统与方法论，以期推动通用人工智能的发展。
在加入斯坦福大学之前，她曾在Google Brain和Anthropic等业界顶尖的人工智能实验室工作多年。
她过往的卓越成就包括：
提出混合专家（MoE）神经架构——目前已被前沿的AI模型广泛应用；领导AlphaChip项目——一项将深度强化学习用于布局优化的开创性工作，并成功应用于谷歌AI加速器（TPU）及数据中心CPU等先进芯片的设计中；在测试时计算的Scaling方面有深入的研究Percy Liang
Percy Liang是斯坦福大学计算机科学副教授，兼任基础模型研究中心（CRFM）主任。同时也是CodaLab Worksheets的创建者，并借此坚定倡导科研工作的可复现性。
他目前专注于通过开源和严格的基准测试，提升基础模型（特别是大语言模型）的可及性与可理解性。
他曾围绕机器学习和自然语言处理领域进行了广泛研究，具体方向包括鲁棒性、可解释性、人机交互、学习理论、知识落地、语义学以及推理等。
此前，他于2004年在MIT获得学士学位，并于2011年在UC伯克利获得博士学位。
内核优化过程
最后，展示一个Conv2D自动生成思路的优化轨迹示例，torch参考基准时间为1.41毫秒。
第0轮：7.02毫秒，达到参考性能的20.1%
思路：给定pytorch代码，用CUDA Kernel替换操作。
第1轮：7.54毫秒，达到参考性能的18.8%
思路：通过使用__ldg加载不变的张量来利用只读缓存。
第2轮：3.46毫秒，达到参考性能的41.0%
思路：将卷积转换为FP16 Tensor-Core GEMM。
作者评论：这是一种算法优化，将卷积转换为隐式GEMM，这对于在Tensor Core上高效运行卷积至关重要
第3轮：3.67毫秒，达到参考性能的38.7%
思路：采用双缓冲cp.async管线，使全局内存加载与Tensor-Core计算重叠。
第4轮：3.46毫秒，达到参考性能的41.0%
思路：给定pytorch代码，使用隐式矩阵乘法（implicit matmul）的CUDA Kernel替换操作。给定的GEMM内核可能会有帮助。
作者评论：因为优化涉及到使用GEMM，所以在这一轮开始时，使用了一个之前生成的现有优秀GEMM内核作为种子，这个思路是手动编写的
第5轮：1.91毫秒，达到参考性能的74.9%
思路：在每个K-tile循环内部，预计算k_idx分解的内核/输入索引并在共享内存中重用，以避免冗余算术运算。
第6轮：1.37毫秒，达到参考性能的103.6%
思路：在共享内存中预计算并缓存N维GEMM索引，以减少输入数据加载循环内的冗余算术运算。
第7轮：1.38毫秒，达到参考性能的102.9%
思路：通过使用专用的每warp共享内存缓冲区来并行化CUDA内核输出，以消除写入全局内存时的warp级串行化。
第8轮：1.37毫秒，达到参考性能的103.6%
思路：在共享内存中预计算并缓存基础输入坐标，以减少输入数据加载过程中的冗余算术运算。
第9轮：1.36毫秒，达到参考性能的105.1%
思路：对B-fragment加载进行软件流水线处理，使下一个B-tile的共享内存读取与当前B-tile的WMMA计算重叠。
第10轮：1.07毫秒，达到参考性能的133.6%
思路：重用共享内存中预计算的N维GEMM分解来进行输出地址计算，避免冗余且高成本的除法/模运算。
第11轮：1.21毫秒，达到参考性能的117.4%
思路：在half WMMA运算中移除hi/lo分解，如果最终精度可接受，则依赖标准的FP16累积来提高性能。
第12轮：1.01毫秒，达到参考性能的141.2%
思路：通过在流水线的加载阶段动态计算K维索引，使用双缓冲使Asub（权重）和Bsub（输入）的K循环全局内存加载与MMA计算重叠。
第13轮：0.795毫秒，达到参考性能的179.9%
思路：通过使用像half2这样的更宽数据类型，为加载Asub_pipe和Bsub_pipe实现向量化的共享内存写入。
最终代码
最终生成的Conv2D内核代码，使用了先进的CUDA技术，就是人类自己写起来都很有挑战性的那种！
来源：新智元一点号
