用AI進(jìn)行GPU編程，GEAK-Agent讓小白也能上手

始智AI wisemodel.cn社區(qū)是源自中國(guó)的中立開(kāi)放的AI開(kāi)源社區(qū)。正在，歡迎加入共同成長(zhǎng)。wisemodel推出邀請(qǐng)注冊(cè)獎(jiǎng)勵(lì)活動(dòng)，最高可得算力券+token包380元獎(jiǎng)勵(lì)，歡迎參與和支持！

AI 任務(wù)越來(lái)越復(fù)雜，硬件也越來(lái)越多樣，現(xiàn)在大家都急需一種智能系統(tǒng)，不用手動(dòng)調(diào)試，就能直接生成高性能 GPU 內(nèi)核。不管是運(yùn)行頂尖大語(yǔ)言模型的超大規(guī)模數(shù)據(jù)中心，還是學(xué)術(shù)研究用的高性能計(jì)算集群，軟件和硬件得高效配合才能落地，這種需求就更迫切了。

雖然現(xiàn)在大語(yǔ)言模型在代碼生成上挺厲害，但要寫(xiě)出能正確運(yùn)行、還高效的 GPU 代碼，依然是個(gè)大難題。而 AI 輔助開(kāi)發(fā) GPU 內(nèi)核（尤其是針對(duì) AMD Instinct? MI300X 這類新型硬件），能大大提升開(kāi)發(fā)效率，加快部署速度。

AMD 推出了 GEAK（生成高效 AI 中心 GPU 內(nèi)核），一款專門(mén)針對(duì) AMD Instinct? GPU 的自動(dòng) Triton 內(nèi)核生成智能體框架。它把最先進(jìn)的大語(yǔ)言模型和結(jié)構(gòu)化推理、反饋循環(huán)結(jié)合起來(lái)，就是要把 AI 輔助代碼生成的能力再提一個(gè)檔次。

同時(shí)還發(fā)布了兩個(gè)評(píng)估基準(zhǔn)：一個(gè)是 TritonBench 修訂版（從 TritonBench-G 改編而來(lái)，包含 184 個(gè)內(nèi)核，測(cè)試工具更嚴(yán)格），另一個(gè)是全新的 ROCm Triton 基準(zhǔn)（從 AMD 開(kāi)源 ROCm 倉(cāng)庫(kù)挑選 30 個(gè)真實(shí)場(chǎng)景內(nèi)核），用來(lái)測(cè)試 AI 生成內(nèi)核的正確性和運(yùn)行性能。

實(shí)際測(cè)試結(jié)果很亮眼：GEAK 比直接用大語(yǔ)言模型生成代碼強(qiáng)太多 ——TritonBench 修訂版上正確生成率高達(dá) 54.89%，ROCm Triton 基準(zhǔn)上更是達(dá)到 63.33%，而直接提示大語(yǔ)言模型不做智能體反饋的話，正確率還不到 15%。另外，GEAK 生成的內(nèi)核在 TritonBench 修訂版上，比原來(lái)的參考版本平均快 2.59 倍。還對(duì) ROCm Triton 基準(zhǔn)中的一個(gè)高性能內(nèi)核做了詳細(xì)研究，并且已經(jīng)開(kāi)源了 GEAK 智能體和評(píng)估框架。

目前，GEAK-agent已開(kāi)源上線始智AI-wisemodel開(kāi)源社區(qū)，8月底舉辦的AMD AI 速訓(xùn)營(yíng)24小時(shí) GPU黑客松的一個(gè)賽題就是基于GEAK-Agent設(shè)計(jì)的，幾十個(gè)團(tuán)隊(duì)已經(jīng)在參賽時(shí)深度體驗(yàn)和嘗試，歡迎更多感興趣的朋友前往體驗(yàn)和嘗試。

模型地址

https://wisemodel.cn/codes/motata/GEAK-agent

01.

挑戰(zhàn)與突破

隨著程序合成技術(shù)的發(fā)展，大家對(duì)“嚴(yán)格且多維度”的基準(zhǔn)測(cè)試需求越來(lái)越迫切。現(xiàn)在大部分編碼基準(zhǔn)測(cè)試都只盯著“功能對(duì)不對(duì)”，比如HumanEval是專家寫(xiě)的編程任務(wù)加配套測(cè)試用例，MBPP是靠眾包收集問(wèn)題，都是用手工設(shè)計(jì)的測(cè)試或獨(dú)立執(zhí)行環(huán)境來(lái)評(píng)估。

