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