CUDA-L2：強化学習による行列乗算のcuBLAS性能超越

要旨

本論文では、大規模言語モデル（LLM）と強化学習（RL）を組み合わせて、半精度汎用行列乗算（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}ライブラリに問い合わせてヒューリスティックの提案に基づきアルゴリズムを選択する{\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