本文作者包括明尼蘇達(dá)大學(xué)的李世陽（共同第一作者），張子?。ü餐谝蛔髡撸琖inson Chen，羅越波，洪明毅，丁才文。

現(xiàn)有的 LLM 自動化 CUDA 方法大多只能優(yōu)化單個 Kernel，面對完整的端到端 GPU 程序（如整個 VisionTransformer 推理）往往束手無策。

本文中，StitchCUDA 提出了一個根本性的問題轉(zhuǎn)向：從優(yōu)化單個 Kernel，到生成完整的端到端 GPU 程序。通過多智能體協(xié)作框架與基于 Rubric Reward 的 Agentic RL，StitchCUDA 在 KernelBench Level 3 端到端任務(wù)上實(shí)現(xiàn)了90% 的成功率和 1.50× 的平均加速比，分別比多智能體基線高出 1.72× 和 RL 模型基線高出 2.73×。

• 論文標(biāo)題：StitchCUDA: An Automated Multi-Agents End-to-End GPU Programming Framework with Rubric-based Agentic Reinforcement Learning
• 論文鏈接：http://arxiv.org/abs/2603.02637

背景與動機(jī)：從單 Kernel 優(yōu)化到端到端 GPU 編程

CUDA 代碼的性能對當(dāng)今模型訓(xùn)練與推理至關(guān)重要。近年來，基于 LLM 的 CUDA 代碼生成取得了不少進(jìn)展：多智能體框架（如 CUDAForge、QiMeng）和基于 RL 的方法（如 Kevin-32B、CUDA-L1/L2）在 KernelBench Level 1/2 的單 Kernel 任務(wù)上均表現(xiàn)出色。

然而，真正的挑戰(zhàn)在于端到端 GPU 程序的生成。KernelBench Level 3 的任務(wù)涉及完整的模型架構(gòu)（如 MiniGPTBlock 推理代碼），影響性能的因素不僅僅包括單個 kernel runtime，還由算子融合、Launch 配置、CPU-GPU 同步、數(shù)據(jù)搬運(yùn)等系統(tǒng)級因素共同決定。如下圖所示，現(xiàn)有方法在 KernelBench Level 3 上的表現(xiàn)遠(yuǎn)不理想：

• GPT-5.2（前沿 LLM）：成功率僅 20%，加速比 0.48×
• Kevin-32B（RL 模型）：成功率 20%，加速比 0.34×
• CUDAForge（多智能體）：成功率 60%，加速比 0.87×
• StitchCUDA（本文方法）：成功率 90%，加速比 1.50×

研究團(tuán)隊(duì)總結(jié)了使用 LLM 進(jìn)行端到端 CUDA 生成與優(yōu)化的三大核心挑戰(zhàn)：

（C1）端到端程序需要全局協(xié)調(diào)。不同于單 Kernel 優(yōu)化，端到端 GPU 程序的性能由 Kernel 融合邊界、跨 Kernel 內(nèi)存布局、CPU-GPU 同步等系統(tǒng)級決策主導(dǎo)，無法通過逐一處理單個 Kernel 來解決。

（C2）Coder 的 CUDA 編程能力需要在 Prompt 工程以外進(jìn)一步提升。多智能體框架可以從其他 Agent 獲取反饋來引導(dǎo) Coder，但如果沒有參數(shù)更新，Coder 往往無法可靠地執(zhí)行復(fù)雜的 CUDA 變換（例如根據(jù)性能分析提示推導(dǎo)出正確的 Tiling 策略），成為實(shí)際中的主要瓶頸。

（C3）現(xiàn)有的 RL 方法存在諸多挑戰(zhàn)?，F(xiàn)有的 RLVR 方法容易出現(xiàn) Reward Hacking（如直接抄寫 PyTorch 代碼或硬編碼輸出）和退化行為（只替換簡單的 ReLU 而不碰關(guān)鍵的 Conv/GEMM）；同時，Coder 也沒有被訓(xùn)練去理解結(jié)構(gòu)化的執(zhí)行反饋并實(shí)施有針對性的優(yōu)化，進(jìn)而使得訓(xùn)練的模型不適配多智能體框架。

StitchCUDA 方法介紹

