機(jī)器之心編輯部

近日，來自字節(jié)跳動 Seed 團(tuán)隊(duì)和清華大學(xué) AIR的新研究CUDA Agent，在 AI 領(lǐng)域引發(fā)了不小的轟動。

研究人員訓(xùn)練了一個(gè)能夠編寫快速 CUDA 內(nèi)核的模型：不只是正確的內(nèi)核，而是真正經(jīng)過優(yōu)化的內(nèi)核。

在簡單/中等內(nèi)核上，它的性能比 torch.compile高出 2 倍；在復(fù)雜內(nèi)核上，它的性能比 torch.compile高出約 92%；即使在最難的設(shè)置下，它的性能也比 Claude Opus 4.5 和 Gemini 3 Pro高出約 40%

• 論文鏈接：https://arxiv.org/abs/2602.24286
• 項(xiàng)目主頁：https://cuda-agent.github.io/
• Github鏈接：https://github.com/BytedTsinghua-SIA/CUDA-Agent
• 數(shù)據(jù)集鏈接：https://huggingface.co/datasets/BytedTsinghua-SIA/CUDA-Agent-Ops-6K

在此之前，GPT、Claude 等大模型已經(jīng)能寫出「正確」的 CUDA 代碼，AI 生成的代碼也已獲得了一定程度的應(yīng)用，但能跑通和跑得快完全是兩碼事。

GPU 內(nèi)核優(yōu)化是現(xiàn)代深度學(xué)習(xí)的基礎(chǔ)，但它仍然是一項(xiàng)高度專業(yè)化的工作，需要深厚的硬件知識。現(xiàn)有的 AI 生成 CUDA 代碼通常依賴無訓(xùn)練的提示工程（Prompting）或多輪執(zhí)行反饋微調(diào)的機(jī)制。這導(dǎo)致模型只能解決表面上的語法錯(cuò)誤，無法真正理解底層硬件邏輯，一定程度上限制了其內(nèi)在的優(yōu)化能力。

真正極致的 CUDA 優(yōu)化需要處理的任務(wù)，是只有在性能分析器中才能看到的硬件級指標(biāo)。人們一直期待能出現(xiàn)一個(gè)像人類 CUDA 專家一樣思考的 AI。

針對這一矛盾，CUDA Agent 的核心理念簡單而巧妙：CUDA 性能并非取決于正確性，而是取決于硬件。線程束、內(nèi)存帶寬、內(nèi)存沖突——這些只有在性能分析器中才能看到的東西。

研究人員不再獎勵(lì)「是否編譯成功」，而是獎勵(lì)實(shí)際的GPU速度。真實(shí)的性能分析數(shù)據(jù)。強(qiáng)化學(xué)習(xí)直接基于性能進(jìn)行訓(xùn)練。

產(chǎn)生的效果出乎人們的預(yù)料。

在 KernelBench 基準(zhǔn)測試上，CUDA Agent 取得了 SOTA 的成績：在 Level-1、Level-2 和 Level-3 三個(gè)劃分上，相比 torch.compile 分別實(shí)現(xiàn)了 100%、100% 和 92% 的加速比例（faster rate）。

CUDA Agent 與 torch.compile 和強(qiáng)大的專有模型在 KernelBench 上的對比。

簡而言之，CUDA Agent是一個(gè)大規(guī)模的智能體強(qiáng)化學(xué)習(xí)系統(tǒng)，包含三個(gè)核心組成部分：可擴(kuò)展的數(shù)據(jù)合成機(jī)制、一個(gè)集成技能增強(qiáng)且具備可靠驗(yàn)證與性能分析能力的 CUDA 開發(fā)環(huán)境，以及用于穩(wěn)定長上下文訓(xùn)練的強(qiáng)化學(xué)習(xí)算法技術(shù)。

此外，研究團(tuán)隊(duì)同時(shí)發(fā)布了CUDA-Agent-Ops-6K，一個(gè)經(jīng)過嚴(yán)格篩選與數(shù)據(jù)污染控制的高質(zhì)量合成訓(xùn)練數(shù)據(jù)集，可支持基于強(qiáng)化學(xué)習(xí)的 CUDA 內(nèi)核優(yōu)化研究的復(fù)現(xiàn)。

系統(tǒng)管線設(shè)計(jì)

數(shù)據(jù)合成

研究團(tuán)隊(duì)通過一個(gè)三階段的管線來構(gòu)建訓(xùn)練任務(wù)：種子問題爬取、基于 LLM 的組合式合成，以及基于執(zhí)行結(jié)果的篩選。

• 從 torch 和 transformers中挖掘種子算子。每個(gè)算子都以一個(gè) Python 類的形式表示，包含初始化和前向傳播方法。
• 在組合式合成階段，最多采樣 5 個(gè) torch 算子，并將它們按順序組合，構(gòu)造成融合任務(wù)。
• 篩選階段僅保留那些在 eager 模式和 compile 模式下都能正常運(yùn)行的任務(wù)，同時(shí)移除包含隨機(jī)性的算子。
• 為防止「投機(jī)取巧」，剔除在不同輸入下輸出為常數(shù)或無法區(qū)分的任務(wù)。
• 在工作負(fù)載控制方面，將 eager 模式下的運(yùn)行時(shí)間限制在 1ms–100ms 區(qū)間內(nèi)，并移除與 KernelBench 高度相似的樣本。

三階段數(shù)據(jù)收集管線

