AI自動生成的蘋果芯片Metal內(nèi)核，比官方的還要好？

Gimlet Labs的最新研究顯示，在蘋果設(shè)備上，AI不僅能自動生成Metal內(nèi)核，還較基線內(nèi)核實(shí)現(xiàn)了87%的PyTorch推理速度提升。

更驚人的是，AI生成的Metal內(nèi)核還在測試的215個PyTorch模塊上實(shí)現(xiàn)了平均1.87倍的加速，其中一些工作負(fù)載甚至比基準(zhǔn)快了數(shù)百倍。

真就AI Make蘋果AI Great Again？

用AI為蘋果設(shè)備生成內(nèi)核

先說結(jié)論：通過AI自動實(shí)現(xiàn)內(nèi)核優(yōu)化，可以在無需修改用戶代碼、無需新框架或移植的情況下，顯著提升模型性能。

為了證明這一點(diǎn)，研究人員選取了來自Anthropic、DeepSeek和OpenAI的8個頂尖模型，讓它們?yōu)樘O果設(shè)備生成優(yōu)化的GPU內(nèi)核，以加速PyTorch推理速度。

至于為什么是蘋果？別問——問就全球最大硬件供應(yīng)商（doge）

接下來，讓我們看看研究人員是怎么做的：

實(shí)驗(yàn)設(shè)置

首先，在模型選擇方面，參與測試的模型包括：claude-sonnet-4、claude-opus-4；gpt-4o、gpt-4.1、gpt-5、o3；deepseek-v3、deepseek-r1。

其次，在測試輸入方面，研究使用了KernelBench數(shù)據(jù)集中定義的PyTorch模塊，并選取了其中215個模塊進(jìn)行測試。

這些被選取的模塊被劃分為三個等級，分別是第一級的簡單操作（如矩陣乘法、卷積）；第二級是由第一級操作組成的多操作序列；第三級是完整的模型架構(gòu)（如 AlexNet、VGG）。

再次，在評估指標(biāo)方面，研究人員主要關(guān)注兩個指標(biāo)：一是AI生成內(nèi)核的正確性，二是其相較于基準(zhǔn)PyTorch的性能提升。

最后，研究使用的蘋果硬件為Mac Studio (Apple M4 Max chip)，Baseline為PyTorch eager mode（劃重點(diǎn)，一會要考）

實(shí)驗(yàn)測試

在上述準(zhǔn)備完畢后，研究團(tuán)隊展開了測試。

測試流程如下：

• 接收提示（prompt）和PyTorch代碼；
• 生成 Metal 內(nèi)核；
• 評估其是否在正確性（correctness^4）上與基準(zhǔn)PyTorch一致；
• 如果編譯失敗或不正確，則將錯誤信息回傳給智能體重試，最多允許重試5次。

如上所說，研究者首先關(guān)注AI生成內(nèi)核的正確性。

實(shí)驗(yàn)表明，正確性會隨著嘗試次數(shù)的增加而提升。以o3為例：第一次嘗試就有約60%的概率得到可用實(shí)現(xiàn)，到第5次嘗試時可用實(shí)現(xiàn)比例達(dá)到94%。

此外，研究還發(fā)現(xiàn)推理模型非常擅長跨層級生成正確的內(nèi)核，盡管非推理模型有時也能做到這一點(diǎn)。

那么，AI生成的內(nèi)核表現(xiàn)如何呢？

實(shí)驗(yàn)結(jié)果相當(dāng)驚艷，幾乎每個模型都生成了一些比基準(zhǔn)更快的內(nèi)核。

例如，GPT-5在一個Mamba 25狀態(tài)空間模型上實(shí)現(xiàn)了4.65倍的加速，其主要通過內(nèi)核融合（kernel fusion） 來減少內(nèi)核調(diào)用的開銷，并改善內(nèi)存訪問模式。

在一些案例中，o3甚至將延遲提升了超過9000倍！

總體而言，GPT-5平均可以帶來約20%的加速，其他模型則落后。

不過，GPT并非是門門最優(yōu)，研究人員發(fā)現(xiàn)GPT-5在34%的問題上生成了最優(yōu)解。

但在另外30%的問題上，其他模型生成的解比GPT-5更優(yōu)！

