用AI進行GPU編程，GEAK-Agent讓小白也能上手

始智AI wisemodel.cn社區是源自中國的中立開放的AI開源社區。正在，歡迎加入共同成長。wisemodel推出邀請注冊獎勵活動，最高可得算力券+token包380元獎勵，歡迎參與和支持！

AI 任務越來越復雜，硬件也越來越多樣，現在大家都急需一種智能系統，不用手動調試，就能直接生成高性能 GPU 內核。不管是運行頂尖大語言模型的超大規模數據中心，還是學術研究用的高性能計算集群，軟件和硬件得高效配合才能落地，這種需求就更迫切了。

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

AMD 推出了 GEAK（生成高效 AI 中心 GPU 內核），一款專門針對 AMD Instinct? GPU 的自動 Triton 內核生成智能體框架。它把最先進的大語言模型和結構化推理、反饋循環結合起來，就是要把 AI 輔助代碼生成的能力再提一個檔次。

同時還發布了兩個評估基準：一個是 TritonBench 修訂版（從 TritonBench-G 改編而來，包含 184 個內核，測試工具更嚴格），另一個是全新的 ROCm Triton 基準（從 AMD 開源 ROCm 倉庫挑選 30 個真實場景內核），用來測試 AI 生成內核的正確性和運行性能。

實際測試結果很亮眼：GEAK 比直接用大語言模型生成代碼強太多 ——TritonBench 修訂版上正確生成率高達 54.89%，ROCm Triton 基準上更是達到 63.33%，而直接提示大語言模型不做智能體反饋的話，正確率還不到 15%。另外，GEAK 生成的內核在 TritonBench 修訂版上，比原來的參考版本平均快 2.59 倍。還對 ROCm Triton 基準中的一個高性能內核做了詳細研究，并且已經開源了 GEAK 智能體和評估框架。

目前，GEAK-agent已開源上線始智AI-wisemodel開源社區，8月底舉辦的AMD AI 速訓營24小時 GPU黑客松的一個賽題就是基于GEAK-Agent設計的，幾十個團隊已經在參賽時深度體驗和嘗試，歡迎更多感興趣的朋友前往體驗和嘗試。

模型地址

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

01.

挑戰與突破

隨著程序合成技術的發展，大家對“嚴格且多維度”的基準測試需求越來越迫切。現在大部分編碼基準測試都只盯著“功能對不對”，比如HumanEval是專家寫的編程任務加配套測試用例，MBPP是靠眾包收集問題，都是用手工設計的測試或獨立執行環境來評估。

為了覆蓋更多場景、提升擴展性，近幾年開始流行“自動生成測試用例”方法，這也讓它在軟件工程、開發者輔助這些領域用得更廣泛了。 但除了“能不能用”，“跑得夠不夠快”也成了基準測試的關鍵維度，越來越受重視。可現有框架大多是針對“競爭式任務”或“順序執行代碼”，對并行編程、硬件相關的評估支持特別少。

雖然有少數基準測試針對CPU并行編程，但GPU相關的寥寥無幾——要知道現在部署大規模深度學習模型，硬件效率可是越來越重要了。

為了補上這個缺口，KernelBench和TritonBench這類基準測試就應運而生了，專門評估大語言模型生成GPU內核的能力。其中KernelBench重點看不同任務的正確性和運行效率，還搞了“fast_p”這類指標；TritonBench則聚焦大語言模型用Triton領域特定語言（DSL，簡單說就是專門針對某類場景設計的語言）生成高性能代碼的難點。

這些測試都暴露了一個問題：普通代碼生成和“高性能、懂硬件”的代碼生成差距太大，必須要有專門的、面向計算的評估框架。

用于代碼生成的 LLM

近幾年很多研究都在試讓大模型寫代碼，像DeepSeekCoder、Qwen-Coder這些模型在通用編碼基準測試里表現不錯，但遇到“性能敏感的底層代碼”就不行了。

