新智元報導
編輯：Aeneas 好困
【新智元導讀】原本只想小試身手生成一些數據，沒想到卻不經意地擊敗了PyTorch的專業核心程式！史丹佛華人團隊用純粹的CUDA-C寫出的AI生成核心程式，瞬間驚艷了業界並登上Hacker News熱門排行。團隊甚至表示：原本不想公開這個成果的。
就在剛才，史丹佛HAI的華人頂尖團隊又帶來了驚人新作。
他們用純粹的CUDA-C語言編寫的快速AI生成核心程式，竟然超越了PyTorch！
在這個過程中，完全不用借助CUTLASS和Triton等函式庫和領域特定語言（DSL），就能讓效能表現接近PyTorch內建、經過專家最佳化的標準生產級核心程式，甚至在某些情況下還更勝一籌。
作者團隊都是我們熟悉的名字——Anne Ouyang、Azalia Mirhoseini和Percy Liang，有趣的是，他們甚至直言，這個結果其實原本不想拿出來發表。
一經發表，這個發現就引爆了技術圈，現在已經登上Hacker News總榜第二名。
說起來，這個發現還有很多意外的成分。
原本，他們的目標是生成合成數據，來訓練更好的核心程式生成模型，合成數據生成的設計也十分簡單。
然而，意想不到的事情發生了，僅用於測試的合成數據生成本身，竟然開始生成非常優秀的核心程式，甚至超越了人類專家最佳化的PyTorch基準，而且還利用了高階最佳化和硬體特性。
而在此之前，這是一項非常艱鉅的挑戰。
因此，研究者們決定提前撰寫部落格文章，分享他們的發現。
總結來說，研究的亮點成果如下：
矩陣乘法（Matmul, FP32）：效能達到PyTorch FP32 torch.matmul的101.3%
二維卷積（Conv2D, FP32）：效能達到PyTorch FP32 torch.nn.Conv2D的179.9%
Softmax（FP32）：效能達到PyTorch FP32 torch.softmax的111.8%
層正規化（LayerNorm, FP32）：效能達到PyTorch FP32 torch.nn.LayerNorm的484.4%
二維卷積 + ReLU + 最大池化（Conv2D + ReLU + MaxPool, FP32）：效能達到PyTorch FP32參考實作的290.1%，達到PyTorch FP32 torch.compile()參考實作的189.0%
以上結果在英偉達L40S GPU上進行了基準測試，效能百分比定義為參考時間除以生成的核心程式時間。
網友：強制LLM推論，實在太有趣了
在Hacker News上，網友們也對此展開了熱烈討論。
例如為什麼使用FP32核心程式會比PyTorch更容易實現效能提升，理由就相當有趣。
如果AI真的能以更低的成本，實現更最佳化的核心程式，的確潛力巨大。
最令人震撼的就是，無論是最近谷歌的AlphaEvolve，還是o3在Linux核心程式中發現了零日漏洞，都在提醒我們——
Gemini Pro 2.5和o3已經達到了一個全新的能力水準，那些曾經在其他模型上嘗試失敗的想法，現在突然奏效了。
可以說，我們已經到達了一個節點，大型語言模型（LLM）能比人類快得多的速度進行迭代和測試，資訊組合、進步和智慧應用的蠻力，似乎正在成功！
接下來，我們來看看史丹佛研究者們部落格中的具體內容。
部落格全文
方法
研究者們採用了KernelBench的任務設定（這是他們在2024年12月發布的一款基於AI的核心程式生成基準測試）。
具體來說，給定一段torch程式碼，大型語言模型（LLM）會編寫自定義核心程式來取代原有的torch運算子，目標是實現加速。
依照KernelBench最初的設計，參考程式碼預設使用FP32精度；在給定的容許誤差閾值（1e-02）下，採用較低精度的解決方案也是被允許的。
此外，由於存在大量針對特定規模的最佳化手段，KernelBench中的每個問題都設定了具體的輸入大小。
因此，該基準測試旨在找出針對特定問題規模的最快核心程式，而非一個適用於任意問題規模的高速核心程式。
而且，研究者會同時運行torch參考程式碼和生成的程式碼，並透過在多種隨機輸入下比較兩者輸出的數值是否一致，來檢驗其正確性。
目前，在最佳化核心程式這個問題上，業界擴展測試時計算資源最常用的方法是「序列修訂（sequential revision）」。
這是一種多輪迭代的循環：模型首先對核心程式進行增量式修改，接著檢查其正確性和效能，然後根據結果再次嘗試。
也就是說，要嘛修復有問題的核心程式，要嘛進一步提升現有核心程式的效能。
這個循環過程非常直觀，也容易實現。模型會修復失效的核心程式，微調可用的核心程式，一步步最佳化出效能更佳的版本。
這種方法的主要局限，在於最佳化思路缺乏多樣性。
序列循環往往容易陷入局部最佳的困境，例如反覆嘗試同類型的轉換，或是在缺乏潛力的最佳化路徑上無休止地調整。
其結果便是測試時計算資源的低效利用，並且難以促使模型產生具有根本性創新的最佳化思路。
為解決這一問題，研究者引入了兩項關鍵改變：
運用自然語言對最佳化思路進行推論
他們不再於每一步直接生成新的核心程式，而是以先前嘗試過的思路為條件，用自然語言生成最佳化思路，隨後將這些思路具體化為新的程式碼變體。
在每個最佳化步驟進行分支擴展
他們不是每步只改進一個候選方案，而是進行分支擴展，讓每個思路都能派生出多種實現版本，其中效能最佳的核心程式將作為下一輪最佳化的種子。
（研究者也會保留一個表現優異的現有核心程式庫，用於提供種子）。
這種方式解鎖了大規模的平行處理能力，使他們能夠在每一輪探索截然不同的最佳化方向，避免陷入狹窄的最佳化路徑。
其結果是，這種測試時循環不再像序列修訂那般，僅僅是與編譯器「對話」，而是更接近一種結構化的探索性搜尋。
這種搜尋由明確的最佳化假設指導，並採用大規模平行評估的方式進行。
研究者運行了KernelBench第1級的10個問題，以進行測試。
他們調整了問題規模，以確保核心程式啟動開銷相對於問題的整體運行時間而言可以忽略不計。
然後，使用OpenAI o3和Gemini 2.5 Pro模型進行了5輪實驗。
下圖展示了首次發現效能最佳核心程式所在的輪次分布情況。
可以看到，大多數最佳結果出現在靠後的輪次（總共5輪），其中絕大部分出現在第4輪或第5輪。
隨著擴大搜尋範圍，研究者還發現：許多高效能核心程式的最佳化策略高度相似，集中在少數幾種常見的模式上，這與他們手動編寫核心程式的經驗也是一致的。
主要的最佳化類別歸納如下——
記憶體存取最佳化：提升不同記憶體層級（全域記憶體、共享記憶體、暫存器）之間數據遷移的效率，並確保數據存取方式能夠最大化頻寬、最小化衝突。
非同步操作與延遲隱藏：透過將耗時較長的操作（例如全域記憶體存取）與計算或其他記憶體傳輸重疊執行，來隱藏其帶來的延遲。
資料類型與精度最佳化：在允許的條件下，盡可能使用較低精度的資料類型（如FP16或BF16），以降低記憶體頻寬需求，提升快取效率，並有望利用專門的硬體加速單元。
計算與指令最佳化：提升算術運算本身的效率，削減指令數量，或利用專門的硬體指令。
平行性與佔用率增強：最大化串流多處理器（SM）上活躍執行緒束（warp）的數量，以便更好地隱藏延遲，提高整體吞吐率。
控制流與循環最佳化：減少由循環、分支及索引計算等引入的額外開銷。
總結
這次研究者採用的方法，與AI研究中一個日益顯著的趨勢不謀而合——
將強大的推論能力與對多個假設的平行探索相結合，能夠帶來效能的提升。
正如一些近期研究（例如AlphaEvolve、Gemini 2.5 Pro Deep Think）所強調的，我們並不總是需要大規模的重新訓練。
有時，巧妙的搜尋和分支策略便足以催生科學創新、攻克複雜難題，而借助驗證器進行廣泛搜尋，則可能帶來更大的收益。
然而，這並不意味著我們不需要進一步的訓練。
恰恰相反，研究者的這種方法，也有助於生成更優質的合成數據，用以改進未來的模型訓練（這需要更多的問題實例）。
因此，它既是一種強大的測試時擴展方法，也是我們邁向更智慧、數據效率更高的模型開發之路的一步。
而且，這次研究者展現的僅僅是初步的成果。這些最佳化結果的品質看起來相當可觀，但仍有廣闊的提升空間，例如產生更優的最佳化思路、生成更高品質的最終程式碼，以及將此方法應用於日益複雜的核心程式。
目前，研究者仍在積極改進的兩個具體例子包括：
FP16 Matmul：效能達到torch.matmul的52%
FP16 Flash Attention：效能達到torch.nn.functional.scaled_dot_product_attention的9%
在現代機器學習任務中，FP32的應用不如FP16或BF16普遍，並且在較新的硬體上，針對FP32的最佳化往往也更少。
這或許能部分解釋，為何基於FP32的核心程式更容易在效能上超越PyTorch。
作者介紹
Anne Ouyang
Anne Ouyang目前是史丹佛大學電腦科學（CS）博士生，在Scaling Intelligence Lab（可擴展智慧實驗室）進行研究。
她的研究興趣主要集中在可擴展的自我改進機器學習系統，同時也廣泛關注實證機器學習（empirical ML）和效能工程（performance engineering）。
此前，她在MIT獲得學士和碩士學位，並曾在NVIDIA cuDNN團隊工作，負責編寫CUDA核心程式，用於加速GPU上的深度學習工作負載。
Azalia Mirhoseini
Azalia Mirhoseini是史丹佛大學電腦科學助理教授，也是Scaling Intelligence Lab（可擴展智慧實驗室）的創始人，並在Google DeepMind兼任高級研究科學家。
她的實驗室致力於開發可擴展的自主演進人工智慧系統與方法論，以期推動通用人工智慧的發展。
在加入史丹佛大學之前，她曾在Google Brain和Anthropic等業界頂尖的人工智慧實驗室工作多年。
她過往的卓越成就包括：
提出混合專家（MoE）神經架構——目前已被前瞻的AI模型廣泛應用；
領導AlphaChip專案——一項將深度強化學習用於佈局最佳化的開創性工作，並成功應用於Google AI加速器（TPU）及資料中心CPU等先進晶片的設計中；
在測試時計算的Scaling方面有深入的研究
Percy Liang
Percy Liang是史丹佛大學電腦科學副教授，兼任基礎模型研究中心（CRFM）主任。同時也是CodaLab Worksheets的創建者，並藉此堅定倡導科研工作的可重現性。
他目前專注於透過開源和嚴格的基準測試，提升基礎模型（特別是大型語言模型）的可及性與可理解性。
他曾圍繞機器學習和自然語言處理領域進行了廣泛研究，具體方向包括穩健性、可解釋性、人機互動、學習理論、知識落地、語義學以及推論等。
此前，他於2004年在MIT獲得學士學位，並於2011年在UC柏克萊獲得博士學位。
核心程式最佳化過程
最後，展示一個Conv2D自動生成思路的最佳化軌跡範例，torch參考基準時間為1.41毫秒。
第0輪：7.02毫秒，達到參考效能的20.1%
思路：給定pytorch程式碼，用CUDA核心程式取代操作。
第1輪：7.54毫秒，達到參考效能的18.8%
思路：透過使用__ldg載入不變的張量來利用唯讀快取。
第2輪：3.46毫秒，達到參考效能的41.0%
思路：將卷積轉換為FP16 Tensor-Core GEMM。
作者評論：這是一種演算法最佳化，將卷積轉換為隱式GEMM，這對於在Tensor Core上高效運行卷積至關重要
第3輪：3.67毫秒，達到參考效能的38.7%
思路：採用雙緩衝cp.async管線，使全域記憶體載入與Tensor-Core計算重疊。
第4輪：3.46毫秒，達到參考效能的41.0%
思路：給定pytorch程式碼，使用隱式矩陣乘法（implicit matmul）的CUDA核心程式取代操作。給定的GEMM核心程式可能會有所幫助。
作者評論：因為最佳化涉及到使用GEMM，所以在這一輪開始時，使用了一個之前生成的現有優秀GEMM核心程式作為種子，這個思路是手動編寫的
第5輪：1.91毫秒，達到參考效能的74.9%
思路：在每個K-tile循環內部，預先計算k_idx分解的核心程式/輸入索引並在共享記憶體中重用，以避免冗餘算術運算。
第6輪：1.37毫秒，達到參考效能的103.6%
思路：在共享記憶體中預先計算並快取N維GEMM索引，以減少輸入數據載入循環內的冗餘算術運算。
第7輪：1.38毫秒，達到參考效能的102.9%
思路：透過使用專用的每執行緒束共享記憶體緩衝區來平行化CUDA核心程式輸出，以消除寫入全域記憶體時的執行緒束級序列化。
第8輪：1.37毫秒，達到參考效能的103.6%
思路：在共享記憶體中預先計算並快取基礎輸入座標，以減少輸入數據載入過程中的冗餘算術運算。
第9輪：1.36毫秒，達到參考效能的105.1%
思路：對B-fragment載入進行軟體管線處理，使下一個B-tile的共享記憶體讀取與當前B-tile的WMMA計算重疊。
第10輪：1.07毫秒，達到參考效能的133.6%
思路：重用共享記憶體中預先計算的N維GEMM分解來進行輸出地址計算，避免冗餘且高成本的除法/模運算。
第11輪：1.21毫秒，達到參考效能的117.4%
思路：在half WMMA運算中移除hi/lo分解，如果最終精度可接受，則依賴標準的FP16累積來提高效能。
第12輪：1.01毫秒，達到參考效能的141.2%
思路：透過在管線的載入階段動態計算K維索引，使用雙緩衝使Asub（權重）和Bsub（輸入）的K循環全域記憶體載入與MMA計算重疊。
第13輪：0.795毫秒，達到參考效能的179.9%
思路：透過使用像half2這樣的更寬資料類型，為載入Asub_pipe和Bsub_pipe實現向量化的共享記憶體寫入。
最終程式碼
最終生成的Conv2D核心程式碼，使用了先進的CUDA技術，就是人類自己寫起來都很有挑戰性的那種！
參考資料：
https://crfm.stanford.edu/2025/05/28/fast-kernels.html
https://news.ycombinator.com/item?id=44139454