中科院團隊提出SparseRL，深度強化學習可自動生成高性能CUDA代碼

如何讓 AI 不僅寫出「能跑」的代碼，還能寫出「跑得快」的代碼？這個問題困擾了 AI 系統研究者很久。

近日，中科院計算所團隊提出了一種名為 SparseRL 的新框架，首次將深度強化學習引入稀疏 CUDA 代碼生成任務。簡單來說，就是讓 AI 學會根據稀疏矩陣的結構，自動生成最優的 CUDA 實現代碼。

實驗顯示，在經典的 SpMV 任務上，這種方法能讓編譯成功率提升 20%，代碼執行速度提升 30%。

目前，該項成果已入選 ICLR 2026 Oral。

• 論文地址：https://openreview.net/pdf?id=VdLEaGPYWT
• 代碼鏈接：https://github.com/QiWu-NCIC/SparseRL

為什么稀疏代碼這么難寫？

要理解這項工作的價值，得先說說稀疏矩陣運算的特殊性。

稀疏矩陣在 LLM 推理、圖神經網絡、科學計算中無處不在。但和稠密矩陣不同，稀疏矩陣的非零元素分布是不規則的，這導致一個核心問題：最優的 CUDA 代碼實現取決于矩陣的具體結構，而這個結構只有在運行時才能知道。

換句話說，沒有一種「萬能」的高性能實現能應對所有稀疏矩陣。工程師們不得不針對不同的稀疏模式手動調優，這個過程既耗時又依賴經驗。

現有的 AI 代碼生成方法也幫不上太大忙。原因有三：

• 第一，傳統監督學習只關心代碼「對不對」，不關心「快不快」。同一個稀疏矩陣可能有多種正確的 CUDA 實現，但執行速度可能相差數倍，監督學習無法區分這種差異。
• 第二，執行效率這個核心指標是「不可微」的，沒法通過傳統的反向傳播來優化。
• 第三，稀疏矩陣的輸入（行列索引序列）和 CUDA 代碼之間存在巨大的語義鴻溝，模型很難理解矩陣結構和最優代碼策略之間的關聯。

圖片 1：展示不同稀疏矩陣需要不同 CUDA 實現策略的示例

SparseRL 怎么做到的？

研究團隊的思路很巧妙：既然執行效率不可微，那就用強化學習來優化。

SparseRL 把預訓練語言模型當作一個策略網絡，每生成一個 token 就是一次動作，而代碼的編譯結果和執行時間就是獎勵信號。

整個訓練過程分為三個階段：

• 第一階段是預訓練：在大量 CUDA 代碼語料上訓練語言模型，讓它建立對 GPU 編程的基礎認知；
• 第二階段是監督微調：用「稀疏矩陣 - 正確代碼」的配對數據教模型生成語法正確、功能正確的代碼；
• 第三階段是強化學習優化：這一步是關鍵 —— 引入深度強化學習，以編譯正確性和執行效率為獎勵，讓模型學會生成高性能代碼。

圖片 2：展示三階段訓練流程的整體框架圖

為了讓模型真正「看懂」稀疏矩陣的結構，研究團隊設計了一個關鍵技術：正弦位置嵌入。

稀疏矩陣的輸入是非零元素的行列索引序列，傳統的 token 嵌入無法捕捉這種二維坐標之間的空間關系。SparseRL 對行列索引分別進行正弦 / 余弦編碼，類似于 Transformer 的位置編碼，但專門針對二維坐標做了定制。

用通俗的話說，這就像給模型裝上了一副「坐標眼鏡」，讓它能看見非零元素在哪里、是怎么分布的。

另一個核心創新是層級獎勵函數。這個獎勵函數同時考慮兩個層面：正確性獎勵確保代碼能編譯、結果正確；效率獎勵則優化執行速度。設計邏輯是先保證「對」，再追求「快」。

效果如何？

研究團隊在 SpMV（稀疏矩陣 - 向量乘法）和 SpMM（稀疏矩陣 - 稠密矩陣乘法）兩個任務上驗證了方法的有效性。

在 SpMV 任務上，SparseRL 相比傳統監督學習方法，編譯成功率提升了 20%，平均執行速度提升了 30%。更重要的是，模型能根據不同的稀疏結構自動選擇不同的代碼策略，在對角型、帶狀型、隨機稀疏型等多種矩陣上都有優勢，部分場景下生成的代碼甚至接近或超越了手工調優的水平。

圖片 3：展示 SparseRL 與基線方法的差距

團隊還做了消融實驗來驗證各個組件的必要性。

結果顯示，去掉 RL 階段后性能顯著下降，說明強化學習確實是關鍵；去掉正弦嵌入后模型難以理解輸入結構，編譯率下降；只用正確性獎勵而不用效率獎勵，代碼能跑但不夠快。

當然，這個方法也有局限。論文提到，RL 訓練需要大量的編譯 - 執行反饋循環，計算成本較高；模型是針對特定 GPU 架構訓練的，遷移到新硬件可能需要重新微調；生成的代碼可能缺乏人類工程師的編碼風格，可解釋性不足。

意義與展望

SparseRL 的價值在于它代表了一個范式轉變：代碼生成的目標從「生成能運行的代碼」轉向「生成高性能代碼」。

對于 HPC 工程師和 AI 基礎設施開發者來說，這項工作展示了一種新可能 ——讓 AI 來處理那些繁瑣的性能優化工作，而人類可以把精力放在更高層次的設計上。

研究團隊表示，未來計劃將方法擴展到多 GPU 分布式稀疏計算，探索與傳統 AutoTuning 技術的結合，并支持更多類型的稀疏算子。同時，他們也在研究如何降低 RL 訓練成本，讓這種方法更實用。

作者介紹

王耀宇，中國科學院計算技術研究所博士生（共同一作），主要研究方向為深度學習編譯優化與高性能計算。

譚光明，中國科學院計算技術研究所研究員、博士生導師，主要從事高性能計算、GPU 編譯優化與深度學習系統研究，在多 GPU 分布式計算、稀疏矩陣計算、深度學習編譯器等領域取得多項重要成果，發表多篇高性能計算與機器學習相關論文。