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

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

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

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

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

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

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

稀疏矩陣在 LLM 推理、圖神經(jīng)網(wǎng)絡(luò)、科學計算中無處不在。但和稠密矩陣不同，稀疏矩陣的非零元素分布是不規(guī)則的，這導致一個核心問題：最優(yōu)的 CUDA 代碼實現(xiàn)取決于矩陣的具體結(jié)構(gòu)，而這個結(jié)構(gòu)只有在運行時才能知道。

換句話說，沒有一種「萬能」的高性能實現(xiàn)能應(yīng)對所有稀疏矩陣。工程師們不得不針對不同的稀疏模式手動調(diào)優(yōu)，這個過程既耗時又依賴經(jīng)驗。

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

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

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

SparseRL 怎么做到的？

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

SparseRL 把預(yù)訓練語言模型當作一個策略網(wǎng)絡(luò)，每生成一個 token 就是一次動作，而代碼的編譯結(jié)果和執(zhí)行時間就是獎勵信號。

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

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

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

為了讓模型真正「看懂」稀疏矩陣的結(jié)構(gòu)，研究團隊設(shè)計了一個關(guān)鍵技術(shù)：正弦位置嵌入。

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

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

另一個核心創(chuàng)新是層級獎勵函數(shù)。這個獎勵函數(shù)同時考慮兩個層面：正確性獎勵確保代碼能編譯、結(jié)果正確；效率獎勵則優(yōu)化執(zhí)行速度。設(shè)計邏輯是先保證「對」，再追求「快」。

效果如何？

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

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

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

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

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

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

意義與展望

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

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

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

作者介紹

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

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