斯坦福意外用AI生成超強(qiáng)CUDA內(nèi)核，性能比人類專家優(yōu)化得還要好！翻倍碾壓原生PyTorch，華人主創(chuàng)

克雷西 2025-05-31 11:49:49 來源：量子位

明敏 克雷西 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI

好家伙，AI意外生成的內(nèi)核（kernel），性能比人類專家專門優(yōu)化過的還要好！

斯坦福最近披露了一組新發(fā)現(xiàn)，結(jié)果真的太亮眼了。

由AI優(yōu)化的內(nèi)核，在常見深度學(xué)習(xí)操作上，翻倍超越原生PyTorch，性能至多可以提升近400%——

• • 矩陣乘法（Matmul，F(xiàn)P32）：性能達(dá)到PyTorch?torch.matmul的101.3%。
• • 二維卷積（Conv2D）：性能達(dá)到?torch.nn.Conv2D的179.9%。
• • Softmax：性能達(dá)到?torch.softmax的111.8%。
• • 層歸一化（LayerNorm）：性能達(dá)到torch.nn.LayerNorm的484.4%。
• • Conv2D+ReLU+MaxPool組合操作：性能達(dá)到PyTorch參考實現(xiàn)的290.1%，以及torch.compile()參考實現(xiàn)的189.0%。

（在NVIDIA L40S GPU上進(jìn)行基準(zhǔn)測試，性能百分比定義為參考時間除以生成的kernel_size時間）

更驚人的是，這一切都是意外實現(xiàn)的。

研究團(tuán)隊本來的目標(biāo)是生成合成數(shù)據(jù)以訓(xùn)練內(nèi)核生成模型。

結(jié)果發(fā)現(xiàn)，僅在測試階段生成的合成數(shù)據(jù)本身，竟然可以生成性能非常優(yōu)秀的內(nèi)核。

圍觀網(wǎng)友：沒想到AI也要取代內(nèi)核工程師了。

還有人發(fā)現(xiàn)，除了性能大幅提升外，研究團(tuán)隊采用的方法也非常有趣：

他們沒有簡單的在操作上逐步優(yōu)化（類似于爬坡算法），而是在每次迭代之間加入了一個語言推理的步驟，通過這種方式鼓勵搜索過程更加多樣化。

也就是說，他們是讓系統(tǒng)在每次改進(jìn)時通過類似“思考”的方式產(chǎn)生更多想法，從而找到更好的解決方案。

具體如何實現(xiàn)，一起來看。

改代碼前先生成自然語言優(yōu)化思想

按照斯坦福團(tuán)隊博客的描述，這種內(nèi)核生成的思路非常簡單——給定torch代碼，然后告訴大模型編寫自定義內(nèi)核來替換torch算子。

這些內(nèi)核是用純CUDA-C編寫，無需使用CUTLASS和Triton等庫和DSL（Domain-Specific Language，領(lǐng)域?qū)Ｓ谜Z言）。

不同于傳統(tǒng)方法的是，模型并不是一上來就直接改代碼，而是先用自然語言生成優(yōu)化思想，然后再將這些思想轉(zhuǎn)化為新的代碼變體。

團(tuán)隊這樣做的理由是，“按順序修改”式的優(yōu)化思路缺乏多樣性，導(dǎo)致陷入局部極小值，重復(fù)訪問同一類轉(zhuǎn)換或無休止地優(yōu)化沒有前景的軌跡。

為了進(jìn)一步增強(qiáng)思路多樣性，斯坦福團(tuán)隊還使用了多分支的探索模式。

具體來說，他們的方法并非每一步都只優(yōu)化一個候選方案，而是將每個想法分散開來，使其衍生出多個實現(xiàn)，并使用性能最高的內(nèi)核作為下一輪的種子。

團(tuán)隊使用OpenAI o3和Gemini 2.5 Pro挑戰(zhàn)KernelBench 1級中的10個問題，運行多輪后，最佳內(nèi)核開始出現(xiàn)。

其中大多數(shù)最佳結(jié)果出現(xiàn)在后續(xù)輪次（總共5輪），并且主要是第4輪或第5輪。

KernelBench是斯坦福團(tuán)隊自己提出的一套AI生成內(nèi)核測試基準(zhǔn)，基準(zhǔn)中的任務(wù)分為3個級別，其中1級是指單一原始操作（Single primitive operation），包括AI的基礎(chǔ)構(gòu)建塊（例如卷積、矩陣-向量與矩陣-矩陣乘法、損失函數(shù)、激活函數(shù)以及層歸一化）。

這一發(fā)現(xiàn)再加上之前DeepMind的AlphaEvolve，以及o3發(fā)現(xiàn)Linux的0day漏洞等一系列事件，讓網(wǎng)友們認(rèn)為Gemini 2.5Pro和o3的能力水平已經(jīng)達(dá)到了新的層級。

回到斯坦福的項目，在生成過程當(dāng)中，可以看到模型的生成思路開始顯現(xiàn)出與人類的經(jīng)驗相似之處——