為了覆蓋更多場(chǎng)景、提升擴(kuò)展性，近幾年開(kāi)始流行“自動(dòng)生成測(cè)試用例”方法，這也讓它在軟件工程、開(kāi)發(fā)者輔助這些領(lǐng)域用得更廣泛了。 但除了“能不能用”，“跑得夠不夠快”也成了基準(zhǔn)測(cè)試的關(guān)鍵維度，越來(lái)越受重視。可現(xiàn)有框架大多是針對(duì)“競(jìng)爭(zhēng)式任務(wù)”或“順序執(zhí)行代碼”，對(duì)并行編程、硬件相關(guān)的評(píng)估支持特別少。

雖然有少數(shù)基準(zhǔn)測(cè)試針對(duì)CPU并行編程，但GPU相關(guān)的寥寥無(wú)幾——要知道現(xiàn)在部署大規(guī)模深度學(xué)習(xí)模型，硬件效率可是越來(lái)越重要了。

為了補(bǔ)上這個(gè)缺口，KernelBench和TritonBench這類基準(zhǔn)測(cè)試就應(yīng)運(yùn)而生了，專門(mén)評(píng)估大語(yǔ)言模型生成GPU內(nèi)核的能力。其中KernelBench重點(diǎn)看不同任務(wù)的正確性和運(yùn)行效率，還搞了“fast_p”這類指標(biāo)；TritonBench則聚焦大語(yǔ)言模型用Triton領(lǐng)域特定語(yǔ)言（DSL，簡(jiǎn)單說(shuō)就是專門(mén)針對(duì)某類場(chǎng)景設(shè)計(jì)的語(yǔ)言）生成高性能代碼的難點(diǎn)。

這些測(cè)試都暴露了一個(gè)問(wèn)題：普通代碼生成和“高性能、懂硬件”的代碼生成差距太大，必須要有專門(mén)的、面向計(jì)算的評(píng)估框架。

用于代碼生成的 LLM

近幾年很多研究都在試讓大模型寫(xiě)代碼，像DeepSeekCoder、Qwen-Coder這些模型在通用編碼基準(zhǔn)測(cè)試?yán)锉憩F(xiàn)不錯(cuò)，但遇到“性能敏感的底層代碼”就不行了。

也有專門(mén)為“內(nèi)核生成”訓(xùn)練的模型，比如KernelLLM，但它只能把PyTorch代碼轉(zhuǎn)換成Triton內(nèi)核，通用性太差。而我們的智能體不一樣，既能靠自然語(yǔ)言指令生成Triton代碼，也能參考已有代碼來(lái)寫(xiě)，靈活性更高，還能結(jié)合上下文。另外，KernelLLM這類專用模型，沒(méi)有GPT-4.1、Claude 4這些頂尖大模型的強(qiáng)大推理和問(wèn)題解決能力，而處理復(fù)雜優(yōu)化任務(wù)時(shí)，往往要權(quán)衡各種約束，這些能力恰恰是關(guān)鍵。

推理時(shí)計(jì)算擴(kuò)展

“推理時(shí)計(jì)算擴(kuò)展”簡(jiǎn)單說(shuō)就是大模型在“干活（推理）的時(shí)候”，多分配點(diǎn)計(jì)算資源（比如給更長(zhǎng)的上下文、多跑幾次推理、并行生成），不用額外訓(xùn)練就能提升復(fù)雜任務(wù)的表現(xiàn)。 比如Chain-of-Thought提示法，就是讓模型“一步步思考”來(lái)增強(qiáng)推理能力；Reflexion框架則是把反饋記下來(lái)，指導(dǎo)后續(xù)決策，讓智能體越用越聰明。

GEAK智能體也用了這種思路，還加了類似Reflexion的調(diào)試模塊，能根據(jù)之前的執(zhí)行反饋，一次次優(yōu)化生成的代碼。

02.

GEAK技術(shù)介紹

如上圖所示，該智能AI系統(tǒng)包含生成器、反思器、評(píng)估器和優(yōu)化器四個(gè)核心模塊：生成器根據(jù)用戶查詢與上下文信息生成代碼；評(píng)估器采用級(jí)聯(lián)設(shè)計(jì)，先通過(guò)功能測(cè)試驗(yàn)證代碼正確性，未通過(guò)則將錯(cuò)誤跟蹤反饋給反思器深入分析，通過(guò)后繼續(xù)評(píng)估延遲、內(nèi)存效率等性能指標(biāo)；反思器會(huì)針對(duì)執(zhí)行不正確的代碼，結(jié)合代碼本身與錯(cuò)誤跟蹤定位潛在問(wèn)題；優(yōu)化器則以功能正確的代碼為輸入，制定策略提升其延遲表現(xiàn)與計(jì)算效率，同時(shí)系統(tǒng)還集成了多種可組合、可配置的技術(shù)以進(jìn)一步提升智能體性能。

