Sakana AI 推出自動化 CUDA 最佳化智慧體:PyTorch 程式碼自動轉寫加速
Sakana AI 推出一套自動化 CUDA 核心發現與最佳化框架及基準 robust-kbench。該框架能將 PyTorch 程式碼轉寫為 CUDA 核心,並通過演化生成機制和大模型驗證器迭代最佳化底層效能。這意味著寫底層運算元最佳化的門檻大幅降低,繁重的工程調優可以交給 AI 自動完成。

寫 Python 或 PyTorch 已經成了做 AI 的標配,但要真正榨乾 GPU 的每一滴算力,工程師們還是不得不去啃最底層的 CUDA。
這活兒門檻極高,不僅要精通硬體架構,還要反覆微調記憶體訪問和執行緒排程。現在,Sakana AI 決定把這個重擔丟給大模型,釋出了一套能夠自動化發現、驗證並最佳化 CUDA 核心的智慧體框架。
讓大模型啃下最硬的底層程式碼
最近不少團隊都在嘗試利用測試期計算擴充套件讓大模型寫程式碼,但大部分成果都集中在高層應用,極少有人敢深入底層的 CUDA 核心最佳化。
此外,現有的核心評估基準也存在漏洞,場景多樣性不足,很多模型生成的程式碼看似得分高,實際泛化能力很差。
為了解決這些痛點,Sakana AI 首先推出了名為 robust-kbench 的全新基準測試,用來在多樣化場景下嚴格評估核心的效能與正確性。
在此基礎上,他們搭建了一套自動化的智慧體流水線。前沿大模型會先接收 PyTorch 程式碼,將其直接翻譯成功能等價的 CUDA 核心程式碼。
演化生成加驗證器,速度超過原生實現
翻譯只是第一步,真正的核心在於效能調優。
這套系統引入了一種專為 CUDA 生態定製的演化元生成流程,通過多輪迭代不斷最佳化核心的執行時效能。
為了防止生成出一堆看似跑得快但結果全錯的程式碼,流程中嵌入了基於大模型的驗證器。驗證器負責提前篩查和過濾錯誤核心,大大減輕了直接在硬體上測試的驗證壓力。
在 robust-kbench 上的測試顯示,這套系統生成的 CUDA 核心在實際應用中跑贏了 PyTorch 官方實現,並且這一優勢同時涵蓋了前向傳播與反向傳播。
不僅如此,該框架還學會了運算元融合等高階執行時最佳化策略,能將多個計算步驟合併以減少訪存開銷。
為什麼重要
對從事高效能運算和模型最佳化的開發者來說,手寫底層 CUDA 運算元通常是最耗時、最繁瑣的環節。如果後續能將 PyTorch 原型程式碼穩定自動轉換為高效率的底層實現,開發與部署週期將被大幅縮短。
在行業層面,頂尖的系統級效能最佳化工程師極度稀缺,這也一直是阻礙新架構快速落地的大痛點。Sakana AI 證明了大模型不僅能寫 CRUD,還能在硬核的系統級程式設計中通過演化搜尋完成超乎尋常的微觀最佳化。
我們認為這套方案最聰明的地方在於閉環設計。它沒有依賴大模型單次生成的運氣,而是用演化搜尋配合專屬驗證器,在進入物理硬體前就把不合規的程式碼過濾掉,極大降低了試錯成本。不過看起來目前的測試主要集中在自身構建的基準中,未來面對跨硬體架構和極度複雜的冷門運算元時,這套自動化流水線能否持續穩定輸出,依然需要實踐檢驗。
來源 (1)
- sakana.ai(在新分頁開啟)原文sakana.ai