
要約
本稿では、大規模言語モデル(LLM)と強化学習(RL)を統合したシステム「CUDA-L2」を提案する。このシステムは、半精度一般行列乗算(HGEMM)のCUDAカーネルを自動最適化することを目的としており、CUDAの実行速度をRLの報酬関数として用いる。CUDA-L2は、1,000種類のカーネル設定を対象に自動最適化を実行し、これまでに報告された主要な行列乗算基準(matmul baselines)を体系的に上回る性能を発揮する。広く使用されているこの http URL から、最新のNVIDIA社の非公開ライブラリであるcuBLAS、cuBLASLtに至るまで、あらゆる基準を上回る成果を達成した。オフラインモード(カーネルが時間的間隔を置かずに連続実行される状況)では、CUDA-L2はこの http URL に対して平均で+22.0%、最適なレイアウト設定(normal-normal NN および transposed-normal TN)を用いたcuBLASに対して+19.2%、heuristicに基づくアルゴリズム選択を行うcuBLASLt-heuristicに対して+16.8%、cuBLASLtの提案から最大100件の候補から最速のアルゴリズムを自動選択するcuBLASLt-AutoTuningに対して+11.4%の速度向上を達成した。サーバーモード(リアルタイム推論を模倣するように、カーネルがランダムな間隔で実行される状況)では、各基準に対する速度向上はさらに顕著に増加し、それぞれこの http URL で+28.7%、cuBLASで+26.0%、cuBLASLt-heuristicで+22.4%、cuBLASLt-AutoTuningで+15.9%を達成した。本研究は、HGEMMのような性能が極めて重要で、長年にわたり高度に最適化されたカーネルであっても、LLMを用いた強化学習による自動化によって、人間が扱うには現実的でない規模の設定空間を体系的に探索することで、さらなる性能向上が可能であることを示している。