單樣本提示（1-shot prompting）

為了實(shí)現(xiàn)單樣本提示，會(huì)從現(xiàn)有數(shù)據(jù)集中選取最相似的 Triton 代碼放入提示中。用于單樣本提示的數(shù)據(jù)集與基準(zhǔn)測(cè)試集不重疊。我們發(fā)現(xiàn)，從數(shù)據(jù)集中檢索單樣本時(shí)，基于代碼相似性比基于指令相似性更有效。

知識(shí)注入（Knowledge Injection）

在提示中加入了關(guān)于編寫(xiě)高效 Triton 內(nèi)核的領(lǐng)域特定知識(shí)，包括詳細(xì)的硬件規(guī)格。這種底層優(yōu)化原則的融入顯著提升了大型語(yǔ)言模型（LLM）生成代碼的準(zhǔn)確性和質(zhì)量。

反思機(jī)制（Reflexion）

為了增強(qiáng)系統(tǒng)的自我糾正能力，引入了Reflexion[13]模塊，該模塊支持內(nèi)省式調(diào)試和迭代改進(jìn)。當(dāng) LLM 生成的內(nèi)核未通過(guò)功能測(cè)試時(shí)，產(chǎn)生的錯(cuò)誤跟蹤會(huì)作為反饋提供給反思器，以進(jìn)行進(jìn)一步分析和糾正。智能體會(huì)負(fù)責(zé)分析失敗原因并提出相應(yīng)的解決方案，通過(guò)增加迭代次數(shù)來(lái)改進(jìn)內(nèi)核生成。

內(nèi)核智能體的 LLM 選擇

嘗試了多種 LLM，如 GPT-4.1、O1 和 Gemini 2.5 Pro，觀察到不同模型產(chǎn)生的輸出存在顯著差異，這表明底層 LLM 的能力會(huì)對(duì)結(jié)果產(chǎn)生重大影響。

作為優(yōu)化器的 LLM

優(yōu)化器 LLM 的任務(wù)是根據(jù)之前的代碼生成結(jié)果及其相應(yīng)的性能指標(biāo)，確定潛在的優(yōu)化方向。這些歷史記錄（包括生成的代碼和相關(guān)的性能結(jié)果）會(huì)按性能升序排列。這種結(jié)構(gòu)化的呈現(xiàn)方式有助于引導(dǎo) LLM 提出更有效的優(yōu)化策略[14]。

調(diào)試陷阱（Debugging trap）

當(dāng) LLM 生成的代碼存在錯(cuò)誤時(shí)，錯(cuò)誤跟蹤會(huì)提供給反思器進(jìn)行糾正。然而，發(fā)現(xiàn)有時(shí)代碼經(jīng)過(guò)多次反思循環(huán)后，仍然存在相同的錯(cuò)誤，這就是所說(shuō)的“調(diào)試陷阱”。

為了防止智能體陷入調(diào)試陷阱，通過(guò)參數(shù)max_perf_debug_num對(duì)每個(gè)代碼片段的調(diào)試嘗試次數(shù)進(jìn)行限制。如果代碼在達(dá)到該閾值后仍然失敗，智能體會(huì)被要求放棄當(dāng)前方法，生成新的策略和全新的代碼。

并行縮放（Parallel Scaling）

為了確保可靠的內(nèi)核生成，會(huì)并行且獨(dú)立地運(yùn)行多個(gè) GEAK 實(shí)例多次。為了增加生成代碼的多樣性，LLM 輸出會(huì)通過(guò)將溫度參數(shù)設(shè)置為 1 來(lái)進(jìn)行采樣。

實(shí)驗(yàn)表明，這種生成多樣性可以產(chǎn)生原本未被充分探索的正確且更快的內(nèi)核。此外，進(jìn)一步的實(shí)驗(yàn)表明，結(jié)合順序縮放和并行縮放可以帶來(lái)額外的性能提升。為了獲得無(wú)偏的準(zhǔn)確率估計(jì)，使用 pass@k 指標(biāo)進(jìn)行評(píng)估。

03.

基準(zhǔn)測(cè)試

團(tuán)隊(duì)實(shí)驗(yàn)了 Triton 內(nèi)核評(píng)估基準(zhǔn)測(cè)試，包括從 TritonBench-G 修改的 184 個(gè)內(nèi)核集和從各種開(kāi)源 ROCm 存儲(chǔ)庫(kù)中選取的 30 個(gè)內(nèi)核集。

TritonBench-revised 基準(zhǔn)測(cè)試

