中科院團(tuán)隊(duì)提出SparseRL，深度強(qiáng)化學(xué)習(xí)可自動(dòng)生成高性能CUDA代碼

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

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

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

目前，該項(xiàng)成果已入選 ICLR 2026 Oral。

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

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

要理解這項(xiàng)工作的價(jià)值，得先說說稀疏矩陣運(yùn)算的特殊性。

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

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

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

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

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

SparseRL 怎么做到的？

研究團(tuán)隊(duì)的思路很巧妙：既然執(zhí)行效率不可微，那就用強(qiáng)化學(xué)習(xí)來優(yōu)化。

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

整個(gè)訓(xùn)練過程分為三個(gè)階段：

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

圖片 2：展示三階段訓(xùn)練流程的整體框架圖

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

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

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

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

效果如何？

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

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

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

團(tuán)隊(duì)還做了消融實(shí)驗(yàn)來驗(yàn)證各個(gè)組件的必要性。

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

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

意義與展望

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

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

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

作者介紹

王耀宇，中國科學(xué)院計(jì)算技術(shù)研究所博士生（共同一作），主要研究方向?yàn)樯疃葘W(xué)習(xí)編譯優(yōu)化與高性能計(jì)算。

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