arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2501.08071cs.ARcs.LG

CuAsmRL:通过深度强化学习优化GPU SASS调度

CuAsmRL: Optimizing GPU SASS Schedules via Deep Reinforcement Learning

University of Cambridge(剑桥大学)

Guoliang He, Eiko Yoneki

AI总结:

针对专用CUDA内核仍有性能提升空间的问题,提出CuAsmRL方法,通过构建汇编游戏训练强化学习智能体优化GPU SASS调度,可透明提升现有内核性能最高26%、平均9%,还能揭示潜在优化操作。

摘要:

大语言模型(LLM)以其巨大的计算需求著称。为降低成本,研究人员开发了专用CUDA内核,这类内核通常融合多个张量运算以尽可能最大化GPU利用率。然而,这些专用内核仍存在性能提升空间:CUDA汇编专家表明,手动优化GPU SASS调度可带来更优性能,且人们主要通过试错法手动寻找最佳GPU SASS调度。\n 在本研究中,我们采用一种自动化方法优化GPU SASS调度,该方法可集成至现有编译器框架中。自动化优化的核心是训练一个RL(强化学习)智能体以模仿人类专家执行手动调度的方式。为此,我们构建了一个汇编游戏,RL智能体可在其中探索以寻找最佳GPU SASS调度。该汇编游戏从经-O3优化的SASS调度启动,RL智能体可迭代施加动作来变异当前调度。若变异后的调度在GPU上执行时获得更高吞吐量,则生成正奖励。实验表明,CuAsmRL可透明地进一步提升现有专用CUDA内核的性能,最高达26%,平均提升9%。此外,它还可作为工具,揭示自动学习到的潜在优化操作。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