CUDA-L1: コントラスティブ強化学習を用いたCUDA最適化の向上

大規模言語モデルの急速な進歩により、GPU計算リソースへの需要が指数関数的に増加しており、自動化的なCUDA最適化戦略の緊急な必要性が生じています。最近のLLM(Large Language Model)の進展はコード生成に有望であることが示されていますが、現行の最先端モデル(例:R1, o1)はCUDA速度向上において低い成功率しか達成していません。 本論文では、CUDA最適化のための自動強化学習フレームワークであるCUDA-L1を紹介します。CUDA-L1はCUDA最適化タスクで性能向上を達成しています:NVIDIA A100上で訓練された結果、KernelBenchの250個すべてのCUDAカーネルに対して平均で17.7倍の高速化を実現し、最大で449倍の高速化も達成しました。さらに、このモデルはGPUアーキテクチャ間での優れた移植性も示しており、A100向けに最適化されたにもかかわらずH100では平均17.8倍、RTX 3090では19.0倍、L40では16.5倍、H800では14.7倍、H20では13.9倍の高速化を達成しています。 これらのベンチマーク結果以外にも、CUDA-L1は以下の特徴的な性質を示しています:1) 多様なCUDA最適化技術を見出し、それらを戦略的に組み合わせて最適な性能を達成する方法を学習します;2) CUDA最適化の基本原理を見出します;3) 非自明なパフォーマンスボトルネックを特定し、パフォーマンスに悪影響を与えると思われる最適化手法を拒否します。 CUDA-L1の能力は、強化学習が人間の専門知識やドメイン知識なしに速度向上に基づく報酬信号のみで初期的に性能が低いLLMを効果的なCUDAオプティマイザに変えることができることを示しています。より重要なのは、訓練されたRL(Reinforcement Learning)モデルが獲得した推論能力を新しいカーネルに拡張できることです。このパラダイムは自動化的なCUDA操作最適化の可能性を開き、GPU効率の大幅な向上とGPU計算リソースへの増大する圧力緩和につながる可能性があります。