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

​AI 生成优化 Metal 内核,PyTorch 推理速度提升惊人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 技术在硬件优化领域的一个重要进展。

github:https://github.com/ScalingIntelligence/KernelBench/

划重点:  

🌟 AI 自动生成 Metal 内核,提升 PyTorch 推理速度87%。  

⚡️ 在215个 PyTorch 模块上实现平均1.87倍的加速,部分工作负载速度提升数百倍。  

🔍 研究旨在验证 AI 在内核生成的可行性,助力硬件优化。