![]()
本文作者包括明尼蘇達大學的李世陽(共同第一作者),張子健(共同第一作者),Winson Chen,羅越波,洪明毅,丁才文。
現有的 LLM 自動化 CUDA 方法大多只能優化單個 Kernel,面對完整的端到端 GPU 程序(如整個 VisionTransformer 推理)往往束手無策。
本文中,StitchCUDA 提出了一個根本性的問題轉向:從優化單個 Kernel,到生成完整的端到端 GPU 程序。通過多智能體協作框架與基于 Rubric Reward 的 Agentic RL,StitchCUDA 在 KernelBench Level 3 端到端任務上實現了90% 的成功率和 1.50× 的平均加速比,分別比多智能體基線高出 1.72× 和 RL 模型基線高出 2.73×。
![]()
- 論文標題:StitchCUDA: An Automated Multi-Agents End-to-End GPU Programming Framework with Rubric-based Agentic Reinforcement Learning
- 論文鏈接:http://arxiv.org/abs/2603.02637
背景與動機:從單 Kernel 優化到端到端 GPU 編程
CUDA 代碼的性能對當今模型訓練與推理至關重要。近年來,基于 LLM 的 CUDA 代碼生成取得了不少進展:多智能體框架(如 CUDAForge、QiMeng)和基于 RL 的方法(如 Kevin-32B、CUDA-L1/L2)在 KernelBench Level 1/2 的單 Kernel 任務上均表現出色。
然而,真正的挑戰在于端到端 GPU 程序的生成。KernelBench Level 3 的任務涉及完整的模型架構(如 MiniGPTBlock 推理代碼),影響性能的因素不僅僅包括單個 kernel runtime,還由算子融合、Launch 配置、CPU-GPU 同步、數據搬運等系統級因素共同決定。如下圖所示,現有方法在 KernelBench Level 3 上的表現遠不理想:
![]()
- GPT-5.2(前沿 LLM):成功率僅 20%,加速比 0.48×
- Kevin-32B(RL 模型):成功率 20%,加速比 0.34×
- CUDAForge(多智能體):成功率 60%,加速比 0.87×
- StitchCUDA(本文方法):成功率 90%,加速比 1.50×
研究團隊總結了使用 LLM 進行端到端 CUDA 生成與優化的三大核心挑戰:
(C1)端到端程序需要全局協調。不同于單 Kernel 優化,端到端 GPU 程序的性能由 Kernel 融合邊界、跨 Kernel 內存布局、CPU-GPU 同步等系統級決策主導,無法通過逐一處理單個 Kernel 來解決。
(C2)Coder 的 CUDA 編程能力需要在 Prompt 工程以外進一步提升。多智能體框架可以從其他 Agent 獲取反饋來引導 Coder,但如果沒有參數更新,Coder 往往無法可靠地執行復雜的 CUDA 變換(例如根據性能分析提示推導出正確的 Tiling 策略),成為實際中的主要瓶頸。
(C3)現有的 RL 方法存在諸多挑戰。現有的 RLVR 方法容易出現 Reward Hacking(如直接抄寫 PyTorch 代碼或硬編碼輸出)和退化行為(只替換簡單的 ReLU 而不碰關鍵的 Conv/GEMM);同時,Coder 也沒有被訓練去理解結構化的執行反饋并實施有針對性的優化,進而使得訓練的模型不適配多智能體框架。
StitchCUDA 方法介紹
為了解決上述挑戰,StitchCUDA 提出了一套多智能體框架 + 基于 Rubric Reward 的 Agentic RL方案,核心包含兩大模塊:
多智能體協作框架
![]()
StitchCUDA 將端到端 GPU 編程任務分解為三個專門的 Agent,通過迭代式「計劃 — 編碼 — 分析 — 優化」循環協作完成:
Planner(規劃器):解析 PyTorch 參考代碼,使用 Nsys Profiling 進行性能分析,識別耗時最長的 Kernel 和系統瓶頸。然后在系統層面將任務分解為多個子任務,同時考慮 Kernel 效率和 Host 端編排 —— 例如:子任務 1「使用 Pinned Memory 分配連續張量」,子任務 2「用 cuBLASLt 融合 Epilogue 替換 FC 層 GEMM」,子任務 3「用定制的 fast in-place ReLU kernel 替換 pointwise Kernel」。
Coder(編碼器):按照 Planner 的規劃,逐個子任務生成 CUDA 實現(源代碼、構建文件、Pybind 接口),并調用 nvcc 編譯。在收到 Verifier 的反饋后,對當前子任務進行迭代優化。
Verifier(驗證器):負責正確性驗證和性能分析。編譯失敗時,分析錯誤日志并返回具體修復指導。測試通過時,從兩個層面分析程序:Nsys用于識別最耗時的 GPU Kernel 和系統級瓶頸(如 CPU-GPU 數據傳輸、Kernel Launch、同步開銷),NCU用于分析具體的瓶頸 Kernel(判斷是 Memory-bound 還是 Compute-bound),最終生成可執行的優化建議。
此外,Planner 和 Verifier 還集成了RAG 模塊,從 NVIDIA 官方文檔(CUDA-12.9、cuBLAS、CUTLASS、Nsys/NCU 指南、Hopper/Blackwell 架構白皮書)中檢索最新的 API 規范和用法指南,避免 LLM 因預訓練知識過時而產生幻覺。
基于 Rubric Reward 的 Agentic RL
為了提升 Coder 的端到端 GPU 編程能力,StitchCUDA 引入了一種創新的 Agentic RL 訓練方案:
![]()
將多輪交互分解為原子技能
標準的多輪 Agentic RL 需要收集完整的交互軌跡(15 輪迭代 × 每輪 4-5 分鐘環境交互),單條軌跡就需要 60-75 分鐘,整體訓練預估需要約8 卡 H200 訓練 1200-1500 小時。StitchCUDA 將其分解為兩個原子技能的單輪 RL 訓練:
- Skill 1(從零生成):給定參考 PyTorch 代碼和子任務需求,生成正確的 CUDA 實現
- Skill 2(反饋驅動優化):根據結構化的執行反饋(編譯診斷、性能瓶頸分析),修復 Bug 并提升性能
通過在工作流執行過程中收集單輪訓練數據(每個 Skill 各 200 個樣本),然后合并用 GRPO 聯合優化, 訓練一個基于 Qwen-32B 的 Coder 僅需約160 H200-Hour,相比多輪 Agentic RL 減少了約60-75 倍。
Rubric Reward:解決 Reward Hacking 和退化行為
現有 RL 方法的核心問題在于獎勵設計:簡單的「正確性 + 加速比」獎勵容易被 LLM 利用,比如說,直接復制 PyTorch 代碼就能獲得高獎勵,而激進優化若產生微小錯誤則獎勵為零,這推動模型走向保守、退化的行為。此外,模型也會通過直接復制 pytorch 代碼的形式來 hacking 評測程序,從而獲得高 reward。
StitchCUDA 引入了由 CUDA 專家設計的Rubric Reward(評分準則獎勵),從四個維度對生成代碼進行綜合評估:
- Anti-Hacking(反作弊):懲罰 Reward Hacking 行為(如復制 PyTorch 代碼、硬編碼輸出)
- CUDA Engineering(工程質量):獎勵高級優化技術的使用(Tiling + Shared Memory、cuBLASLt Epilogue、Tensor Core、混合精度等)
- Operator Coverage(算子覆蓋):鼓勵覆蓋更多關鍵算子的優化,而非只替換簡單的 ReLU
- Skill Compliance(技能遵循):確保遵循任務需求(Skill 1)或反饋指令(Skill 2)
最終獎勵公式將 Rubric Reward 與規則獎勵(正確性 × 加速比)相結合,同時通過 Reward Clipping(R_max=5)防止極端獎勵對訓練的干擾,增強訓練的穩定性。
實驗評估
實驗在 KernelBench Level 1/2/3 上進行,測試硬件覆蓋兩代 NVIDIA 架構:H200(Hopper)和RTX PRO 6000(Blackwell),以驗證方法的跨架構泛化能力。對比方法包括前沿 LLM(GPT-5.2、Claude-4-sonnet)、開源基座(Qwen3-32B)、RL 方法(Kevin-32B)和多智能體框架(CUDAForge),以及 StitchCUDA 的多個變體。
下表展示了所有方法在兩個硬件平臺上的完整結果(正確率 / 平均加速比 / Fast1):
![]()
關鍵發現
多智能體框架大幅提升端到端正確性。以 Qwen3-32B 為例,單次生成在 Level 3 上失敗(0/10),而 StitchCUDA 多智能體框架(不含 RL)將其提升到 3/10。即使是更強的 GPT-5.2,多智能體框架也帶來顯著提升。
Agentic RL 是實現系統級加速的關鍵。對比 StitchCUDA 和無 RL 變體(StitchCUDA-Q),RL 在 Level 3 上將正確率從 3/10 提升至 9/10,加速比從 0.24× 提升至 1.50×,Fast1 從 10% 提升至 70%。
Agentic RL 超越更強模型的效果。即使與使用 GPT-5.2 作為所有 Agent 的 StitchCUDA-G 相比,使用 Qwen-32B 作為 Coder 的 StitchCUDA 在 Level 3 上仍然全面領先。 這說明 RL 訓練帶來的能力提升是模型規模難以替代的。
超越 torch.compile。在 H200 上,StitchCUDA 對比啟用 torch.compile 的參考代碼仍然實現了 1.29× 的加速,表明其手動的系統級優化(自定義 Kernel 融合、數據搬運優化)能夠超越編譯器的自動優化。
Hacking 檢測
Reward Hacking 是 CUDA RL 訓練中的重要挑戰之一。模型會因此學會 hack 測評程序而不是進行 CUDA 優化,我們對 50 個測試任務進行了系統性的 hacking 檢測,結果如下:
![]()
格式檢查(Format Check)為什么不夠?一種直覺的解決方案是用規則檢測 reward hacking, 比如檢查生成代碼中是否包含 torch.nn.functional 調用。但我們發現這存在一個 trade-off:
- 檢查過嚴 → 誤殺合法實現。在端到端 GPU 程序中,合理地在 CUDA Kernel 內部調用 PyTorch 子函數是完全合法的策略(例如用 cuDNN 的 torch.conv2d 處理卷積,同時自定義融合后續的 Bias+ReLU)。過嚴的格式檢查會將這類正確且高效的實現判定為 Hacking。
- 檢查過松 → 漏過作弊。放寬檢查標準又會讓模型輕松繞過,比如將 PyTorch 調用封裝在一層 wrapper 函數中。
Rubric Reward 如何解決 Reward Hacking 的問題?StitchCUDA 的 Rubric Reward 不依賴硬編碼的格式規則,而是使用推理模型按照專家設計的評分準則從語義層面評估代碼質量。Anti-Hacking 維度會判斷「生成代碼是否真正實現了 CUDA 優化,還是本質上仍在調用 PyTorch」,這種語義級評估天然地避免了格式檢查的 false positive/negative 困境。
結果是顯著的:StitchCUDA 將 Hacking 率從 Kevin-32B 的 52% 降至 16%, Hacking 從 4 次降至 0 次。而去除 Rubric 的 StitchCUDA-A 變體,Hacking 率回升至 32%,進一步驗證了 Rubric Reward 的因果效應。
消融實驗
去除 Rubric Reward 后(StitchCUDA-A),Level 3 成功率從 90% 降至 50%,加速比從 1.50× 降至 0.46×,進一步確認了 Rubric Reward 對有效 RL 訓練的關鍵作用。失敗原因包括:退化的保守實現、反饋遵循能力下降、以及無 Rubric 懲罰導致的 Reward Hacking。
案例展示
以 Level 3 Task 44(GPT-2 Transformer Block)為例,StitchCUDA 實現了3.75× 加速比。
Planner 在系統層面提出了混合精度計算(LayerNorm / 殘差用 fp32,MLP 用 fp16)和連續數據布局優化。在 Kernel 層面,Coder 實現了cuBLASLt Epilogue 融合(將GEMM+Bias 和 GEMM+Bias+GELU 融合為單次 Launch)、描述符 / 算法緩存(避免重復的 Heuristic 查詢)、以及按 Stream 持久化 Workspace(減少 cudaMalloc/cudaFree 開銷)。
這些系統級 + Kernel 級協同優化是單 Kernel 優化方法無法實現的。
總結
StitchCUDA 提出了首個面向端到端 GPU 程序生成的完整解決方案,通過:
- 多智能體協作框架:將復雜的端到端任務分解為「計劃 — 編碼 — 分析 — 優化」的迭代循環
- 原子技能分解:將昂貴的多輪 Agentic RL 轉化為高效的單輪訓練,降低約 60-75 倍計算開銷
- Rubric Reward:從反作弊、工程質量、算子覆蓋、技能遵循四維度全面評估,有效解決 Reward Hacking 和退化行為,鼓勵模型優化更多的算子,使用更高級的技術。
在 KernelBench 上,StitchCUDA 在端到端任務上實現了近 100% 的成功率和1.5× 的平均加速比,顯著超越所有現有方法,為 LLM 驅動的自動化 GPU 編程開辟了新的方向。
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.