AI自動(dòng)生成的蘋果芯片Metal內(nèi)核，表現(xiàn)比官方的還要出色？Gimlet Labs的最新研究顯示，在蘋果設(shè)備上，AI不僅能夠自動(dòng)生成Metal內(nèi)核，還讓PyTorch推理速度相較于基線內(nèi)核提升了87%。更令人驚訝的是，AI生成的Metal內(nèi)核在215個(gè)受測(cè)的PyTorch模塊上，實(shí)現(xiàn)了平均1.87倍的加速，部分工作負(fù)載的速度甚至比基準(zhǔn)快了數(shù)百倍。

真的是AI讓蘋果AI再次偉大？

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

結(jié)論是：通過(guò)AI自動(dòng)進(jìn)行內(nèi)核優(yōu)化，無(wú)需修改用戶代碼，也無(wú)需新框架或移植，就能夠顯著提升模型性能。為了驗(yàn)證這一點(diǎn)，研究人員挑選了來(lái)自Anthropic、DeepSeek和OpenAI的8個(gè)頂尖模型，讓它們?yōu)樘O果設(shè)備生成優(yōu)化的GPU內(nèi)核，以加快PyTorch推理速度。

為什么選擇蘋果設(shè)備呢？因?yàn)樗侨蜃畲蟮挠布?yīng)商。接下來(lái)，看看研究人員的具體做法。

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

在模型選擇上，參與測(cè)試的模型有claude - sonnet - 4、claude - opus - 4、gpt - 4o、gpt - 4.1、gpt - 5、o3、deepseek - v3、deepseek - r1。測(cè)試輸入方面，研究采用了KernelBench數(shù)據(jù)集中定義的PyTorch模塊，并選取其中215個(gè)模塊進(jìn)行測(cè)試。這些模塊被劃分為三個(gè)等級(jí)，第一級(jí)是簡(jiǎn)單操作，如矩陣乘法、卷積；第二級(jí)是由第一級(jí)操作組成的多操作序列；第三級(jí)是完整的模型架構(gòu)，如AlexNet、VGG。評(píng)估指標(biāo)主要關(guān)注兩個(gè)方面，一是AI生成內(nèi)核的正確性，二是其相較于基準(zhǔn)PyTorch的性能提升。研究使用的蘋果硬件是Mac Studio (Apple M4 Max chip)，Baseline為PyTorch eager mode。

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

準(zhǔn)備工作完成后，研究團(tuán)隊(duì)開(kāi)始測(cè)試。測(cè)試流程如下：接收提示（prompt）和PyTorch代碼，生成Metal內(nèi)核，評(píng)估其在正確性上是否與基準(zhǔn)PyTorch一致，如果編譯失敗或不正確，則將錯(cuò)誤信息回傳給智能體重試，最多重試5次。實(shí)驗(yàn)表明，內(nèi)核正確性會(huì)隨著嘗試次數(shù)的增加而提高。以o3為例，第一次嘗試約有60%的概率得到可用實(shí)現(xiàn)，到第5次嘗試時(shí)，可用實(shí)現(xiàn)比例達(dá)到94%。此外，推理模型很擅長(zhǎng)跨層級(jí)生成正確的內(nèi)核，非推理模型有時(shí)也能做到。

實(shí)驗(yàn)結(jié)果令人驚艷，幾乎每個(gè)模型都生成了比基準(zhǔn)更快的內(nèi)核。例如，GPT - 5在一個(gè)Mamba 25狀態(tài)空間模型上實(shí)現(xiàn)了4.65倍的加速，主要通過(guò)內(nèi)核融合減少內(nèi)核調(diào)用開(kāi)銷，并改善內(nèi)存訪問(wèn)模式。在一些案例中，o3甚至將延遲提升了超過(guò)9000倍。總體而言，GPT - 5平均能帶來(lái)約20%的加速，其他模型稍遜一籌。不過(guò)，GPT并非在所有問(wèn)題上都能生成最優(yōu)內(nèi)核，研究人員發(fā)現(xiàn)GPT - 5在34%的問(wèn)題上生成了最優(yōu)解，而在另外30%的問(wèn)題上，其他模型生成的解更優(yōu)。這意味著沒(méi)有單一模型能在所有問(wèn)題上都生成最優(yōu)內(nèi)核，因此研究人員開(kāi)展了智能體群體實(shí)驗(yàn)。

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