Triton 內(nèi)核和相應(yīng)的評(píng)估代碼庫(kù)借鑒并改編自 TritonBench，但做了以下增強(qiáng)：TritonBench-G 的184個(gè)內(nèi)核中有37個(gè)在AMD GPU上運(yùn)行時(shí)有錯(cuò)誤（如共享內(nèi)存錯(cuò)誤、無(wú)效的 HIP 參數(shù)（HIP是AMD的異構(gòu)計(jì)算接口，類似于NVIDIA的CUDA）、ModuleNotFound），修復(fù)了這些內(nèi)核以兼容 AMD GPU。

對(duì)于一些內(nèi)核，原始 TritonBench 存儲(chǔ)庫(kù)缺少調(diào)用函數(shù)，導(dǎo)致執(zhí)行準(zhǔn)確性腳本實(shí)際上比較的是空字符串輸出。這一問(wèn)題通過(guò)以下方式得到糾正：為缺少調(diào)用函數(shù)的內(nèi)核添加測(cè)試函數(shù)調(diào)用。開(kāi)發(fā)基于容差的測(cè)試函數(shù)輸出張量比較，代替 STDOUT 字符串比較。確保單元測(cè)試中用于生成輸入的隨機(jī)張量生成具有一致的種子。對(duì)于一些內(nèi)核，單元測(cè)試編寫(xiě)不一致，導(dǎo)致輸出不符合預(yù)期，修復(fù)了這些內(nèi)核，使其具有一致的單元測(cè)試。

ROCm Triton 基準(zhǔn)測(cè)試

以下內(nèi)核由 Triton 專家工程師編寫(xiě)并公開(kāi)，以幫助圍繞AMD GPU高效運(yùn)行AI工作負(fù)載構(gòu)建生態(tài)系統(tǒng)，借助前沿LLM重構(gòu)這些內(nèi)核及其相應(yīng)單元測(cè)試，使其格式與TritonBench-revised基準(zhǔn)測(cè)試一致。這樣做能夠以一致的方式比較智能體的準(zhǔn)確性和效率。

評(píng)估指標(biāo)

團(tuán)隊(duì)采用了與 TritonBench 一致的評(píng)估指標(biāo)：調(diào)用準(zhǔn)確性指 AI 生成內(nèi)核可成功編譯且無(wú)錯(cuò)誤運(yùn)行的比例；執(zhí)行準(zhǔn)確性僅針對(duì)編譯成功的內(nèi)核計(jì)算，即滿足所有單元測(cè)試的內(nèi)核占比；加速比則僅針對(duì)通過(guò)全部單元測(cè)試的內(nèi)核，定義為所有單元測(cè)試中參考基準(zhǔn)內(nèi)核延遲中位數(shù)與生成內(nèi)核延遲中位數(shù)之比的期望值，用于衡量生成內(nèi)核相對(duì)參考內(nèi)核的執(zhí)行時(shí)間改進(jìn)效果。

04.

實(shí)驗(yàn)結(jié)果概要

以 GPT-4.1、Gemini 2.5 Pro 等前沿 LLM 直接提示生成內(nèi)核作為基準(zhǔn)，其執(zhí)行準(zhǔn)確率低、性能不佳，非 CUDA 平臺(tái)（如 ROCm）挑戰(zhàn)尤為突出；引入單樣本提示后，準(zhǔn)確率和部分場(chǎng)景加速比有所提升，但存在 Gemini 2.5 Pro 加速比下降的例外情況。

GEAK 系統(tǒng)表現(xiàn)顯著優(yōu)于直接 LLM 提示：在 MI300 上，TritonBench-revised 基準(zhǔn)測(cè)試中實(shí)現(xiàn) 54.89% 執(zhí)行準(zhǔn)確率和 2.59 倍加速比，ROCm 基準(zhǔn)測(cè)試中達(dá) 63.33% 執(zhí)行準(zhǔn)確率和 0.92 倍加速比。

研究顯示，GEAK 生成的triton_test_flip.py內(nèi)核比專家編寫(xiě)版本快2.26倍，核心優(yōu)勢(shì)在于優(yōu)化內(nèi)存訪問(wèn)模式（單遍操作、合并訪問(wèn)）、降低寄存器壓力、提升緩存效率及顯式處理邊界條件。

