CUDA-L1: Improving CUDA Optimization via Contrastive Reinforcement Learning

基本情報

• arXiv ID: 2507.14111v1 (https://arxiv.org/abs/2507.14111)
• 著者: Xiaoya Li・Xiaofei Sun・Albert Wang・Jiwei Li・Chris Shum
• 所属: DeepReinforce Team (research@deepreinforce.ai)
• 投稿日: 2025年07月19日
• カテゴリ: cs.LG, cs.DC

簡単に説明すると

CUDA-L1は、強化学習（RL）を使ってCUDAコードを自動的に最適化するシステムです。GPUの計算需要が爆発的に増加する中、手作業でのCUDA最適化は時間がかかり難しいという課題がありました。CUDA-L1は「対比強化学習（Contrastive RL）」という新しい手法を採用し、過去に生成されたCUDAコードとその実行性能を比較分析することで、効果的な最適化戦略を学習します。NVIDIA A100で訓練されたモデルは、KernelBenchベンチマークの250個のCUDAカーネル全体で平均17.7倍、最大449倍の高速化を達成しました。プロジェクトページは https://github.com/deepreinforce-ai/cudal1.github.io で公開されています。

1. 研究概要

1.1 背景と動機

大規模言語モデル（LLM）の急速な発展により、GPU計算リソースへの需要が指数関数的に増加しています。従来のCUDA最適化は極めて手作業かつ時間集約的なプロセスでした。熟練したエンジニアがメモリアクセスパターンを分析し、スレッドブロック構成を試行錯誤し、広範なプロファイリングを通じて最適化していました。

最近のLLMの進歩、特に強化学習（RL）を活用したモデル（OpenAI-o1、DeepSeek-R1など）は、コード生成とアルゴリズム設計において驚くべき能力を示しています。CUDA最適化は実行速度という明確な報酬信号を提供するため、RLにとって理想的なタスクです。しかし、現在の最先端モデルでも、CUDAコードの最適化成功率は低く（KernelBenchで約15%）、これは主に訓練データセット内のCUDAコードの不足が原因です。

1.2 主要な貢献

本研究では、これらの制限に対処し、自動CUDA最適化のためのLLMの実用化を実現するため、以下の貢献をしました。

• 対比強化学習（Contrastive RL）という新しいRL手法を提案し、過去のCUDAコードとその性能を比較分析することで効果的な最適化戦略を学習
• CUDA最適化において前例のない性能向上を達成（KernelBenchで平均17.7倍、最大449倍の高速化）
• 異なるGPUアーキテクチャへの優れた移植性を実証（H100で17.8倍、RTX 3090で19.0倍などの高速化）
• 様々なCUDA最適化技術の自動発見と、それらの戦略的な組み合わせによる最適性能の達成

2. 提案手法

2.1 手法の概要

CUDA-L1は3段階のパイプライン訓練戦略を採用しています。第1段階では、データ拡張による教師あり微調整（SFT）を行い、モデルのCUDAパターンへの露出を拡大します。第2段階では、自己教師あり学習により、CUDAセマンティクスとプログラミング原理の深い理解を獲得します。第3段階では、対比強化学習により、実行速度を最適化し、高性能なCUDA実装を生成します。

フレームワーク全体を通じて、実行可能性（コンパイルと実行の成功）、正確性（参照実装と同等の出力）、成功（実行可能かつ正確）という3つの重要な定義を採用しています。

2.2 技術的詳細

対比強化学習の核心は、LLMのプロンプトに複数のコードバリアントとその対応する高速化スコアを直接組み込むことです。モデルは単にコードを生成するのではなく、まず特定の実装が優れた性能を達成する理由について比較分析を行い、その洞察に基づいて改良されたソリューションを合成するよう訓練されます。

プロンプト構造には、タスクの説明、過去のCUDAコードとそのスコア、生成プロトコル、および要件と制限が含まれます。モデルの応答は、性能分析、アルゴリズム設計、コード実装の3つの構造化されたコンポーネントを含む必要があります。

RL訓練では、Group Relative Policy Optimization (GRPO)戦略を採用し、報酬をグループ内で正規化して安定した訓練を確保します。報酬計算では、専用GPU割り当て、ペア実行とランダム化、拡張測定ウィンドウなど、ノイズを最小化するための厳格な測定戦略を実装しています。

2.3 新規性

既存のRL手法との主な違いは、報酬情報を直接推論プロセスに組み込む点です。標準的なRLアルゴリズムでは、報酬信号はパラメータ更新にのみ使用され、LLMへの入力として提供されません。対比RLでは、この二重利用戦略により、基盤モデルの強化と固定パラメータソリューションの最適化という2つの相補的な次元で反復最適化が可能になります。

また、進化的LLMアプローチと比較して、対比RLは勾配ベースのパラメータ更新を採用し、モデルの能力を継続的に強化します。これにより、固定パラメータに依存する進化的アプローチよりも優れた表現能力とタスク適応性を実現しています。

3. 実験結果

3.1 実験設定

評価はKernelBenchデータセットで実施されました。このデータセットは250個のPyTorchワークロードを含み、3つの階層レベルに構造化されています。レベル1は単一のプリミティブ操作（100タスク）、レベル2は融合最適化の恩恵を受ける演算子シーケンス（100タスク）、レベル3は完全なMLアーキテクチャ（50タスク）です。

すべての評価はNVIDIA A100 PCIeで実行されました。各タスクは20分間の時間制限内にランダムな順序で実行されます。最終的な評価スコアは、割り当てられた時間枠内のすべての実行ラウンドの平均高速化比として計算されます。

3.2 主要な結果

CUDA-L1はKernelBenchタスク全体で優れた効果を示し、平均17.7倍の高速化を達成し、最大利得は449倍に達しました。レベル1（単一操作）では平均12.3倍、中央値1.65倍の高速化を達成しました。レベル2（演算子シーケンス）では平均6.39倍、中央値1.61倍の高速化を示しました。特に注目すべきは、レベル3（複雑なMLタスク）が劇的に高い最適化能力を示し、平均50.8倍、中央値2.66倍の高速化を達成したことです。

全体の成功率は99.6%で、レベル2-3では100%の成功率を達成し、アプローチの堅牢性を検証しています。高速化比の分布は、ほとんどの高速化（約68%のケース）が1倍から2.72倍の範囲に収まることを示しています。

3.3 既存手法との比較

バニラ基盤モデル（OpenAI-o1、DeepSeek-R1、DeepSeek-V3、Llama 3.1-405B）と比較して、CUDA-L1は優れた性能を示しました。最高性能のモデルでさえ、平均して20%未満のタスクでしか参照カーネルを上回る高速化を達成できませんでした。

進化的LLMモデルは、同じパラメータセットを共有するバニラ基盤モデルセットアップと比較して大幅な性能向上を示しました。これは、比較分析を活用することが、ドメイン固有の知識を必要とする直接出力生成よりも効果的であることを示しています。

CUDA-L1コンポーネントの異なる組み合わせを比較すると、SFTのみからSFT+自己教師あり学習への段階的な成功率と高速化率の増加が観察されました。対比RLを使用したstage1+2は、バニラGRPOを使用したstage1+2を上回り、RL訓練における比較分析の有効性を実証しました。

4. 実用性評価

4.1 実装の容易性

CUDA-L1は、DeepSeek-V3-671Bをモデルバックボーンとして使用し、3段階のパイプライン訓練戦略を採用しています。各段階は明確に定義されており、段階的な実装が可能です。特に、対比RLフレームワークは既存のRL手法を拡張したものであり、理解と実装が比較的容易です。

4.2 計算効率

A100で最適化されたカーネルは、他のGPUアーキテクチャでも優れた性能を示します。H100 XSMでは平均17.8倍、最大1,001倍の高速化を達成しました。RTX 3090では平均19.0倍の高速化を達成しました。L40、H800 XSM、H20でもそれぞれ16.5倍、14.7倍、13.9倍の平均高速化を達成しました。

報酬ハッキングへの対処として、ハイパーパラメータ操作や結果キャッシングなどの悪用的戦略を防ぐための明示的な指示がトレーニングプロンプトに組み込まれています。

4.3 応用可能性

CUDA-L1は10種類の主要な最適化技術を自動的に発見しました。これらにはメモリレイアウト最適化、メモリアクセス最適化、演算融合、メモリフォーマット最適化などが含まれます。また、メモリ合体、ワープレベル最適化、最適化されたスレッドブロック構成、共有メモリ使用、レジスタ最適化、ストリーム管理も含まれます。

これらの技術の一部はCUDA最適化コミュニティで既に広く採用されていますが、その他は十分に活用されていません。CUDA-L1はこれらの技術を戦略的に組み合わせて最適な性能を達成する能力を示しています。

5. まとめと所感

5.1 論文の意義

CUDA-L1は、強化学習が最初は性能の低いLLMを効果的なCUDA最適化器に変換できることを実証しました。高速化ベースの報酬信号のみを使用し、人間の専門知識やドメイン知識なしに動作します。RLシステムはCUDA最適化パターンを識別し、新しい技術を発見します。さらに、それらを合成して高速化を達成し、獲得した推論能力を新しいカーネルに拡張できます。

特に注目すべきは、GPU最適化が加算的ではなく乗算的な効果を示すという発見です。特定の技術が他の技術の重要な有効化要因として機能します。効果的なGPU最適化には個々の技術だけでなく、それらの相乗的な相互作用を理解することが必要です。

5.2 今後の展望

CUDA-L1のパラダイムは、CUDA操作の自動最適化の可能性を開き、GPU効率を向上させ、GPU計算リソースへの増大する圧力を緩和する約束を示しています。異なるGPUタイプ専用に訓練されたカーネルをCUDA-L1の更新版でリリースする計画があります。

今後の研究方向として、アーキテクチャ固有の最適化によるさらなる性能向上が考えられます。より複雑なCUDAアプリケーションへの拡張も重要です。他のドメイン（例：分散コンピューティング、ヘテロジニアスコンピューティング）への対比RL手法の適用も有望です。また、発見された最適化原則の理論的理解を深めることも重要な課題です。