也有專門為“內核生成”訓練的模型，比如KernelLLM，但它只能把PyTorch代碼轉換成Triton內核，通用性太差。而我們的智能體不一樣，既能靠自然語言指令生成Triton代碼，也能參考已有代碼來寫，靈活性更高，還能結合上下文。另外，KernelLLM這類專用模型，沒有GPT-4.1、Claude 4這些頂尖大模型的強大推理和問題解決能力，而處理復雜優化任務時，往往要權衡各種約束，這些能力恰恰是關鍵。

推理時計算擴展

“推理時計算擴展”簡單說就是大模型在“干活（推理）的時候”，多分配點計算資源（比如給更長的上下文、多跑幾次推理、并行生成），不用額外訓練就能提升復雜任務的表現。 比如Chain-of-Thought提示法，就是讓模型“一步步思考”來增強推理能力；Reflexion框架則是把反饋記下來，指導后續決策，讓智能體越用越聰明。

GEAK智能體也用了這種思路，還加了類似Reflexion的調試模塊，能根據之前的執行反饋，一次次優化生成的代碼。

02.

GEAK技術介紹

如上圖所示，該智能AI系統包含生成器、反思器、評估器和優化器四個核心模塊：生成器根據用戶查詢與上下文信息生成代碼；評估器采用級聯設計，先通過功能測試驗證代碼正確性，未通過則將錯誤跟蹤反饋給反思器深入分析，通過后繼續評估延遲、內存效率等性能指標；反思器會針對執行不正確的代碼，結合代碼本身與錯誤跟蹤定位潛在問題；優化器則以功能正確的代碼為輸入，制定策略提升其延遲表現與計算效率，同時系統還集成了多種可組合、可配置的技術以進一步提升智能體性能。

單樣本提示（1-shot prompting）

為了實現單樣本提示，會從現有數據集中選取最相似的 Triton 代碼放入提示中。用于單樣本提示的數據集與基準測試集不重疊。我們發現，從數據集中檢索單樣本時，基于代碼相似性比基于指令相似性更有效。

知識注入（Knowledge Injection）

在提示中加入了關于編寫高效 Triton 內核的領域特定知識，包括詳細的硬件規格。這種底層優化原則的融入顯著提升了大型語言模型（LLM）生成代碼的準確性和質量。

反思機制（Reflexion）

為了增強系統的自我糾正能力，引入了Reflexion[13]模塊，該模塊支持內省式調試和迭代改進。當 LLM 生成的內核未通過功能測試時，產生的錯誤跟蹤會作為反饋提供給反思器，以進行進一步分析和糾正。智能體會負責分析失敗原因并提出相應的解決方案，通過增加迭代次數來改進內核生成。

內核智能體的 LLM 選擇

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

作為優化器的 LLM

優化器 LLM 的任務是根據之前的代碼生成結果及其相應的性能指標，確定潛在的優化方向。這些歷史記錄（包括生成的代碼和相關的性能結果）會按性能升序排列。這種結構化的呈現方式有助于引導 LLM 提出更有效的優化策略[14]。

調試陷阱（Debugging trap）

當 LLM 生成的代碼存在錯誤時，錯誤跟蹤會提供給反思器進行糾正。然而，發現有時代碼經過多次反思循環后，仍然存在相同的錯誤，這就是所說的“調試陷阱”。

為了防止智能體陷入調試陷阱，通過參數max_perf_debug_num對每個代碼片段的調試嘗試次數進行限制。如果代碼在達到該閾值后仍然失敗，智能體會被要求放棄當前方法，生成新的策略和全新的代碼。

并行縮放（Parallel Scaling）

為了確保可靠的內核生成，會并行且獨立地運行多個 GEAK 實例多次。為了增加生成代碼的多樣性，LLM 輸出會通過將溫度參數設置為 1 來進行采樣。

實驗表明，這種生成多樣性可以產生原本未被充分探索的正確且更快的內核。此外，進一步的實驗表明，結合順序縮放和并行縮放可以帶來額外的性能提升。為了獲得無偏的準確率估計，使用 pass@k 指標進行評估。

03.

基準測試

團隊實驗了 Triton 內核評估基準測試，包括從 TritonBench-G 修改的 184 個內核集和從各種開源 ROCm 存儲庫中選取的 30 個內核集。

TritonBench-revised 基準測試

