SparseRL 框架的核心创新 - 中科院计算所团队提出名为 SparseRL 的新框架，首次将深度强化学习引入稀疏 CUDA 代码生成任务，旨在让 AI 根据稀疏矩阵结构自动生成最优的 CUDA 实现代码[2] - 框架将预训练语言模型作为策略网络，将生成 token 视为动作，以代码编译结果和执行时间作为奖励信号[12] - 框架采用三阶段训练流程：在大量 CUDA 代码语料上进行预训练；用“稀疏矩阵-正确代码”配对数据进行监督微调；引入深度强化学习，以编译正确性和执行效率为奖励进行优化[18] 技术实现的关键组件 - 设计了正弦位置嵌入技术，对稀疏矩阵的非零元素行列索引分别进行正弦/余弦编码，使模型能捕捉二维坐标间的空间关系，理解矩阵结构[13] - 设计了层级奖励函数，同时考虑正确性奖励（确保代码能编译、结果正确）和效率奖励（优化执行速度），遵循先保证“对”再追求“快”的逻辑[14] - 消融实验证明，去掉强化学习阶段后性能显著下降，去掉正弦嵌入后模型编译率下降，仅用正确性奖励则代码执行速度不足，验证了各核心组件的必要性[20] 实验性能与效果评估 - 在经典的稀疏矩阵-向量乘法任务上，相比传统监督学习方法，该方法使编译成功率提升 20%，平均代码执行速度提升 30%[3][16] - 在稀疏矩阵-向量乘法和稀疏矩阵-稠密矩阵乘法两个任务上验证有效，模型能根据不同的稀疏结构自动选择不同的代码策略，在对角型、带状型、随机稀疏型等多种矩阵上均有优势[16] - 部分场景下生成的代码性能接近或超越了手工调优的水平[16] 研究背景与问题挑战 - 稀疏矩阵在大型语言模型推理、图神经网络、科学计算中无处不在，但其非零元素分布不规则，导致最优的 CUDA 代码实现取决于仅在运行时可知的具体矩阵结构[5] - 不存在一种“万能”的高性能实现能应对所有稀疏矩阵，工程师需针对不同稀疏模式手动调优，过程耗时且依赖经验[6] - 现有 AI 代码生成方法存在三大局限：传统监督学习只关心代码正确性不关心执行速度；执行效率指标不可微，无法通过传统反向传播优化；稀疏矩阵输入与 CUDA 代码间存在巨大语义鸿沟[17] 研究意义与未来展望 - 该研究代表了一种范式转变，将代码生成的目标从“生成能运行的代码”转向“生成高性能代码”[22] - 对于高性能计算工程师和 AI 基础设施开发者而言，展示了让 AI 处理繁琐性能优化工作的新可能，使人类能聚焦于更高层次的设计[22] - 未来计划将方法扩展到多 GPU 分布式稀疏计算，探索与传统自动调优技术的结合，支持更多类型的稀疏算子，并研究如何降低强化学习训练成本[22] 研究成果与作者信息 - 该项研究成果已入选 ICLR 2026 的口头报告[4] - 论文共同一作为中国科学院计算技术研究所博士生王耀宇，主要研究方向为深度学习编译优化与高性能计算[24] - 通讯作者为中国科学院计算技术研究所研究员谭光明，主要从事高性能计算、GPU 编译优化与深度学习系统研究，在相关领域取得多项重要成果[24]