CuAsmRL: optimizing GPU SASS schedules via deep reinforcement learning
Large language models (LLMs) are remarked by their sub- stantial 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 op- timize 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 formu- late an assembly game, where RL agents can play to find the best GPU SASS schedules. The assembly game starts from a -O3 optimized SASS schedule, and the RL agents can itera- tively 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 can be used as a tool to reveal potential optimization moves learned automatically.
Tue 4 MarDisplayed time zone: Pacific Time (US & Canada) change
14:00 - 15:00 | |||
14:00 20mTalk | Code Generation for Cryptographic Kernels Using Multi-word Modular Arithmetic on GPU Main Conference | ||
14:20 20mTalk | CuAsmRL: optimizing GPU SASS schedules via deep reinforcement learning Main Conference | ||
14:40 20mTalk | Proteus: Portable Runtime Optimization of GPU Kernel Execution with Just-In-Time Compilation Main Conference Giorgis Georgakoudis Lawrence Livermore National Laboratory, Konstantinos Parasyris Lawrence Livermore National Laboratory, David Beckingsale Lawrence Livermore National Laboratory |