消融實(shí)驗(yàn)結(jié)果顯示：順序縮放方面，增加迭代次數(shù)可使GEAK的準(zhǔn)確率和性能穩(wěn)步提升，19次迭代時(shí)調(diào)用準(zhǔn)確率達(dá)63.04%、執(zhí)行準(zhǔn)確率44.02%，加速比長(zhǎng)期保持在1.4倍以上；并行縮放方面，固定10次順序迭代后，并行運(yùn)行次數(shù)的增加使準(zhǔn)確率呈近乎對(duì)數(shù)線性增長(zhǎng)，與順序改進(jìn)形成互補(bǔ)，且在MI250、MI300上均展現(xiàn)出良好的魯棒性和可移植性；模塊影響方面，知識(shí)注入、單樣本提示、優(yōu)化器三大模塊均發(fā)揮積極作用，其中知識(shí)注入顯著提升準(zhǔn)確率，單樣本提示進(jìn)一步優(yōu)化正確性，優(yōu)化器則大幅提升加速比，三者結(jié)合時(shí)執(zhí)行準(zhǔn)確率達(dá)40.76%，加速比達(dá)1.45倍。

ROCm 基準(zhǔn)測(cè)試中的推理時(shí)間計(jì)算擴(kuò)展性研究

TritonBench 修改基準(zhǔn)測(cè)試的推理時(shí)間計(jì)算擴(kuò)展性研究

----- END -----

wisemodel相關(guān)：

系列模型：

關(guān)于wisemodel更多

1

歡迎持續(xù)關(guān)注和支持

開(kāi)源社區(qū)建設(shè)需要長(zhǎng)期堅(jiān)持和投入，更需要廣大用戶的積極參與、貢獻(xiàn)和維護(hù)，歡迎大家加入wisemodel開(kāi)源社區(qū)的志愿者計(jì)劃和開(kāi)源共創(chuàng)計(jì)劃。期待更多開(kāi)發(fā)者將開(kāi)源成果，包括模型、數(shù)據(jù)集和代碼等發(fā)布到 wisemodel.cn 社區(qū)，共建中立、開(kāi)放的AI開(kāi)源社區(qū)生態(tài)。歡迎掃碼添加wisemodel微信，申請(qǐng)加入wisemodel社群，持續(xù)關(guān)注wisemodel.cn開(kāi)源社區(qū)動(dòng)態(tài)。

2

歡迎加盟wisemodel開(kāi)源社區(qū)

始智AI wisemodel社區(qū)自2023年9月上線以來(lái)，逐漸成為影響力日益擴(kuò)大的中立開(kāi)放的AI開(kāi)源社區(qū)，為了加快公司發(fā)展，我們長(zhǎng)期需要技術(shù)、運(yùn)營(yíng)等人才加盟，技術(shù)側(cè)重在AI infra、后端開(kāi)發(fā)，熟悉K8S、模型訓(xùn)練和推理等技術(shù)， 以及熟悉開(kāi)發(fā)者生態(tài)運(yùn)營(yíng)的成員，歡迎感興趣的朋友加盟，可以通過(guò)添加wisemodel微信，或者將簡(jiǎn)歷投遞到郵箱：liudaoquan@wisemodel.cn

3

歡迎投稿優(yōu)質(zhì)內(nèi)容

歡迎投稿分享人工智能領(lǐng)域相關(guān)的優(yōu)秀研究成果，鼓勵(lì)高校實(shí)驗(yàn)室、大企業(yè)研究團(tuán)隊(duì)、個(gè)人等，在wisemodel平臺(tái)上分享各類優(yōu)質(zhì)內(nèi)容，可以是AI領(lǐng)域最新論文解讀、最新開(kāi)源成果介紹，也可以是關(guān)于AI技術(shù)實(shí)踐、應(yīng)用和總結(jié)等。投稿可以發(fā)郵件到liudaoquan@wisemodel.cn，也可以掃碼添加wisemodel微信。

4

關(guān)于wisemodel開(kāi)源社區(qū)

始智AI wisemodel.cn開(kāi)源社區(qū)由清華校友總會(huì)AI大數(shù)據(jù)專委會(huì)副秘書(shū)長(zhǎng)劉道全創(chuàng)立，旨在打造和建設(shè)中立開(kāi)放的AI開(kāi)源創(chuàng)新社區(qū)，將打造成“HuggingFace”之外最活躍的AI開(kāi)源社區(qū)，匯聚主要AI開(kāi)源模型、數(shù)據(jù)集和代碼等，歡迎高校科研院所、大型互聯(lián)網(wǎng)公司、創(chuàng)新創(chuàng)業(yè)企業(yè)、廣大個(gè)人開(kāi)發(fā)者，以及政府部門(mén)、學(xué)會(huì)協(xié)會(huì)、聯(lián)盟、基金會(huì)等，還有投資機(jī)構(gòu)、科技媒體等，共同參與建設(shè)AI開(kāi)源創(chuàng)新生態(tài)。

向上滑動(dòng)查看