优化Metal内核通过AI显著提升PyTorch推理效率达87%
发布时间：2025-09-05 09:05:43 | 责任编辑：张毅 | 浏览量：48 次
在苹果设备上，AI 技术正展现出惊人的潜力。根据 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 技术在硬件优化领域的一个重要进展。
github:https://github.com/ScalingIntelligence/KernelBench/
划重点:
🌟 AI 自动生成 Metal 内核，提升 PyTorch 推理速度87%。
⚡️ 在215个 PyTorch 模块上实现平均1.87倍的加速，部分工作负载速度提升数百倍。
🔍 研究旨在验证 AI 在内核生成的可行性，助力硬件优化。
。
关键细节如下
-
实验设置
-
参测模型：Claude-4系列、GPT-4o/4.1/5、o3、DeepSeek-V3/R1等8款最新大模型
-
硬件/软件：Apple M4 Max + macOS Metal，PyTorch eager做baseline
-
测试集：KernelBench的三级算子/子图/完整网络共215题
-
-
生成流程
-
给模型一段PyTorch代码 + 简短prompt，要求输出Metal kernel
-
编译→正确性校验，失败可把报错回传模型，最多retry 5次
-
o3首次通过率约60%，第5次达到94%
-
-
性能结果
-
单模型：GPT-5平均约1.2×，o3在个别Mamba层跑出9000+×
-
多模型“agent swarm”：平均1.31×；再喂入CUDA参考与M4 GPU-trace后，平均1.87×，对应整体87%耗时削减
-
最实用收益集中在Level-2多算子融合与内存访问优化
-
-
技术亮点
-
kernel fusion：把多个小kernel合并，减少调度与显存往返
-
自动贴合M4的tile/size限制，手写汇编级细节由LLM一次性生成
-
无需改用户代码、无需新框架，直接替换.msl文件即可
-
-
注意点
-
baseline是PyTorch eager而非生产级图编译/ONNX，所以87%属于“原型验证”幅度，落地到已充分优化的Metal backend后增益会收窄
-
研究侧重“AI能否写出可正确、可提速的内核”，并非宣告所有场景都能再快87%
-
。
本网站(https://aigc.izzi.cn)刊载的所有内容，包括文字、图片、音频、视频等均在网上搜集。
访问者可将本网站提供的内容或服务用于个人学习、研究或欣赏，以及其他非商业性或非盈利性用途，但同时应遵守著作权法及其他相关法律的规定，不得侵犯本网站及相关权利人的合法权利。除此以外，将本网站任何内容或服务用于其他用途时，须征得本网站及相关权利人的书面许可，并支付报酬。
本网站内容原作者如不愿意在本网站刊登内容，请及时通知本站，予以删除。
