世界首個「AI CUDA工程師」誕生！AI自己寫代碼優(yōu)化CUDA內(nèi)核，性能狂飆100倍
就在剛剛，由Transformer八大金剛之一Llion Jones創(chuàng)辦的Sakana AI，推出了世界上首個「AI CUDA工程師」——The AI CUDA Engineer。
同時，還發(fā)布了一個由「AI CUDA工程師」生成的，包含17,000多個已驗(yàn)證CUDA內(nèi)核的數(shù)據(jù)集。
論文地址：https://pub.sakana.ai/static/paper.pdf
項目主頁：https://pub.sakana.ai/ai-cuda-engineer
自從去年提出了世界上首個「AI科學(xué)家」之后，一個疑問一直縈繞在團(tuán)隊心頭：如果可以使用AI來進(jìn)行AI研究，是不是也能用AI來研究如何讓AI運(yùn)行得更快？
團(tuán)隊認(rèn)為，AI系統(tǒng)也應(yīng)該像人腦一樣高效地利用資源。而實(shí)現(xiàn)的最佳途徑，就是利用AI來讓AI更高效！
于是，「AI CUDA工程師」就這樣誕生了。
這位AI工程師，能自動生成高度優(yōu)化的CUDA內(nèi)核，不僅比PyTorch中常見的機(jī)器學(xué)習(xí)操作快10-100倍，而且也比生產(chǎn)中廣泛使用的CUDA內(nèi)核快最高5倍。
通過LLM驅(qū)動的進(jìn)化代碼優(yōu)化技術(shù)，「AI CUDA工程師」能將PyTorch代碼轉(zhuǎn)換為CUDA內(nèi)核，并通過進(jìn)化算法優(yōu)化 CUDA內(nèi)核的執(zhí)行效率，實(shí)現(xiàn)多個運(yùn)算操作的融合。
更重要的是，它能夠通過學(xué)習(xí)歷史創(chuàng)新成果，發(fā)現(xiàn)編寫高效CUDA運(yùn)算的全新方法。
最終，開啟一個AI驅(qū)動加速AI的新紀(jì)元，實(shí)現(xiàn)自動化推理時優(yōu)化。
英偉達(dá)高級研究科學(xué)家Jim Fan表示：「干得漂亮，Sakana AI！這是我最近看到的最酷的自主編程AI智能體。AutoML王者歸來！」
對計算資源最有價值的使用方式，就是提高相同計算資源的未來生產(chǎn)力。它具備了所有超越人類表現(xiàn)的關(guān)鍵要素：
- CUDA調(diào)試異常繁瑣，需要大量領(lǐng)域?qū)I(yè)知識，有時候就像是煉金術(shù)；
- 獎勵函數(shù)定義得非常清晰：數(shù)值正確性和執(zhí)行時間；
- 進(jìn)化周期很快，因?yàn)檎_性檢查很輕量，獎勵評估也很簡單；
- 可以進(jìn)行完全并行化的搜索。
而且，「創(chuàng)新檔案」這個概念也很有趣，跟英偉達(dá)「我的世界」Voyage智能體中使用的「技能庫」非常相似。
本質(zhì)上，這是一個完全由LLM在探索CUDA設(shè)計空間時編寫的代碼倉庫，還保存了最佳代碼片段。這是進(jìn)化搜索過程中產(chǎn)生的學(xué)習(xí)成果，無需更新參數(shù)權(quán)重。
Jim Fan順便推薦了一下，英偉達(dá)最近也在展開的一項相關(guān)研究，同樣也是用強(qiáng)大了的LLM推理來改進(jìn)CUDA內(nèi)核。
「AI CUDA工程師」：基于智能體的CUDA內(nèi)核自動發(fā)現(xiàn)、優(yōu)化、組合
CUDA是一個底層軟件層，能夠直接訪問英偉達(dá)GPU的硬件指令集來進(jìn)行并行計算。CUDA內(nèi)核，則是使用CUDA語言編寫的、在GPU上運(yùn)行的函數(shù)。
通過在CUDA內(nèi)核層面直接編寫指令，我們能夠顯著提升AI算法的性能。
然而，使用CUDA需要扎實(shí)的GPU專業(yè)知識，而在實(shí)踐中，大多數(shù)機(jī)器學(xué)習(xí)算法都是在PyTorch或JAX等更高層次的抽象層中編寫的。
「AI CUDA工程師」是一個基于前沿LLM的AI智能體框架，目的是將標(biāo)準(zhǔn)PyTorch代碼自動轉(zhuǎn)換為高度優(yōu)化的CUDA內(nèi)核。
通過運(yùn)用進(jìn)化優(yōu)化技術(shù)，并借鑒進(jìn)化計算中的「交叉」操作和「創(chuàng)新檔案」等概念來發(fā)現(xiàn)具有潛力的過渡性內(nèi)核，研究人員提出的框架不僅能夠自動將PyTorch模塊轉(zhuǎn)換為CUDA內(nèi)核，而且他們的高度優(yōu)化的CUDA內(nèi)核通常能夠?qū)崿F(xiàn)顯著的性能提升。
研究團(tuán)隊相信這項技術(shù)能夠帶來性能加速，從而加快LLM或其他生成式AI模型等基礎(chǔ)大模型的訓(xùn)練和運(yùn)行（推理）速度，最終讓AI模型在英偉達(dá)GPU上運(yùn)行得更快。
- 第1和第2階段（轉(zhuǎn)換和翻譯）：「AI CUDA工程師」首先將PyTorch代碼轉(zhuǎn)換為可運(yùn)行的CUDA內(nèi)核。即便在沒有專門優(yōu)化的情況下，它們就已經(jīng)觀察到初步的性能提升。
- 第3階段（進(jìn)化優(yōu)化）：借鑒生物進(jìn)化中的適者生存法則，框架采用進(jìn)化的方法來確保只保留最優(yōu)秀的CUDA內(nèi)核。同時，團(tuán)隊創(chuàng)新性地提出了內(nèi)核交叉提示策略，能夠?qū)⒍鄠€優(yōu)化后的內(nèi)核進(jìn)行有效組合。
- 第4階段（創(chuàng)新檔案）：正如文化通過千年的文明傳承來塑造人類智慧一樣，「AI CUDA工程師」也會利用在過往創(chuàng)新和發(fā)現(xiàn)中積累的經(jīng)驗(yàn)，構(gòu)建一個由高性能CUDA內(nèi)核構(gòu)成的創(chuàng)新檔案庫，并以這些積累的優(yōu)化經(jīng)驗(yàn)為基礎(chǔ)，實(shí)現(xiàn)更進(jìn)一步的轉(zhuǎn)換和性能突破。
提升內(nèi)核運(yùn)行速度
研究團(tuán)隊相信，這僅僅是人工智能偉大優(yōu)化的開始！
在論文中，他們引入了一種端到端的智能體工作流，能夠?qū)yTorch代碼轉(zhuǎn)換為可工作的CUDA內(nèi)核，優(yōu)化CUDA運(yùn)行時性能，并自動融合多個內(nèi)核。
此外，他們構(gòu)建了各種技術(shù)來增強(qiáng)流程的一致性和性能，包括LLM集成、迭代分析反饋循環(huán)、本地內(nèi)核代碼編輯以及交叉內(nèi)核優(yōu)化。
AI發(fā)現(xiàn)的部分CUDA內(nèi)核在KernelBench上實(shí)現(xiàn)了SOTA的性能
研究人員表示，「AI CUDA工程師」能夠穩(wěn)健地轉(zhuǎn)換250個考慮的Torch操作中的230多個，并且對于大多數(shù)內(nèi)核實(shí)現(xiàn)了顯著的運(yùn)行時性能提升。此外，該方法還能夠有效地融合各種內(nèi)核操作，并且可以超越幾種現(xiàn)有的加速操作。
具體而言，在229個考慮的任務(wù)中，他們在81%的任務(wù)上優(yōu)于PyTorch原生運(yùn)行。此外，所有發(fā)現(xiàn)的CUDA內(nèi)核中有20%至少比其PyTorch實(shí)現(xiàn)快兩倍。
下圖展示了PyTorch原生內(nèi)核與經(jīng)過優(yōu)化處理后的內(nèi)核在不同級別（Level 1、Level 2、Level 3）下的加速對比。
從下圖展示的內(nèi)核可以看出，「AI CUDA工程師」能實(shí)現(xiàn)各種操作，包括歸一化方法、損失函數(shù)、特殊矩陣乘法，甚至是完整的神經(jīng)網(wǎng)絡(luò)架構(gòu)。
17,000+個已驗(yàn)證CUDA內(nèi)核數(shù)據(jù)集
研究團(tuán)隊隨論文一同發(fā)布的「資料庫」（The AI CUDA Engineer Archive）是一個包含超過3萬個CUDA內(nèi)核的數(shù)據(jù)集，這些內(nèi)核均由「AI CUDA工程師」生成。
數(shù)據(jù)集采用CC-By-4.0許可證發(fā)布，可通過HuggingFace平臺訪問，并可用于LLM的進(jìn)階微調(diào)。
其中包括，PyTorch參考實(shí)現(xiàn)、PyTorch/NCU/Clang-tidy性能分析數(shù)據(jù)、每個任務(wù)的多個內(nèi)核版本、錯誤信息，以及與PyTorch原生和編譯運(yùn)行時間的性能對比數(shù)據(jù)。
團(tuán)隊希望這個數(shù)據(jù)集能夠促進(jìn)開源模型的后續(xù)訓(xùn)練，從而開發(fā)更出色的CUDA功能模塊。包括離線強(qiáng)化學(xué)習(xí)、偏好優(yōu)化和標(biāo)準(zhǔn)監(jiān)督微調(diào)等技術(shù)。
數(shù)據(jù)集地址：https://huggingface.co/datasets/SakanaAI/AI-CUDA-Engineer-Archive
統(tǒng)計數(shù)據(jù)顯示，「資料庫」共包含超過3萬個內(nèi)核，其中已驗(yàn)證正確的實(shí)現(xiàn)超過1.7萬個，約有50%的內(nèi)核性能優(yōu)于PyTorch原生運(yùn)行時間。
基于文本嵌入進(jìn)一步分析發(fā)現(xiàn)，這些內(nèi)核可根據(jù)不同的任務(wù)類型（如矩陣乘法、池化操作、卷積運(yùn)算）和實(shí)現(xiàn)策略（包括循環(huán)展開、運(yùn)算融合、向量化處理）自然地聚類。
探索內(nèi)核資料庫
除此之外，研究人員還開發(fā)了一個交互式網(wǎng)站，供用戶實(shí)時查看超過1.7萬個已驗(yàn)證內(nèi)核及其配置信息，包括PyTorch、NCU和Clang-Tidy的相關(guān)數(shù)據(jù)。
在網(wǎng)站中，用戶可以探索涵蓋230個任務(wù)的各類高性能內(nèi)核。并且可以動手定制一個排行榜，查看不同實(shí)驗(yàn)和LLM間的相關(guān)內(nèi)核。
內(nèi)核性能排行榜
并且，還可以直觀地查看內(nèi)核可視化結(jié)果、搜索相關(guān)內(nèi)核、下載驗(yàn)證代碼和性能提升數(shù)據(jù)，以及查看詳細(xì)的性能分析結(jié)果。并且還可以深入了解優(yōu)化實(shí)驗(yàn)的具體細(xì)節(jié)。
性能分析數(shù)據(jù)、評估腳本下載、相關(guān)內(nèi)核信息和實(shí)驗(yàn)發(fā)現(xiàn)細(xì)節(jié)的實(shí)例歸一化內(nèi)核
AI革命才剛剛起步
研究人員認(rèn)為，今天的LLM就像是我們這一代的「大型機(jī)」。
我們?nèi)匀惶幵贏I發(fā)展的早期階段，在市場競爭和全球創(chuàng)新的推動下（尤其是那些在資源受限條件下進(jìn)行創(chuàng)新的團(tuán)隊），這項技術(shù)必將實(shí)現(xiàn)百萬量級的效率提升。
目前，我們的AI系統(tǒng)消耗著海量資源，如果這項技術(shù)繼續(xù)擴(kuò)張而不考慮效率和能源消耗問題，其發(fā)展將難以持續(xù)。
從根本上來說，沒有任何理由表明我們的AI系統(tǒng)不能達(dá)到甚至超越人類智能的效率水平。我們相信，要實(shí)現(xiàn)這種效率提升，最佳途徑就是利用AI來優(yōu)化AI本身。
這正是Sakana AI努力的方向，而這個項目正是朝著實(shí)現(xiàn)AI百萬倍性能提升目標(biāo)邁出的關(guān)鍵一步。
就像從早期龐大臃腫的大型機(jī)演變到現(xiàn)代計算機(jī)一樣，相比當(dāng)今「臃腫」低效的LLM，未來幾年我們使用AI的方式將會發(fā)生翻天覆地的變化。
