1. 首頁 > 生活百科 > >

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

人工智能aigpucuda

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

文章圖片

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

文章圖片

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



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

【中科院團隊提出SparseRL，深度強化學習可自動生成高性能CUDA代碼】近日 ， 中科院計算所團隊提出了一種名為 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 分布式計算、稀疏矩陣計算、深度學習編譯器等領域取得多項重要成果 ， 發表多篇高性能計算與機器學習相關論文 。

    推薦閱讀