斯坦福团队意外生成超快AI内核，性能超越PyTorch标准基线

斯坦福大学的研究人员在生成AI优化内核方面取得了令人瞩目的成果。这些内核是通过纯CUDA-C编写而成的，没有使用如CUTLASS和Triton等现有库和领域特定语言（DSL），但其性能却接近或超过了PyTorch中专家手动优化的标准内核。研究团队使用了KernelBench作为基准测试工具，该工具于2024年12月发布。通过OpenAI的o3和Gemini 2.5 Pro模型运行了10个KernelBench Level 1的问题，并进行了多轮搜索和优化。 以下是部分优化内核的性能对比（所有测试均在Nvidia L40S GPU上进行）： 矩阵乘法（Matmul）：101.3%的PyTorch标准性能，问题规模为4096x4096的方阵。 卷积运算（Conv2D）：179.9%的PyTorch标准性能，输入张量大小为(100, 3, 224, 224)，卷积参数设置为in_channels=3, out_channels=96, kernel_size=11, stride=4, padding=2。 Softmax：111.8%的PyTorch标准性能，输入张量大小为(4096, 65536)。 层归一化（LayerNorm）：484.4%的PyTorch标准性能，输入张量大小为(16, 64, 256, 256)。 卷积+ReLU+最大池化（Conv2D + ReLU + MaxPool）：290.1%的PyTorch基准性能，189.0%的PyTorch编译后性能，输入张量大小与上述Conv2D相同，最大池化参数设置为kernel_size=3, stride=2。 研究方法主要包括两个关键改进： 自然语言推理优化策略：在每一步生成新的内核时，先在自然语言中生成优化策略想法，再根据这些想法实现具体代码变体。 多分支并行探索：每一步生成多种不同的内核实现，选择性能最高的内核继续下一轮优化。这种策略允许大规模并行搜索，避免陷入局部最优解。 以Conv2D内核为例，初始版本的运行时间为7.02毫秒，仅相当于PyTorch参考时间的20.1%。通过逐步应用优化策略，到第13轮时，最终时间缩短到0.795毫秒，性能达到PyTorch标准性能的179.9%，显著优于PyTorch原生实现。优化策略包括使用读取缓存、利用半精度张量计算、双缓冲加载管道、预计算共享内存索引以及矢量化共享内存写入等高级技术。 研究人员指出，这种方法的成功在于结合了强大的推理能力和并行探索多个假设的策略，从而实现了显著的性能提升。虽然目前的结果非常有潜力，但仍有许多改进空间，例如生成更高质量的优化策略和代码，以及应用于更为复杂的内核。此外，FP32格式在现代机器学习工作负载中不如FP16或BF16常用，这也是实现性能提升的一个原因。 业内人士对这一研究普遍持积极态度，认为这是人工智能在高性能计算领域的一个重要突破，展示了通过优化搜索和并行策略可以解锁更多的创新可能性。该研究成果对于未来的模型训练也有重要意义，能够生成更好的合成数据来提高模型的效率和准确性。斯坦福大学在这方面具有深厚的研究背景和技术积累，此次发布的成果进一步巩固了其在机器学习和高性能计算领域的领先地位。 感谢Christopher Rinard、Saman Amarasinghe和Allen Nie在讨论中的帮助，以及Standard Kernel Co.和Prime Intellect对这项工作的支持。