CUDA-L2: 강화 학습을 통한 행렬 곱셈 성능의 cuBLAS 능가

초록

본 논문에서는 대규모 언어 모델(LLM)과 강화 학습(RL)을 결합하여 Half-precision General Matrix Multiply(HGEMM) CUDA 커널을 자동으로 최적화하는 시스템인 CUDA-L2를 제안한다. CUDA 실행 속도를 RL 보상으로 활용하여 CUDA-L2는 1,000가지 구성에 걸쳐 HGEMM 커널을 자동 최적화한다. CUDA-L2는 널리 사용되는 {\it torch.matmul}부터 최신 Nvidia의 클로즈드 소스 라이브러리인 {\it cuBLAS}, {\it cuBLASLt}에 이르기까지 현재까지의 주요 행렬 곱셈 기준선을 체계적으로 능가한다. 오프라인 모드(커널을 시간 간격 없이 연속 실행)에서 CUDA-L2는 평균적으로 {\it torch.matmul} 대비 +22.0%, 최적 레이아웃 구성(normal-normal NN 및 transposed-normal TN)을 사용한 {\it cuBLAS} 대비 +19.2%, 휴리스틱 제안을 기반으로 알고리즘을 선택하는 {\it cuBLASLt-heuristic} 대비 +16.8%, 그리고 {\it cuBLASLt}의 제안 후보 중 최대 100개에서 가장 빠른 알고리즘을 선택하는 가장 경쟁력 있는 {\it cuBLASLt-AutoTuning} 모델 대비 +11.4%의 성능 향상을 보인다. 서버 모드(실시간 추론을 시뮬레이션하기 위해 커널을 무작위 간격으로 실행)에서는 이러한 속도 향상이 {\it torch.matmul}, {\it cuBLAS}, {\it cuBLASLt-heuristic}, {\it cuBLASLt-AutoTuning}에 대해 각각 +28.7%, +26.0%, +22.4%, +15.9%로 더욱 증가한다. CUDA-L2는 HGEMM과 같이 성능이 매우 중요하고 철저히 최적화된 커널조차 LLM 주도의 RL 자동화를 통해 인간이 실현하기 어려운 규모로 구성 공간을 체계적으로 탐색함으로써 개선될 수 있음을 보여준다. 프로젝트 및 코드는 github.com/deepreinforce-ai/CUDA-L2에서 확인할 수 있다.

English

In this paper, we propose CUDA-L2, a system that combines large language models (LLMs) and reinforcement learning (RL) to automatically optimize Half-precision General Matrix Multiply (HGEMM) CUDA kernels. Using CUDA execution speed as the RL reward, CUDA-L2 automatically optimizes HGEMM kernels across 1,000 configurations. CUDA-L2 systematically outperforms major matmul baselines to date, from the widely-used {\it torch.matmul} to state-of-the-art Nvidia's closed-source libraries, i.e., {\it cuBLAS}, {\it cuBLASLt}. In offline mode, where kernels are executed consecutively without time intervals, CUDA-L2 yields +22.0\% over {\it torch.matmul} on average; +19.2\% over {\it cuBLAS} using the optimal layout configuration (normal-normal NN and transposed-normal TN); +16.8\% over {\it cuBLASLt-heuristic}, which queries {\it cuBLASLt} library and selects the algorithm based on the heuristic's suggestion; and +11.4\% over the most competitive {\it cuBLASLt-AutoTuning} model, which selects the fastest algorithm from up to 100 candidates from {\it cuBLASLt}'s suggestions. In server mode, where kernels are executed at random intervals simulating real-time inference, the speedups further increase to +28.7\%, +26.0\%, +22.4\%, and +15.9\% for {\it torch.matmul}, {\it cuBLAS}, {\it cuBLASLt-heuristic}, and {\it cuBLASLt-AutoTuning} respectively. CUDA-L2 shows that even the most performance-critical, heavily-optimized kernels like HGEMM can be improved through LLM-guided RL automation by systematically exploring configuration spaces at scales impractical for humans. Project and code can be found at github.com/deepreinforce-ai/CUDA-L2

CUDA-L2: 강화 학습을 통한 행렬 곱셈 성능의 cuBLAS 능가

CUDA-L2: Surpassing cuBLAS Performance for Matrix Multiplication through Reinforcement Learning

초록

Support