為了解決上述挑戰(zhàn)，StitchCUDA 提出了一套多智能體框架 + 基于 Rubric Reward 的 Agentic RL方案，核心包含兩大模塊：

多智能體協(xié)作框架

StitchCUDA 將端到端 GPU 編程任務(wù)分解為三個專門的 Agent，通過迭代式「計劃 — 編碼 — 分析 — 優(yōu)化」循環(huán)協(xié)作完成：

Planner（規(guī)劃器）：解析 PyTorch 參考代碼，使用 Nsys Profiling 進(jìn)行性能分析，識別耗時最長的 Kernel 和系統(tǒng)瓶頸。然后在系統(tǒng)層面將任務(wù)分解為多個子任務(wù)，同時考慮 Kernel 效率和 Host 端編排 —— 例如：子任務(wù) 1「使用 Pinned Memory 分配連續(xù)張量」，子任務(wù) 2「用 cuBLASLt 融合 Epilogue 替換 FC 層 GEMM」，子任務(wù) 3「用定制的 fast in-place ReLU kernel 替換 pointwise Kernel」。

Coder（編碼器）：按照 Planner 的規(guī)劃，逐個子任務(wù)生成 CUDA 實(shí)現(xiàn)（源代碼、構(gòu)建文件、Pybind 接口），并調(diào)用 nvcc 編譯。在收到 Verifier 的反饋后，對當(dāng)前子任務(wù)進(jìn)行迭代優(yōu)化。

Verifier（驗(yàn)證器）：負(fù)責(zé)正確性驗(yàn)證和性能分析。編譯失敗時，分析錯誤日志并返回具體修復(fù)指導(dǎo)。測試通過時，從兩個層面分析程序：Nsys用于識別最耗時的 GPU Kernel 和系統(tǒng)級瓶頸（如 CPU-GPU 數(shù)據(jù)傳輸、Kernel Launch、同步開銷），NCU用于分析具體的瓶頸 Kernel（判斷是 Memory-bound 還是 Compute-bound），最終生成可執(zhí)行的優(yōu)化建議。

此外，Planner 和 Verifier 還集成了RAG 模塊，從 NVIDIA 官方文檔（CUDA-12.9、cuBLAS、CUTLASS、Nsys/NCU 指南、Hopper/Blackwell 架構(gòu)白皮書）中檢索最新的 API 規(guī)范和用法指南，避免 LLM 因預(yù)訓(xùn)練知識過時而產(chǎn)生幻覺。

基于 Rubric Reward 的 Agentic RL

為了提升 Coder 的端到端 GPU 編程能力，StitchCUDA 引入了一種創(chuàng)新的 Agentic RL 訓(xùn)練方案：

將多輪交互分解為原子技能

標(biāo)準(zhǔn)的多輪 Agentic RL 需要收集完整的交互軌跡（15 輪迭代 × 每輪 4-5 分鐘環(huán)境交互），單條軌跡就需要 60-75 分鐘，整體訓(xùn)練預(yù)估需要約8 卡 H200 訓(xùn)練 1200-1500 小時。StitchCUDA 將其分解為兩個原子技能的單輪 RL 訓(xùn)練：

• Skill 1（從零生成）：給定參考 PyTorch 代碼和子任務(wù)需求，生成正確的 CUDA 實(shí)現(xiàn)
• Skill 2（反饋驅(qū)動優(yōu)化）：根據(jù)結(jié)構(gòu)化的執(zhí)行反饋（編譯診斷、性能瓶頸分析），修復(fù) Bug 并提升性能

通過在工作流執(zhí)行過程中收集單輪訓(xùn)練數(shù)據(jù)（每個 Skill 各 200 個樣本），然后合并用 GRPO 聯(lián)合優(yōu)化， 訓(xùn)練一個基于 Qwen-32B 的 Coder 僅需約160 H200-Hour，相比多輪 Agentic RL 減少了約60-75 倍。

Rubric Reward：解決 Reward Hacking 和退化行為

