AI 生成优化 Metal 内核，PyTorch 推理速度提升 87%

根据 Gimlet Labs 的最新研究，AI 能够自动生成优化的 Metal 内核，使得 PyTorch 推理速度提升了87%。这一突破性成果不仅提高了性能，还在测试的215个 PyTorch 模块上实现了平均1.87倍的加速，某些工作负载的速度甚至提高了数百倍。

研究人员选取了来自多个顶尖机构的八个 AI 模型，包括 Anthropic、DeepSeek 和 OpenAI，利用这些模型为苹果设备生成优化的 GPU 内核。这一过程无需修改用户代码或使用新的框架，直接在苹果硬件上提升模型性能。

在实验中，研究团队选择了 Mac Studio （搭载 Apple M4Max 芯片） 进行测试，基准设置为 PyTorch 的 eager 模式。实验采用了 KernelBench 数据集中的215个 PyTorch 模块，这些模块被分为三类，涵盖从简单的矩阵乘法到完整的模型架构。

测试过程包括接收输入和 PyTorch 代码，生成 Metal 内核，并评估其正确性。数据显示，随着尝试次数的增加，AI 生成内核的正确性逐步提升。例如，在第五次尝试时，正确实现的比例达到了94%。此外，模型们在生成内核时表现出了跨层级的能力，尽管非推理模型有时也能生成有效内核。

实验结果表明，GPT-5模型在某些任务上实现了4.65倍的速度提升。更令人惊讶的是，o3模型在某些情况下甚至将延迟降低了9000倍。研究还发现，单一模型在某些任务上并不总是表现最好，多个模型的结合能够生成更优的内核。

为了进一步提升性能，研究者尝试引入额外上下文信息，如 CUDA 实现和 gputrace 的性能分析数据，结果显示这种方法在性能加速方面达到了平均1.87倍，相比于普通智能体的1.31倍提升了三倍。

需要注意的是，研究人员强调，这一工作并不是为了展示最终的性能极限，而是为了验证 AI 在内核生成中的可行性，希望通过自动化减少开发人员的负担。整体而言，这项研究标志着 AI 技术在硬件优化领域的一个重要进展。