Triton 內核和相應的評估代碼庫借鑒并改編自 TritonBench，但做了以下增強：TritonBench-G 的184個內核中有37個在AMD GPU上運行時有錯誤（如共享內存錯誤、無效的 HIP 參數（HIP是AMD的異構計算接口，類似于NVIDIA的CUDA）、ModuleNotFound），修復了這些內核以兼容 AMD GPU。

對于一些內核，原始 TritonBench 存儲庫缺少調用函數，導致執行準確性腳本實際上比較的是空字符串輸出。這一問題通過以下方式得到糾正：為缺少調用函數的內核添加測試函數調用。開發基于容差的測試函數輸出張量比較，代替 STDOUT 字符串比較。確保單元測試中用于生成輸入的隨機張量生成具有一致的種子。對于一些內核，單元測試編寫不一致，導致輸出不符合預期，修復了這些內核，使其具有一致的單元測試。

ROCm Triton 基準測試

以下內核由 Triton 專家工程師編寫并公開，以幫助圍繞AMD GPU高效運行AI工作負載構建生態系統，借助前沿LLM重構這些內核及其相應單元測試，使其格式與TritonBench-revised基準測試一致。這樣做能夠以一致的方式比較智能體的準確性和效率。

評估指標

團隊采用了與 TritonBench 一致的評估指標：調用準確性指 AI 生成內核可成功編譯且無錯誤運行的比例；執行準確性僅針對編譯成功的內核計算，即滿足所有單元測試的內核占比；加速比則僅針對通過全部單元測試的內核，定義為所有單元測試中參考基準內核延遲中位數與生成內核延遲中位數之比的期望值，用于衡量生成內核相對參考內核的執行時間改進效果。

04.

實驗結果概要

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

GEAK 系統表現顯著優于直接 LLM 提示：在 MI300 上，TritonBench-revised 基準測試中實現 54.89% 執行準確率和 2.59 倍加速比，ROCm 基準測試中達 63.33% 執行準確率和 0.92 倍加速比。

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

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

ROCm 基準測試中的推理時間計算擴展性研究

TritonBench 修改基準測試的推理時間計算擴展性研究

----- END -----

wisemodel相關：

系列模型：

關于wisemodel更多

1

歡迎持續關注和支持

開源社區建設需要長期堅持和投入，更需要廣大用戶的積極參與、貢獻和維護，歡迎大家加入wisemodel開源社區的志愿者計劃和開源共創計劃。期待更多開發者將開源成果，包括模型、數據集和代碼等發布到 wisemodel.cn 社區，共建中立、開放的AI開源社區生態。歡迎掃碼添加wisemodel微信，申請加入wisemodel社群，持續關注wisemodel.cn開源社區動態。

2

歡迎加盟wisemodel開源社區

始智AI wisemodel社區自2023年9月上線以來，逐漸成為影響力日益擴大的中立開放的AI開源社區，為了加快公司發展，我們長期需要技術、運營等人才加盟，技術側重在AI infra、后端開發，熟悉K8S、模型訓練和推理等技術， 以及熟悉開發者生態運營的成員，歡迎感興趣的朋友加盟，可以通過添加wisemodel微信，或者將簡歷投遞到郵箱：liudaoquan@wisemodel.cn

3

歡迎投稿優質內容

歡迎投稿分享人工智能領域相關的優秀研究成果，鼓勵高校實驗室、大企業研究團隊、個人等，在wisemodel平臺上分享各類優質內容，可以是AI領域最新論文解讀、最新開源成果介紹，也可以是關于AI技術實踐、應用和總結等。投稿可以發郵件到liudaoquan@wisemodel.cn，也可以掃碼添加wisemodel微信。

4

關于wisemodel開源社區

始智AI wisemodel.cn開源社區由清華校友總會AI大數據專委會副秘書長劉道全創立，旨在打造和建設中立開放的AI開源創新社區，將打造成“HuggingFace”之外最活躍的AI開源社區，匯聚主要AI開源模型、數據集和代碼等，歡迎高校科研院所、大型互聯網公司、創新創業企業、廣大個人開發者，以及政府部門、學會協會、聯盟、基金會等，還有投資機構、科技媒體等，共同參與建設AI開源創新生態。

向上滑動查看