現(xiàn)有 RL 方法的核心問題在于獎勵設(shè)計：簡單的「正確性 + 加速比」獎勵容易被 LLM 利用，比如說，直接復(fù)制 PyTorch 代碼就能獲得高獎勵，而激進(jìn)優(yōu)化若產(chǎn)生微小錯誤則獎勵為零，這推動模型走向保守、退化的行為。此外，模型也會通過直接復(fù)制 pytorch 代碼的形式來 hacking 評測程序，從而獲得高 reward。

StitchCUDA 引入了由 CUDA 專家設(shè)計的Rubric Reward（評分準(zhǔn)則獎勵），從四個維度對生成代碼進(jìn)行綜合評估：

1. Anti-Hacking（反作弊）：懲罰 Reward Hacking 行為（如復(fù)制 PyTorch 代碼、硬編碼輸出）
2. CUDA Engineering（工程質(zhì)量）：獎勵高級優(yōu)化技術(shù)的使用（Tiling + Shared Memory、cuBLASLt Epilogue、Tensor Core、混合精度等）
3. Operator Coverage（算子覆蓋）：鼓勵覆蓋更多關(guān)鍵算子的優(yōu)化，而非只替換簡單的 ReLU
4. Skill Compliance（技能遵循）：確保遵循任務(wù)需求（Skill 1）或反饋指令（Skill 2）

最終獎勵公式將 Rubric Reward 與規(guī)則獎勵（正確性 × 加速比）相結(jié)合，同時通過 Reward Clipping（R_max=5）防止極端獎勵對訓(xùn)練的干擾，增強(qiáng)訓(xùn)練的穩(wěn)定性。

實(shí)驗(yàn)評估

實(shí)驗(yàn)在 KernelBench Level 1/2/3 上進(jìn)行，測試硬件覆蓋兩代 NVIDIA 架構(gòu)：H200（Hopper）和RTX PRO 6000（Blackwell），以驗(yàn)證方法的跨架構(gòu)泛化能力。對比方法包括前沿 LLM（GPT-5.2、Claude-4-sonnet）、開源基座（Qwen3-32B）、RL 方法（Kevin-32B）和多智能體框架（CUDAForge），以及 StitchCUDA 的多個變體。

下表展示了所有方法在兩個硬件平臺上的完整結(jié)果（正確率 / 平均加速比 / Fast1）：

關(guān)鍵發(fā)現(xiàn)

多智能體框架大幅提升端到端正確性。以 Qwen3-32B 為例，單次生成在 Level 3 上失?。?/10），而 StitchCUDA 多智能體框架（不含 RL）將其提升到 3/10。即使是更強(qiáng)的 GPT-5.2，多智能體框架也帶來顯著提升。

Agentic RL 是實(shí)現(xiàn)系統(tǒng)級加速的關(guān)鍵。對比 StitchCUDA 和無 RL 變體（StitchCUDA-Q），RL 在 Level 3 上將正確率從 3/10 提升至 9/10，加速比從 0.24× 提升至 1.50×，F(xiàn)ast1 從 10% 提升至 70%。

Agentic RL 超越更強(qiáng)模型的效果。即使與使用 GPT-5.2 作為所有 Agent 的 StitchCUDA-G 相比，使用 Qwen-32B 作為 Coder 的 StitchCUDA 在 Level 3 上仍然全面領(lǐng)先。 這說明 RL 訓(xùn)練帶來的能力提升是模型規(guī)模難以替代的。

超越 torch.compile。在 H200 上，StitchCUDA 對比啟用 torch.compile 的參考代碼仍然實(shí)現(xiàn)了 1.29× 的加速，表明其手動的系統(tǒng)級優(yōu)化（自定義 Kernel 融合、數(shù)據(jù)搬運(yùn)優(yōu)化）能夠超越編譯器的自動優(yōu)化。

Hacking 檢測

Reward Hacking 是 CUDA RL 訓(xùn)練中的重要挑戰(zhàn)之一。模型會因此學(xué)會 hack 測評程序而不是進(jìn)行 CUDA 優(yōu)化，我們對 50 個測試任務(wù)進(jìn)行了系統(tǒng)性的 hacking 檢測，結(jié)果如下：

格式檢查（Format Check）為什么不夠？一種直覺的解決方案是用規(guī)則檢測 reward hacking， 比如檢查生成代碼中是否包含 torch.nn.functional 調(diào)用。但我們發(fā)現(xiàn)這存在一個 trade-off：