最終整理得到 6000 條訓(xùn)練樣本，構(gòu)建了 CUDA-Agent-Ops-6K 數(shù)據(jù)集，該數(shù)據(jù)集專為可擴(kuò)展的強(qiáng)化學(xué)習(xí)訓(xùn)練而設(shè)計(jì)，兼具廣泛的任務(wù)多樣性和較低的數(shù)據(jù)污染風(fēng)險(xiǎn)。

智能體環(huán)境

智能體循環(huán)管線遵循一種 ReAct 風(fēng)格的工作流，結(jié)合代碼工具與 CUDA Skill 規(guī)范（SKILL.md），支持迭代式的編碼-編譯-調(diào)試循環(huán)，以及基于性能分析器的優(yōu)化過程。

• 標(biāo)準(zhǔn)工作流程：對原生 PyTorch 實(shí)現(xiàn)進(jìn)行性能分析，編寫 CUDA 內(nèi)核及其綁定代碼，在 GPU 沙盒環(huán)境中完成編譯，并不斷迭代優(yōu)化。
• 目標(biāo)要求：通過正確性檢查，并在性能上相對于 torch.compile 實(shí)現(xiàn)超過 5% 的加速。
• 穩(wěn)健的獎勵(lì)機(jī)制采用基于里程碑的離散獎勵(lì)設(shè)計(jì)，根據(jù)正確性達(dá)標(biāo)情況和性能提升幅度分別給予獎勵(lì)。
• 防止獎勵(lì)作弊的控制措施包括：對驗(yàn)證與性能分析腳本進(jìn)行保護(hù)，禁止回退調(diào)用，采用 5 組不同輸入進(jìn)行正確性檢查，在同步預(yù)熱后進(jìn)行性能分析，以及禁止網(wǎng)絡(luò)檢索。

這些約束共同構(gòu)建了可靠的、基于真實(shí)執(zhí)行結(jié)果的反饋機(jī)制，使策略學(xué)習(xí)聚焦于內(nèi)核質(zhì)量的實(shí)質(zhì)性提升，而非依賴取巧或捷徑行為。

訓(xùn)練流程

訓(xùn)練過程采用分階段設(shè)計(jì)，以穩(wěn)定 CUDA 編碼這一長時(shí)序強(qiáng)化學(xué)習(xí)任務(wù)。首先進(jìn)行單輪 PPO 預(yù)熱訓(xùn)練，隨后分別初始化 actor 和 critic，最后進(jìn)入完整的多輪智能體強(qiáng)化學(xué)習(xí)階段。

• 單輪預(yù)熱階段旨在提升基礎(chǔ)的 CUDA 代碼生成能力，為后續(xù)的交互式智能體訓(xùn)練打下基礎(chǔ)。
• 在 actor 初始化階段，采用基于正向結(jié)果軌跡采樣的拒絕式微調(diào)（RFT）。
• RFT 過濾機(jī)制會剔除低效循環(huán)以及無效的工具調(diào)用模式，從而降低策略崩潰的風(fēng)險(xiǎn)。
• critic 初始化階段通過價(jià)值函數(shù)預(yù)訓(xùn)練，使得從訓(xùn)練早期開始，優(yōu)勢估計(jì)就具備較高可靠性。

借助這一多階段訓(xùn)練設(shè)計(jì)，系統(tǒng)在長上下文設(shè)定下（最長 128k 上下文、訓(xùn)練階段最多 150 輪、評估階段最多 200 輪）依然保持穩(wěn)定，從而實(shí)現(xiàn)持續(xù)的獎勵(lì)增長。

核心實(shí)驗(yàn)結(jié)果

研究團(tuán)隊(duì)在 KernelBench 上報(bào)告了針對整體和 Level-3 拆分的完整指標(biāo)，包括通過率、提速率（與 Eager 對比/與 Compile 對比）以及幾何平均加速比（與 Eager 對比/與 Compile 對比）。

與強(qiáng)大的專有基線模型相比，CUDA Agent 在相對于 Compile 的性能優(yōu)化上展現(xiàn)出顯著優(yōu)勢：在整體 KernelBench 基準(zhǔn)測試中，其相對 Compile 的加速達(dá)成率達(dá)到 96.8%，幾何平均加速比為 2.11 倍

這一優(yōu)勢在高難度設(shè)置下尤為明顯：在 Level-3 上，CUDA Agent 相對 compile 的加速達(dá)成率達(dá)到 90%，相比最強(qiáng)的專有基線高出約 40 個(gè)百分點(diǎn)；在 Level-2 的算子序列任務(wù)上，其加速達(dá)成率達(dá)到 100%，幾何平均加速比達(dá)到 2.80 倍。

在 KernelBench 上的整體性能和加速指標(biāo)。

本研究存在兩個(gè)主要局限。

首先，此次研究未將 CUDA Agent 與更為復(fù)雜的編譯器框架（如 TVM）進(jìn)行對比。其次，訓(xùn)練流程依賴于大規(guī)模 GPU 資源池以及進(jìn)程級隔離機(jī)制，這帶來了相當(dāng)可觀的計(jì)算與工程成本。探索更加資源高效的訓(xùn)練策略，將是未來的重要研究方向。

看起來，CUDA Agent 等技術(shù)的出現(xiàn)即將打破傳統(tǒng)編譯器（如torch.compile或Triton）的優(yōu)化瓶頸。它證明了：大語言模型不僅可以學(xué)習(xí)人類自然語言和高級編程語言，還可以通過基于硬件反饋的強(qiáng)化學(xué)習(xí)，內(nèi)化出極高門檻的「硬件直覺」。

一條通向全自動、高度性能優(yōu)化計(jì)算基礎(chǔ)設(shè)施的道路正在出現(xiàn)。