就在剛剛，由Transformer八大金剛之一Llion Jones創辦的Sakana AI，推出了世界上首個「AI CUDA工程師」——The AI CUDA Engineer。

同時，還發布了一個由「AI CUDA工程師」生成的，包含17,000多個已驗證CUDA內核的數據集。

論文地址：https://pub.sakana.ai/static/paper.pdf

項目主頁：https://pub.sakana.ai/ai-cuda-engineer

自從去年提出了世界上首個「AI科學家」之后，一個疑問一直縈繞在團隊心頭：如果可以使用AI來進行AI研究，是不是也能用AI來研究如何讓AI運行得更快？

團隊認為，AI系統也應該像人腦一樣高效地利用資源。而實現的最佳途徑，就是利用AI來讓AI更高效！

于是，「AI CUDA工程師」就這樣誕生了。

這位AI工程師，能自動生成高度優化的CUDA內核，不僅比PyTorch中常見的機器學習操作快10-100倍，而且也比生產中廣泛使用的CUDA內核快最高5倍。

通過LLM驅動的進化代碼優化技術，「AI CUDA工程師」能將PyTorch代碼轉換為CUDA內核，并通過進化算法優化 CUDA內核的執行效率，實現多個運算操作的融合。

更重要的是，它能夠通過學習歷史創新成果，發現編寫高效CUDA運算的全新方法。

最終，開啟一個AI驅動加速AI的新紀元，實現自動化推理時優化。

英偉達高級研究科學家Jim Fan表示：「干得漂亮，Sakana AI！這是我最近看到的最酷的自主編程AI智能體。AutoML王者歸來！」

對計算資源最有價值的使用方式，就是提高相同計算資源的未來生產力。它具備了所有超越人類表現的關鍵要素：

1. CUDA調試異常繁瑣，需要大量領域專業知識，有時候就像是煉金術；
2. 獎勵函數定義得非常清晰：數值正確性和執行時間；
3. 進化周期很快，因為正確性檢查很輕量，獎勵評估也很簡單；
4. 可以進行完全并行化的搜索。

而且，「創新檔案」這個概念也很有趣，跟英偉達「我的世界」Voyage智能體中使用的「技能庫」非常相似。

本質上，這是一個完全由LLM在探索CUDA設計空間時編寫的代碼倉庫，還保存了最佳代碼片段。這是進化搜索過程中產生的學習成果，無需更新參數權重。

Jim Fan順便推薦了一下，英偉達最近也在展開的一項相關研究，同樣也是用強大了的LLM推理來改進CUDA內核。

「AI CUDA工程師」：基于智能體的CUDA內核自動發現、優化、組合

CUDA是一個底層軟件層，能夠直接訪問英偉達GPU的硬件指令集來進行并行計算。CUDA內核，則是使用CUDA語言編寫的、在GPU上運行的函數。

通過在CUDA內核層面直接編寫指令，我們能夠顯著提升AI算法的性能。

然而，使用CUDA需要扎實的GPU專業知識，而在實踐中，大多數機器學習算法都是在PyTorch或JAX等更高層次的抽象層中編寫的。

「AI CUDA工程師」是一個基于前沿LLM的AI智能體框架，目的是將標準PyTorch代碼自動轉換為高度優化的CUDA內核。

通過運用進化優化技術，并借鑒進化計算中的「交叉」操作和「創新檔案」等概念來發現具有潛力的過渡性內核，研究人員提出的框架不僅能夠自動將PyTorch模塊轉換為CUDA內核，而且他們的高度優化的CUDA內核通常能夠實現顯著的性能提升。

研究團隊相信這項技術能夠帶來性能加速，從而加快LLM或其他生成式AI模型等基礎大模型的訓練和運行（推理）速度，最終讓AI模型在英偉達GPU上運行得更快。

• 第1和第2階段（轉換和翻譯）：「AI CUDA工程師」首先將PyTorch代碼轉換為可運行的CUDA內核。即便在沒有專門優化的情況下，它們就已經觀察到初步的性能提升。
• 第3階段（進化優化）：借鑒生物進化中的適者生存法則，框架采用進化的方法來確保只保留最優秀的CUDA內核。同時，團隊創新性地提出了內核交叉提示策略，能夠將多個優化后的內核進行有效組合。
• 第4階段（創新檔案）：正如文化通過千年的文明傳承來塑造人類智慧一樣，「AI CUDA工程師」也會利用在過往創新和發現中積累的經驗，構建一個由高性能CUDA內核構成的創新檔案庫，并以這些積累的優化經驗為基礎，實現更進一步的轉換和性能突破。