• 檢查過嚴(yán) → 誤殺合法實(shí)現(xiàn)。在端到端 GPU 程序中，合理地在 CUDA Kernel 內(nèi)部調(diào)用 PyTorch 子函數(shù)是完全合法的策略（例如用 cuDNN 的 torch.conv2d 處理卷積，同時自定義融合后續(xù)的 Bias+ReLU）。過嚴(yán)的格式檢查會將這類正確且高效的實(shí)現(xiàn)判定為 Hacking。
• 檢查過松 → 漏過作弊。放寬檢查標(biāo)準(zhǔn)又會讓模型輕松繞過，比如將 PyTorch 調(diào)用封裝在一層 wrapper 函數(shù)中。

Rubric Reward 如何解決 Reward Hacking 的問題？StitchCUDA 的 Rubric Reward 不依賴硬編碼的格式規(guī)則，而是使用推理模型按照專家設(shè)計的評分準(zhǔn)則從語義層面評估代碼質(zhì)量。Anti-Hacking 維度會判斷「生成代碼是否真正實(shí)現(xiàn)了 CUDA 優(yōu)化，還是本質(zhì)上仍在調(diào)用 PyTorch」，這種語義級評估天然地避免了格式檢查的 false positive/negative 困境。

結(jié)果是顯著的：StitchCUDA 將 Hacking 率從 Kevin-32B 的 52% 降至 16%， Hacking 從 4 次降至 0 次。而去除 Rubric 的 StitchCUDA-A 變體，Hacking 率回升至 32%，進(jìn)一步驗(yàn)證了 Rubric Reward 的因果效應(yīng)。

消融實(shí)驗(yàn)

去除 Rubric Reward 后（StitchCUDA-A），Level 3 成功率從 90% 降至 50%，加速比從 1.50× 降至 0.46×，進(jìn)一步確認(rèn)了 Rubric Reward 對有效 RL 訓(xùn)練的關(guān)鍵作用。失敗原因包括：退化的保守實(shí)現(xiàn)、反饋?zhàn)裱芰ο陆?、以及無 Rubric 懲罰導(dǎo)致的 Reward Hacking。

案例展示

以 Level 3 Task 44（GPT-2 Transformer Block）為例，StitchCUDA 實(shí)現(xiàn)了3.75× 加速比。

Planner 在系統(tǒng)層面提出了混合精度計算（LayerNorm / 殘差用 fp32，MLP 用 fp16）和連續(xù)數(shù)據(jù)布局優(yōu)化。在 Kernel 層面，Coder 實(shí)現(xiàn)了cuBLASLt Epilogue 融合（將GEMM+Bias 和 GEMM+Bias+GELU 融合為單次 Launch）、描述符 / 算法緩存（避免重復(fù)的 Heuristic 查詢）、以及按 Stream 持久化 Workspace（減少 cudaMalloc/cudaFree 開銷）。

這些系統(tǒng)級 + Kernel 級協(xié)同優(yōu)化是單 Kernel 優(yōu)化方法無法實(shí)現(xiàn)的。

總結(jié)

StitchCUDA 提出了首個面向端到端 GPU 程序生成的完整解決方案，通過：

• 多智能體協(xié)作框架：將復(fù)雜的端到端任務(wù)分解為「計劃 — 編碼 — 分析 — 優(yōu)化」的迭代循環(huán)
• 原子技能分解：將昂貴的多輪 Agentic RL 轉(zhuǎn)化為高效的單輪訓(xùn)練，降低約 60-75 倍計算開銷
• Rubric Reward：從反作弊、工程質(zhì)量、算子覆蓋、技能遵循四維度全面評估，有效解決 Reward Hacking 和退化行為，鼓勵模型優(yōu)化更多的算子，使用更高級的技術(shù)。

在 KernelBench 上，StitchCUDA 在端到端任務(wù)上實(shí)現(xiàn)了近 100% 的成功率和1.5× 的平均加速比，顯著超越所有現(xiàn)有方法，為 LLM 驅(qū)動的自動化 GPU 編程開辟了新的方向。