• 內(nèi)存訪問優(yōu)化：?提高不同內(nèi)存層次結(jié)構(gòu)（全局內(nèi)存、共享內(nèi)存、寄存器）之間數(shù)據(jù)移動的效率，并確保以最大化帶寬和最小化沖突的方式訪問數(shù)據(jù)；
• 異步操作和延遲隱藏：?通過將慢速操作（如全局內(nèi)存訪問）與計算或其他內(nèi)存?zhèn)鬏斨丿B，“隱藏”慢速操作的延遲；
• 數(shù)據(jù)類型和精度優(yōu)化：?盡可能使用低精度數(shù)據(jù)類型（如 FP16 或 BF16）以減少內(nèi)存帶寬要求、提高緩存效率；
• 計算和指令優(yōu)化：提高算術(shù)計算本身的效率，減少指令數(shù)量，或利用專門的硬件指令；
• 并行性和占用率增強(qiáng)：最大化流多處理器（SM）上的活動線程數(shù)量，以更好地隱藏延遲并提高整體吞吐量；
• 控制流和循環(huán)優(yōu)化：減少與循環(huán)、分支和索引計算相關(guān)的開銷。

并且斯坦福團(tuán)隊還展示了一組具體的優(yōu)化軌跡，從中可以看出，并不是每一步優(yōu)化都一定能讓速度更快，但經(jīng)過多個步驟的組合，內(nèi)核的速度能夠得到大幅提升，并最終超越PyTorch。

在具體實現(xiàn)上，有人詢問AI生成CUDA內(nèi)核時的優(yōu)化建議，是否可以被轉(zhuǎn)化為對應(yīng)代碼實現(xiàn)、還是說只是觸發(fā)了隨機(jī)探索？

作者回應(yīng)說，盡管還沒有進(jìn)行更嚴(yán)謹(jǐn)?shù)南到y(tǒng)驗證，但是手動檢查的案例中，生成的CUDA實現(xiàn)與提出的優(yōu)化建議是大致匹配的。

即AI并不是在完全隨機(jī)做優(yōu)化，而是確實在嘗試實現(xiàn)它自己提出的策略。

華人主創(chuàng)團(tuán)隊意外發(fā)現(xiàn)

這項研究共有三位作者：Anne Ouyang、Azalia Mirhoseini和Percy Liang。

Ouyang目前是斯坦福大學(xué)擴(kuò)展智能實驗室的博士生，她本碩畢業(yè)于麻省理工，曾在英偉達(dá)cuDNN團(tuán)隊工作。

Percy Liang是斯坦福大學(xué)計算機(jī)科學(xué)副教授兼統(tǒng)計學(xué)助理教授，目前擔(dān)任斯坦?；A(chǔ)模型研究中心主任。

曾和李飛飛一起發(fā)布、推進(jìn)了多項研究工作。

Azalia Mirhoseini是斯坦福大學(xué)計算機(jī)科學(xué)助理教授、斯坦福擴(kuò)展實驗室創(chuàng)始人。她曾在DeepMind、Google Brain以及Anthropic工作過。

她此前參與的研究包括MoE、芯片設(shè)計算法AlphaChip等。

本次研究，本來是希望生成數(shù)據(jù)來訓(xùn)練內(nèi)核生成模型。

但是在過程中卻出現(xiàn)了意想不到的結(jié)果，僅在測試階段生成的合成數(shù)據(jù)本身，竟然可以生成性能非常優(yōu)秀的內(nèi)核。

因為這些內(nèi)核利用了此前被認(rèn)為很難實現(xiàn)的高級優(yōu)化和硬件特性，所以團(tuán)隊決定以博客形式分享此次成果。

不過具體是如何生成數(shù)據(jù)的，研究團(tuán)隊暫時不對外發(fā)布，只是提到了這種設(shè)計理念也很簡單。

最關(guān)鍵的還是，它已經(jīng)展示出了巨大潛力。

此外，研究團(tuán)隊也認(rèn)為此次發(fā)現(xiàn)也與最近的一些趨勢相呼應(yīng)——大規(guī)模再訓(xùn)練已不是必需。

有時，聰明的搜索和分支策略，可以解鎖科學(xué)創(chuàng)新并解決復(fù)雜問題，通過verifier進(jìn)行廣泛搜索還能有更多收獲。

將強(qiáng)大推理能力與同時探索多個假設(shè)結(jié)合起來，能帶來更好結(jié)果。就像AlphaEvolve、AlphaEvolution、 Gemini 2.5 Pro深度思考一樣。

最后，團(tuán)隊表示這項研究還有很多可優(yōu)化的空間。比如他們手頭上就還在優(yōu)化兩個維度：

• FP16 Matmul：52% performance of torch.matmul
• FP16 Flash Attention:：9% performance of torch.nn.functional.scaled_dot_product_attention

與FP16或BF16相比，F(xiàn)P32在新推出硬件上的優(yōu)化程度通常比較低，這也是為何使用FP32內(nèi)核比PyTorch更容易實現(xiàn)性能提升。

他們表示，雖然現(xiàn)在還有不少限制，但是對于未來前景還是很樂觀的。

畢竟最開始，他們連能正常運行的內(nèi)核都生成不了，但是通過不斷優(yōu)化搜索方法，已經(jīng)能讓flash attention的性能提升到了一個不錯的水平。

值得一提的是，搜索使用的資源也很少，大概只用了300萬token輸入和400萬token輸出。

One More Thing

實際上，不只是一個團(tuán)隊在嘗試開發(fā)內(nèi)核大模型。

就在5月，開發(fā)了Devin的Cognition開源了首個通過強(qiáng)化學(xué)習(xí)即可編寫CUDA內(nèi)核的大模型Kevin-32B。

它基于QwQ-32B在KernelBench數(shù)據(jù)集上使用GRPO，實現(xiàn)了多輪強(qiáng)化學(xué)習(xí)，性能優(yōu)于o3、o4-mini。

參考鏈接：
[1]https://crfm.stanford.edu/2025/05/28/fast-kernels.html
[2]https://x.com/anneouyang/status/1928124885567467768
[3]https://x.com/cognition_labs/status/1919835720493236295