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

好家伙，AI意外生成的內核（kernel），性能比人類專家專門優化過的還要好！

斯坦福最近披露了一組新發現，結果真的太亮眼了。

由AI優化的內核，在常見深度學習操作上，翻倍超越原生PyTorch，性能至多可以提升近400%——

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

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

更驚人的是，這一切都是意外實現的。

研究團隊本來的目標是生成合成數據以訓練內核生成模型。

結果發現，僅在測試階段生成的合成數據本身，竟然可以生成性能非常優秀的內核。

圍觀網友：沒想到AI也要取代內核工程師了。

還有人發現，除了性能大幅提升外，研究團隊采用的方法也非常有趣：

他們沒有簡單的在操作上逐步優化（類似于爬坡算法），而是在每次迭代之間加入了一個語言推理的步驟，通過這種方式鼓勵搜索過程更加多樣化。
也就是說，他們是讓系統在每次改進時通過類似“思考”的方式產生更多想法，從而找到更好的解決方案。

具體如何實現，一起來看。

改代碼前先生成自然語言優化思想

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

這些內核是用純CUDA-C編寫，無需使用CUTLASS和Triton等庫和DSL（Domain-Specific Language，領域專用語言）。

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

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

為了進一步增強思路多樣性，斯坦福團隊還使用了多分支的探索模式。

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

團隊使用OpenAI o3和Gemini 2.5 Pro挑戰KernelBench 1級中的10個問題，運行多輪后，最佳內核開始出現。

其中大多數最佳結果出現在后續輪次（總共5輪），并且主要是第4輪或第5輪。

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

這一發現再加上之前DeepMind的AlphaEvolve，以及o3發現Linux的0day漏洞等一系列事件，讓網友們認為Gemini 2.5Pro和o3的能力水平已經達到了新的層級。

回到斯坦福的項目，在生成過程當中，可以看到模型的生成思路開始顯現出與人類的經驗相似之處——

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

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

在具體實現上，有人詢問AI生成CUDA內核時的優化建議，是否可以被轉化為對應代碼實現、還是說只是觸發了隨機探索？

作者回應說，盡管還沒有進行更嚴謹的系統驗證，但是手動檢查的案例中，生成的CUDA實現與提出的優化建議是大致匹配的。

即AI并不是在完全隨機做優化，而是確實在嘗試實現它自己提出的策略。

華人主創團隊意外發現

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

Ouyang目前是斯坦福大學擴展智能實驗室的博士生，她本碩畢業于麻省理工，曾在英偉達cuDNN團隊工作。

Percy Liang是斯坦福大學計算機科學副教授兼統計學助理教授，目前擔任斯坦福基礎模型研究中心主任。

曾和李飛飛一起發布、推進了多項研究工作。

Azalia Mirhoseini是斯坦福大學計算機科學助理教授、斯坦福擴展實驗室創始人。她曾在DeepMind、Google Brain以及Anthropic工作過。

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

本次研究，本來是希望生成數據來訓練內核生成模型。

但是在過程中卻出現了意想不到的結果，僅在測試階段生成的合成數據本身，竟然可以生成性能非常優秀的內核。

因為這些內核利用了此前被認為很難實現的高級優化和硬件特性，所以團隊決定以博客形式分享此次成果。

不過具體是如何生成數據的，研究團隊暫時不對外發布，只是提到了這種設計理念也很簡單

最關鍵的還是，它已經展示出了巨大潛力。

此外，研究團隊也認為此次發現也與最近的一些趨勢相呼應——大規模再訓練已不是必需。

有時，聰明的搜索和分支策略，可以解鎖科學創新并解決復雜問題，通過verifier進行廣泛搜索還能有更多收獲。
將強大推理能力與同時探索多個假設結合起來，能帶來更好結果。就像AlphaEvolve、AlphaEvolution、 Gemini 2.5 Pro深度思考一樣。

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

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

與FP16或BF16相比，FP32在新推出硬件上的優化程度通常比較低，這也是為何使用FP32內核比PyTorch更容易實現性能提升。

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

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

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

One More Thing

實際上，不只是一個團隊在嘗試開發內核大模型。

就在5月，開發了Devin的Cognition開源了首個通過強化學習即可編寫CUDA內核的大模型Kevin-32B。

它基于QwQ-32B在KernelBench數據集上使用GRPO，實現了多輪強化學習，性能優于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