英偉達護城河被AI攻破，字節清華CUDA Agent，讓人人能搓CUDA內核

機器之心編輯部

近日，來自字節跳動 Seed 團隊和清華大學 AIR的新研究CUDA Agent，在 AI 領域引發了不小的轟動。

研究人員訓練了一個能夠編寫快速 CUDA 內核的模型：不只是正確的內核，而是真正經過優化的內核。

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

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

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

GPU 內核優化是現代深度學習的基礎，但它仍然是一項高度專業化的工作，需要深厚的硬件知識。現有的 AI 生成 CUDA 代碼通常依賴無訓練的提示工程（Prompting）或多輪執行反饋微調的機制。這導致模型只能解決表面上的語法錯誤，無法真正理解底層硬件邏輯，一定程度上限制了其內在的優化能力。

真正極致的 CUDA 優化需要處理的任務，是只有在性能分析器中才能看到的硬件級指標。人們一直期待能出現一個像人類 CUDA 專家一樣思考的 AI。

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

研究人員不再獎勵「是否編譯成功」，而是獎勵實際的GPU速度。真實的性能分析數據。強化學習直接基于性能進行訓練。

產生的效果出乎人們的預料。

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

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

簡而言之，CUDA Agent是一個大規模的智能體強化學習系統，包含三個核心組成部分：可擴展的數據合成機制、一個集成技能增強且具備可靠驗證與性能分析能力的 CUDA 開發環境，以及用于穩定長上下文訓練的強化學習算法技術。

此外，研究團隊同時發布了CUDA-Agent-Ops-6K，一個經過嚴格篩選與數據污染控制的高質量合成訓練數據集，可支持基于強化學習的 CUDA 內核優化研究的復現。

系統管線設計

數據合成

研究團隊通過一個三階段的管線來構建訓練任務：種子問題爬取、基于 LLM 的組合式合成，以及基于執行結果的篩選。

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

三階段數據收集管線

最終整理得到 6000 條訓練樣本，構建了 CUDA-Agent-Ops-6K 數據集，該數據集專為可擴展的強化學習訓練而設計，兼具廣泛的任務多樣性和較低的數據污染風險。

智能體環境

智能體循環管線遵循一種 ReAct 風格的工作流，結合代碼工具與 CUDA Skill 規范（SKILL.md），支持迭代式的編碼-編譯-調試循環，以及基于性能分析器的優化過程。

• 標準工作流程：對原生 PyTorch 實現進行性能分析，編寫 CUDA 內核及其綁定代碼，在 GPU 沙盒環境中完成編譯，并不斷迭代優化。
• 目標要求：通過正確性檢查，并在性能上相對于 torch.compile 實現超過 5% 的加速。
• 穩健的獎勵機制采用基于里程碑的離散獎勵設計，根據正確性達標情況和性能提升幅度分別給予獎勵。
• 防止獎勵作弊的控制措施包括：對驗證與性能分析腳本進行保護，禁止回退調用，采用 5 組不同輸入進行正確性檢查，在同步預熱后進行性能分析，以及禁止網絡檢索。

這些約束共同構建了可靠的、基于真實執行結果的反饋機制，使策略學習聚焦于內核質量的實質性提升，而非依賴取巧或捷徑行為。

訓練流程

訓練過程采用分階段設計，以穩定 CUDA 編碼這一長時序強化學習任務。首先進行單輪 PPO 預熱訓練，隨后分別初始化 actor 和 critic，最后進入完整的多輪智能體強化學習階段。

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

借助這一多階段訓練設計，系統在長上下文設定下（最長 128k 上下文、訓練階段最多 150 輪、評估階段最多 200 輪）依然保持穩定，從而實現持續的獎勵增長。

核心實驗結果

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

與強大的專有基線模型相比，CUDA Agent 在相對于 Compile 的性能優化上展現出顯著優勢：在整體 KernelBench 基準測試中，其相對 Compile 的加速達成率達到 96.8%，幾何平均加速比為 2.11 倍

這一優勢在高難度設置下尤為明顯：在 Level-3 上，CUDA Agent 相對 compile 的加速達成率達到 90%，相比最強的專有基線高出約 40 個百分點；在 Level-2 的算子序列任務上，其加速達成率達到 100%，幾何平均加速比達到 2.80 倍。

在 KernelBench 上的整體性能和加速指標。

本研究存在兩個主要局限。

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

看起來，CUDA Agent 等技術的出現即將打破傳統編譯器（如torch.compile或Triton）的優化瓶頸。它證明了：大語言模型不僅可以學習人類自然語言和高級編程語言，還可以通過基于硬件反饋的強化學習，內化出極高門檻的「硬件直覺」。

一條通向全自動、高度性能優化計算基礎設施的道路正在出現。