提升內核運行速度

研究團隊相信，這僅僅是人工智能偉大優化的開始！

在論文中，他們引入了一種端到端的智能體工作流，能夠將PyTorch代碼轉換為可工作的CUDA內核，優化CUDA運行時性能，并自動融合多個內核。

此外，他們構建了各種技術來增強流程的一致性和性能，包括LLM集成、迭代分析反饋循環、本地內核代碼編輯以及交叉內核優化。

AI發現的部分CUDA內核在KernelBench上實現了SOTA的性能

研究人員表示，「AI CUDA工程師」能夠穩健地轉換250個考慮的Torch操作中的230多個，并且對于大多數內核實現了顯著的運行時性能提升。此外，該方法還能夠有效地融合各種內核操作，并且可以超越幾種現有的加速操作。

具體而言，在229個考慮的任務中，他們在81%的任務上優于PyTorch原生運行。此外，所有發現的CUDA內核中有20%至少比其PyTorch實現快兩倍。

下圖展示了PyTorch原生內核與經過優化處理后的內核在不同級別（Level 1、Level 2、Level 3）下的加速對比。

從下圖展示的內核可以看出，「AI CUDA工程師」能實現各種操作，包括歸一化方法、損失函數、特殊矩陣乘法，甚至是完整的神經網絡架構。

17,000+個已驗證CUDA內核數據集

研究團隊隨論文一同發布的「資料庫」（The AI CUDA Engineer Archive）是一個包含超過3萬個CUDA內核的數據集，這些內核均由「AI CUDA工程師」生成。

數據集采用CC-By-4.0許可證發布，可通過HuggingFace平臺訪問，并可用于LLM的進階微調。

其中包括，PyTorch參考實現、PyTorch/NCU/Clang-tidy性能分析數據、每個任務的多個內核版本、錯誤信息，以及與PyTorch原生和編譯運行時間的性能對比數據。

團隊希望這個數據集能夠促進開源模型的后續訓練，從而開發更出色的CUDA功能模塊。包括離線強化學習、偏好優化和標準監督微調等技術。

數據集地址：https://huggingface.co/datasets/SakanaAI/AI-CUDA-Engineer-Archive

統計數據顯示，「資料庫」共包含超過3萬個內核，其中已驗證正確的實現超過1.7萬個，約有50%的內核性能優于PyTorch原生運行時間。

基于文本嵌入進一步分析發現，這些內核可根據不同的任務類型（如矩陣乘法、池化操作、卷積運算）和實現策略（包括循環展開、運算融合、向量化處理）自然地聚類。

探索內核資料庫

除此之外，研究人員還開發了一個交互式網站，供用戶實時查看超過1.7萬個已驗證內核及其配置信息，包括PyTorch、NCU和Clang-Tidy的相關數據。

在網站中，用戶可以探索涵蓋230個任務的各類高性能內核。并且可以動手定制一個排行榜，查看不同實驗和LLM間的相關內核。

內核性能排行榜

并且，還可以直觀地查看內核可視化結果、搜索相關內核、下載驗證代碼和性能提升數據，以及查看詳細的性能分析結果。并且還可以深入了解優化實驗的具體細節。

性能分析數據、評估腳本下載、相關內核信息和實驗發現細節的實例歸一化內核

AI革命才剛剛起步

研究人員認為，今天的LLM就像是我們這一代的「大型機」。

我們仍然處在AI發展的早期階段，在市場競爭和全球創新的推動下（尤其是那些在資源受限條件下進行創新的團隊），這項技術必將實現百萬量級的效率提升。

目前，我們的AI系統消耗著海量資源，如果這項技術繼續擴張而不考慮效率和能源消耗問題，其發展將難以持續。

從根本上來說，沒有任何理由表明我們的AI系統不能達到甚至超越人類智能的效率水平。我們相信，要實現這種效率提升，最佳途徑就是利用AI來優化AI本身。

這正是Sakana AI努力的方向，而這個項目正是朝著實現AI百萬倍性能提升目標邁出的關鍵一步。

就像從早期龐大臃腫的大型機演變到現代計算機一樣，相比當今「臃腫」低效的LLM，未來幾年我們使用AI的方式將會發生翻天覆地的變化。