這就意味著沒有單一模型能在所有問題上都生成最優(yōu)內(nèi)核。

因此，如果把多個模型組合起來，就能更大概率生成最優(yōu)內(nèi)核。

于是乎，研究人員又展開了智能體群體實(shí)驗(yàn)（Agentic Swarm）。

智能體群體實(shí)驗(yàn)

果不其然，相較于單個模型，智能體群體策略實(shí)現(xiàn)了更高的性能提升。

與GPT-5相比，智能體群體在各層級平均加速31%，在Level 2問題上加速42%。

在幾乎沒有上下文信息的情況下（僅有輸入問題和提示），智能體群體就已經(jīng)表現(xiàn)得相當(dāng)不錯。

接下來，研究人員嘗試為智能體提供更多上下文，以獲取更快的內(nèi)核。

這里主要包含兩個額外的信息來源：

• CUDA實(shí)現(xiàn)（由于 Nvidia GPU的普及，通?？梢垣@得優(yōu)化過的CUDA參考實(shí)現(xiàn)）；
• M4上gputrace 的性能分析信息。(包含Apple Script捕獲的gputrace摘要、內(nèi)存和時間線視圖)

在具體的實(shí)施步驟中，研究者先將截圖處理任務(wù)分配給一個子智能體（subagent），讓它為主模型提供性能優(yōu)化提示。

在收到提示后，主智能體先進(jìn)行一次初步實(shí)現(xiàn)，然后對其進(jìn)行性能分析和計時。

隨后，再將截圖傳給子智能體以生成性能優(yōu)化提示。

實(shí)驗(yàn)表明，在上下文配置方面也沒有所謂的“單一最佳”方案。

不過，在具體的性能加速方面，加入這些額外上下文實(shí)現(xiàn)了平均1.87倍的加速，相較于普通智能體僅實(shí)現(xiàn)的1.31倍的平均加速，額外上下文將提升幅度提高了三倍！

有提升，但看跟誰比

為了更深入地討論，我們有必要先回顧一些背景知識。

在PyTorch中，我們通常會調(diào)用如Sequential、ReLU這樣的函數(shù)。

在具體的執(zhí)行中，PyTorch會先將函數(shù)拆解為張量運(yùn)算（矩陣乘法、加法等），再交給GPU執(zhí)行。

這時就需要GPU內(nèi)核（kernel）負(fù)責(zé)把這些數(shù)學(xué)操作轉(zhuǎn)成GPU可理解的低級并行指令。

因此，在某種程度上，我們可以說GPU內(nèi)核就像C編譯器一樣，其性能對于運(yùn)算效率至關(guān)重要。

而上面這篇工作所做的，就是讓原本必須由工程師手寫的內(nèi)核優(yōu)化交給AI自動完成，并測試它的性能。

不過，問題就來了。

眾所周知，蘋果硬件并不像英偉達(dá)的CUDA一樣，對PyTorch有很好的優(yōu)化。

因此，這篇研究直接拿MPS后端原生實(shí)現(xiàn)和AI生成的內(nèi)核對比是有失公允的。

不少眼尖的網(wǎng)友也是發(fā)現(xiàn)并指出了這一點(diǎn)：文章里所用的baseline是eager mode，這通常只用于訓(xùn)練調(diào)試或指標(biāo)計算，不會被真正部署到設(shè)備上。

在真實(shí)部署中，一般會先把模型導(dǎo)出為ONNX，再編譯成設(shè)備原生格式（Metal、CUDA 或 ROCm 等），這樣效率會比直接用PyTorch eager mode高很多。

所以，無論內(nèi)核是工程師手寫，還是AI自動生成，經(jīng)過優(yōu)化的GPU內(nèi)核都會比未優(yōu)化的PyTorch推理快得多。

因此，拿調(diào)試過的內(nèi)核和eager比，多少有點(diǎn)奇怪。

對此，研究人員回應(yīng)道：

這篇工作不是為了展示部署環(huán)境的最終性能極限，而是展示AI自動生成內(nèi)核的可行性。

研究的目的是在內(nèi)核工程方面獲得人類專家一定程度的效益，而無需開發(fā)人員的額外投入，希望通過A將部分流程自動化。

所以，重點(diǎn)不在于性能提升，而在原型驗(yàn)證。

對此，你怎么看？