用强化学习自动优化GPU汇编调度，性能最高提升26%。
CuAsmRL: Optimizing GPU SASS Schedules via Deep Reinforcement Learning
- 构建汇编游戏，让RL代理通过试错优化SASS调度
- 在现有CUDA核上实现平均9%、最高26%的性能提升
- 可嵌入编译器，适合加速GPU计算的研发人员
大型语言模型计算需求巨大。为降低成本，研究者开发专用CUDA内核，常融合多个张量操作以最大化GPU利用率。然而这些内核仍可能未达最优，因人工手动优化SASS调度可带来更好性能，但依赖试错。本文提出自动优化方法CuAsmRL，通过训练强化学习代理模仿专家调度行为。我们设计一个汇编游戏：从-O3优化的SASS调度开始，代理可迭代执行动作改变调度；若新调度在GPU上运行获得更高吞吐，则获正奖励。实验表明，CuAsmRL可透明地进一步提升现有专用CUDA内核性能，最高达26%，平均提升9%。该方法还可作为工具揭示自动学习到的优化策略。
原文摘要 · Abstract (English)
Large language models (LLMs) are remarked by their substantial computational requirements. To mitigate the cost, researchers develop specialized CUDA kernels, which often fuse several tensor operations to maximize the utilization of GPUs as much as possible. However, those specialized kernels may still leave performance on the table as CUDA assembly experts show that manual optimization of GPU SASS schedules can lead to better performance, and trial-and-error is largely employed to manually find the best GPU SASS schedules. In this work, we employ an automatic approach to optimize GPU SASS schedules, which thus can be integrated into existing compiler frameworks. The key to automatic optimization is training an RL agent to mimic how human experts perform manual scheduling. To this end, we formulate an assembly game, where RL agents can play to find the best GPU SASS schedules. The assembly game starts from a \textit{-O3} optimized SASS schedule, and the RL agents can iteratively apply actions to mutate the current schedules. Positive rewards are generated if the mutated schedules get higher throughput by executing on GPUs. Experiments show that CuAsmRL can further improve the performance of existing specialized CUDA kernels transparently by up to $26\%$, and on average $9\%$. Moreover, it is used as a tool to reveal potential optimization moves learned automatically.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。