CuAsmRL: Optimizing GPU SASS Schedules via Deep Reinforcement Learning

TL;DR

CuAsmRL tackles the expensive and manual process of optimizing GPU SASS schedules for CUDA kernels by formulating scheduling as an assembly game and embedding an RL-based optimizer into the Triton compiler. It combines a hierarchical search (kernel autotuning + SASS optimization) with pre-game static analysis and a PPO-based RL agent that learns to interleave memory and compute instructions while respecting data dependencies. The approach yields up to $26\%$ improvements and an average of $9\%$ across representative LLM kernels, and reveals novel optimization moves (e.g., LDGSTS vs HMMA ordering) that generalize beyond hand-tuned schedules. This work demonstrates that automatic, hardware-native optimization at the SASS level can significantly boost performance while remaining transparent to kernel developers, with practical impact for large-scale LLM training and inference workloads.

Abstract

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.

Figures (13)

• Figure 1: CUDA compilation as documented by NVIDIA cuda_compiler. C++/Python and PTX are highlighted in green, indicating they are the common programming interfaces. SASS is a GPUs-native assembly and is highlighted in red, meaning it is undocumented. Cubin is an executable binary and is in gray.
• Figure 2: Overall workflow of CuAsmRL. CuAsmRL takes as input the source code targeting Triton's programming interface. Then it uses an autotuner to enumerate and find the optimal kernel configurations. Then the code is compiled with the optimal kernel configurations via Triton's compilation pipeline. Finally, an RL agent is trained to play the assembly game to optimize the SASS schedules, which outputs an optimized cubin.
• Figure 3: Assembly Game. At each iteration $i$, the SASS file is embedded, and the embedding is fed to the RL agent as state $S_i$. The RL agent is represented by a deep neural network. The output of the RL agent is an action $A_i$ that changes the SASS file. Then the mutated SASS file is assembled and sent to execution on the target GPU. A reward $R_i$ is sent back to the agent and the mutated SASS file is transitioned to the next state $S_{i+1}$.
• Figure 4: Embedding. Different fields of SASS instruction such as control code, opcode, and operands are embedded individually and then concatenated to a vector. Dummy values (-1) are used for the absent fields and operand padding. Different vectors are concatenated in a row-wise fashion. The final embedding of the assembly file becomes a matrix, which represents the state of the SASS file.
• Figure 5: An example of an action, which reorders the SASS instructions.