AI自动生成苹果Metal内核，显著提升PyTorch推理速度

09-05 07:33

AI自动生成的苹果芯片Metal内核，表现比官方的还要出色？Gimlet Labs的最新研究显示，在苹果设备上，AI不仅能够自动生成Metal内核，还让PyTorch推理速度相较于基线内核提升了87%。更令人惊讶的是，AI生成的Metal内核在215个受测的PyTorch模块上，实现了平均1.87倍的加速，部分工作负载的速度甚至比基准快了数百倍。

真的是AI让苹果AI再次伟大？

用AI为苹果设备生成内核

结论是：通过AI自动进行内核优化，无需修改用户代码，也无需新框架或移植，就能够显著提升模型性能。为了验证这一点，研究人员挑选了来自Anthropic、DeepSeek和OpenAI的8个顶尖模型，让它们为苹果设备生成优化的GPU内核，以加快PyTorch推理速度。

为什么选择苹果设备呢？因为它是全球最大的硬件供应商。接下来，看看研究人员的具体做法。

实验设置

在模型选择上，参与测试的模型有claude - sonnet - 4、claude - opus - 4、gpt - 4o、gpt - 4.1、gpt - 5、o3、deepseek - v3、deepseek - r1。测试输入方面，研究采用了KernelBench数据集中定义的PyTorch模块，并选取其中215个模块进行测试。这些模块被划分为三个等级，第一级是简单操作，如矩阵乘法、卷积；第二级是由第一级操作组成的多操作序列；第三级是完整的模型架构，如AlexNet、VGG。评估指标主要关注两个方面，一是AI生成内核的正确性，二是其相较于基准PyTorch的性能提升。研究使用的苹果硬件是Mac Studio (Apple M4 Max chip)，Baseline为PyTorch eager mode。

实验测试

准备工作完成后，研究团队开始测试。测试流程如下：接收提示（prompt）和PyTorch代码，生成Metal内核，评估其在正确性上是否与基准PyTorch一致，如果编译失败或不正确，则将错误信息回传给智能体重试，最多重试5次。实验表明，内核正确性会随着尝试次数的增加而提高。以o3为例，第一次尝试约有60%的概率得到可用实现，到第5次尝试时，可用实现比例达到94%。此外，推理模型很擅长跨层级生成正确的内核，非推理模型有时也能做到。

实验结果令人惊艳，几乎每个模型都生成了比基准更快的内核。例如，GPT - 5在一个Mamba 25状态空间模型上实现了4.65倍的加速，主要通过内核融合减少内核调用开销，并改善内存访问模式。在一些案例中，o3甚至将延迟提升了超过9000倍。总体而言，GPT - 5平均能带来约20%的加速，其他模型稍逊一筹。不过，GPT并非在所有问题上都能生成最优内核，研究人员发现GPT - 5在34%的问题上生成了最优解，而在另外30%的问题上，其他模型生成的解更优。这意味着没有单一模型能在所有问题上都生成最优内核，因此研究人员开展了智能体群体实验。

智能体群体实验

与单个模型相比，智能体群体策略实现了更高的性能提升。与GPT - 5相比，智能体群体在各层级平均加速31%，在Level 2问题上加速42%。在几乎没有上下文信息的情况下，智能体群体表现就已不错。研究人员尝试为智能体提供更多上下文以获取更快的内核，额外信息来源包括CUDA实现和M4上gputrace的性能分析信息。在具体实施中，先将截图处理任务分配给子智能体，让其为主模型提供性能优化提示，主智能体初步实现后进行性能分析和计时，再将截图传给子智能体生成新的性能优化提示。实验表明，上下文配置没有“单一最佳”方案，但加入额外上下文实现了平均1.87倍的加速，相较于普通智能体的1.31倍平均加速，提升幅度提高了三倍。

有提升，但看跟谁比

在PyTorch中，调用函数时，PyTorch会先将其拆解为张量运算，再交给GPU执行，这就需要GPU内核将数学操作转换为GPU可理解的低级并行指令，所以GPU内核性能对运算效率很关键。这项研究让原本需工程师手写的内核优化由AI自动完成并测试性能。然而，苹果硬件不像英伟达的CUDA那样对PyTorch有很好的优化，直接拿MPS后端原生实现和AI生成的内核对比不太公平。有网友指出，文章使用的baseline是eager mode，通常用于训练调试或指标计算，不会真正部署到设备上。在真实部署中，一般会先将模型导出为ONNX，再编译成设备原生格式，效率比直接用PyTorch eager mode高很多。所以，无论内核是手写还是AI生成，优化后的GPU内核都比未优化的PyTorch推理快。研究人员回应称，这项工作不是为了展示部署环境的最终性能极限，而是展示AI自动生成内核的可行性，目的是在内核工程方面获得人类专家一定程度的效益，无需开发人员额外投入，重点在于原型验证。