與單個(gè)模型相比，智能體群體策略實(shí)現(xiàn)了更高的性能提升。與GPT - 5相比，智能體群體在各層級(jí)平均加速31%，在Level 2問(wèn)題上加速42%。在幾乎沒(méi)有上下文信息的情況下，智能體群體表現(xiàn)就已不錯(cuò)。研究人員嘗試為智能體提供更多上下文以獲取更快的內(nèi)核，額外信息來(lái)源包括CUDA實(shí)現(xiàn)和M4上gputrace的性能分析信息。在具體實(shí)施中，先將截圖處理任務(wù)分配給子智能體，讓其為主模型提供性能優(yōu)化提示，主智能體初步實(shí)現(xiàn)后進(jìn)行性能分析和計(jì)時(shí)，再將截圖傳給子智能體生成新的性能優(yōu)化提示。實(shí)驗(yàn)表明，上下文配置沒(méi)有“單一最佳”方案，但加入額外上下文實(shí)現(xiàn)了平均1.87倍的加速，相較于普通智能體的1.31倍平均加速，提升幅度提高了三倍。

有提升，但看跟誰(shuí)比

在PyTorch中，調(diào)用函數(shù)時(shí)，PyTorch會(huì)先將其拆解為張量運(yùn)算，再交給GPU執(zhí)行，這就需要GPU內(nèi)核將數(shù)學(xué)操作轉(zhuǎn)換為GPU可理解的低級(jí)并行指令，所以GPU內(nèi)核性能對(duì)運(yùn)算效率很關(guān)鍵。這項(xiàng)研究讓原本需工程師手寫的內(nèi)核優(yōu)化由AI自動(dòng)完成并測(cè)試性能。然而，蘋果硬件不像英偉達(dá)的CUDA那樣對(duì)PyTorch有很好的優(yōu)化，直接拿MPS后端原生實(shí)現(xiàn)和AI生成的內(nèi)核對(duì)比不太公平。有網(wǎng)友指出，文章使用的baseline是eager mode，通常用于訓(xùn)練調(diào)試或指標(biāo)計(jì)算，不會(huì)真正部署到設(shè)備上。在真實(shí)部署中，一般會(huì)先將模型導(dǎo)出為ONNX，再編譯成設(shè)備原生格式，效率比直接用PyTorch eager mode高很多。所以，無(wú)論內(nèi)核是手寫還是AI生成，優(yōu)化后的GPU內(nèi)核都比未優(yōu)化的PyTorch推理快。研究人員回應(yīng)稱，這項(xiàng)工作不是為了展示部署環(huán)境的最終性能極限，而是展示AI自動(dòng)生成內(nèi)核的可行性，目的是在內(nèi)核工程方面獲得人類專家一定程度的效益，無(wú)需開(kāi)發(fā)人員額外投入，重點(diǎn)在于原型驗(yàn)證。

參考鏈接：

[1]https://gimletlabs.ai/blog/ai-generated-metal-kernels#user-content-fn-4

[2]https://news.ycombinator.com/item?id=45118111

[3]https://en.wikipedia.org/wiki/Compute_kernel

[4]https://github.com/ScalingIntelligence/KernelBench/

本文來(lái)自微信公眾號(hào)“量子位”，作者：henry，36氪經(jīng)授權(quán)發(fā)布。

本文僅代表作者觀點(diǎn)，版權(quán)歸原創(chuàng)者所有，如需轉(zhuǎn)載請(qǐng)?jiān)谖闹凶⒚鱽?lái)源及作者名字。

免責(zé)聲明：本文系轉(zhuǎn)載編輯文章，僅作分享之用。如分享內(nèi)容、圖片侵犯到您的版權(quán)或非授權(quán)發(fā)布，請(qǐng)及時(shí)與我們聯(lián)系進(jìn)行審核處理或刪除，您可以發(fā)送材料至郵箱：service@tojoy.com