CUDAエージェントがLLMを訓練しGPUカーネルを作成、速度はtorch.compileを超越

はじめに

ByteDance Seedチームは、清華大学人工知能産業研究院（AIR）と共同でCUDA Agentを発表しました。これは、大規模インテリジェントエージェント強化学習システムであり、言語モデルが高性能なCUDAカーネルを生成することを教えることを目的としています。

この研究は、AIによる低レベルコード生成における長年の弱点に対処するものです。最先端の言語モデルは、コンパイル可能で正しく動作するCUDAコードを生成できますが、本番レベルのGPUワークロードにとっては、正しさだけでは不十分です。生成されたカーネルは、高度に最適化されたコンパイラ生成実装とも競合できなければなりません。

CUDA Agentはまさにこの2番目の問題に焦点を当てています。有効なCUDAコードを生成するだけでなく、意味のある実行時改善が達成されるまで、カーネルの性能分析、検証、最適化を反復的に行う方法を学習するのです。

正しいCUDAコードの生成だけでは不十分な理由

エージェント強化学習を導入する前から、Seed1.6ベースモデルはKernelBenchでかなり強力なCUDAコード生成能力を示していました。

研究者らが使用した250のKernelBenchレベル1〜3タスクにおいて、Seed1.6は74.0%の合格率を達成しました。言い換えれば、ベースモデルはベンチマークタスクの大部分に対して機能的に正しいソリューションを生成できるのです。

しかし、パフォーマンスは別の話です。

ベースモデルが生成したカーネルのうち、torch.compileより高速だったのはわずか27.2%で、コンパイラベースラインに対する幾何平均速度は0.69倍にとどまりました。

このギャップは、自動GPUカーネル生成における主要な課題を浮き彫りにしています。

言語モデルはCUDA構文を理解できますが、非効率なメモリアクセス、過剰なカーネル起動、不適切なタイル選択、不要な中間結果を生成する可能性がまだあります。

パフォーマンスプロファイリング

• CUDA向けの最適化ガイダンス
• 制限付きサンドボックス環境
• 実際の実行結果に連動した報酬シグナル

SKILL.mdファイルはまた、ベンチマーク結果を歪める可能性のある近道を制限しています。例えば、エージェントはカスタムカーネル実装内で任意のPyTorch操作に単純にフォールバックすることはできません。

これは重要です。強化学習システムは、意図された最適化問題を解決せずに報酬を最大化する方法を見つける可能性があるからです。

PPOトレーニングがエージェントを長い最適化トラジェクトリへと拡張

研究者らは**近位ポリシー最適化（PPO）**を使用してCUDA Agentを訓練しました。

主な課題は、GPU最適化が本質的に長期的なタスクであることです。モデルは、望ましい結果に到達するために、コード編集、コンパイル、デバッグ、プロファイリング、さらなる最適化を複数回繰り返す必要があるかもしれません。

そのため、この論文では、異なるコンテキスト長とトレーニング段階を用いて学習を安定化させています。

ベースモデルはSeed1.6で、合計2300億パラメータと230億アクティブパラメータを持つ混合エキスパートモデルです。

エージェント強化学習では：

• コンテキストウィンドウは131,072トークン。
• トレーニングリプレイでは最大150ラウンドのエージェントインタラクションを許可。
• 評価では最大200ラウンドのエージェントインタラクションを許可。
• モデルは150ステップのRLで訓練。
• Actorおよびcriticモデルは、完全な長期的最適化の前に多段階のウォームアップ戦略を採用。

これは、CUDAプロンプトと回答のペアでモデルを単純にファインチューニングするよりも複雑です。

目標は、実行フィードバックとの反復的相互作用を通じて、モデルが自身のカーネルを改善する方法を学ぶことです。

サンドボックスがコンパイルとGPUプロファイリングを分離

実行速度が報酬の一部であるため、信頼性の高いパフォーマンス測定が不可欠です。

そのため、研究者らはCPU-GPUリソースを分離したサンドボックスアーキテクチャを構築しました。

DockerベースのターミナルサンドボックスはコンパイルなどのCPU指向のタスクを処理し、検証とパフォーマンスプロファイリングは128基のNVIDIA H20 GPUを含む専用GPUサンドボックスプールに委任されます。

この分離には二重の目的があります。

第一に、コンパイルとエージェントインタラクションをGPUベンチマーキングから隔離します。第二に、専用のGPU割り当てにより、並行ワークロード間の干渉が減少し、レイテンシ測定がより安定します。

これは強化学習にとって重要です。タイミング測定がノイズの影響を受けると、モデルは信頼性の低い報酬シグナルを受け取り、誤った最適化行動を学習する可能性があります。

CUDA Agentが98.8%の合格率を達成し、96.8%のタスクでコンパイラを凌駕

最終システムは、KernelBenchレベル1〜3の250タスクで評価されました。

CUDA Agentは以下を達成しました：

• 98.8%の総合合格率
• PyTorch即時実行より98.4%高速
• torch.compileより96.8%高速
• 即時実行に対して2.60倍の幾何平均高速化
• torch.compileに対して2.11倍の幾何平均高速化

これらの結果は、Seed1.6の出発点から大幅な改善を示しています。

モデルは正しさテストに合格するコードを生成する能力が向上しただけでなく、生成されたカーネルがコンパイラベースラインを上回る可能性も高くなりました。

KernelBench難易度別の結果

KernelBenchの3つの難易度レベルすべてにおいて、パフォーマンス向上は堅調に維持されています。

レベル2のパフォーマンスは特に顕著です。

これらのタスクには演算子シーケンスが含まれており、融合とメモリ移動による最適化の機会は、静的コンパイラのヒューリスティック戦略が常に効果的に活用できるとは限りません。

論文の著者らは、反復学習型オプティマイザーが、ハードウェア固有のメモリアクセスパターン、タイル分割、融合選択を含む、より広範な実装戦略空間を探索できると考えています。

トレーニングデータが近日公開予定

現在、完全なトレーニングモデルの重みは、このプロジェクトの公開リリースには含まれていません。

ただし、研究者らはトレーニングスタックのいくつかの重要なコンポーネントを公開しています。

CUDA-Agent-Ops-6K

CUDA-Agent-Ops-6Kデータセットには、6,000の合成オペレーターレベルのトレーニングタスクが含まれています。

その構築プロセスは以下の通りです：

1. torch、transformersなどのライブラリからシード演算子を収集。
2. LLMを使用して複数の演算子をより複雑な融合タスクに組み合わせる。
3. 実行結果ベースのチェックによって生成されたタスクをフィルタリング。

フィルタリング段階では、ランダム性がある、単純すぎる、無効、または評価タスクと類似しすぎるケースを除外します。

このデータセットはHugging Faceで公開されており、CC BY 4.0ライセンスが適用されています。

SKILL.mdとエージェント環境

GitHubリポジトリには、CUDA向けのSKILL.md命令ファイルとエージェント作業環境も含まれています。

これらの資料は、最適化ワークフロー、利用可能なツール、制約条件、そしてエージェントをガイドするための実行環境設定を文書化しています。

報酬メカニズムとトレーニングスキーム

論文とリポジトリでは、報酬設計、ウォームアップ段階、クリティック初期化、およびPPOベースのトレーニングスキームが説明されており、長周期の最適化プロセスを安定化させます。

これは、最終的なベンチマークスコアを再現するだけでなく、システムプログラミングエージェントのトレーニングに関心のある研究者にとって特に有用です。

なぜこれがKernelBenchを超えて重要なのか

CUDAカーネル最適化は、現代のAIインフラストラクチャの大部分を支えています。

より高速なカーネルは以下を向上させることができます：

• モデルトレーニングのスループット
• LLM推論のレイテンシ性能
• GPU利用率
• サービスコスト効率
• リアルタイムAIパイプライン
• 高性能数値計算ワークロード

AIBaseの元記事は、そのAIにおける意義を指摘しています。

インフラ、大規模モデル推論サービス、自動運転、アルゴリズム取引など、わずかなレイテンシ差が極めて重要となる分野における潜在的な応用が考えられる。

CUDA Agent の成果は、従来のコンパイラが時代遅れになったことを意味するわけではない。

torch.compile は依然として強力な自動ベースラインであり、CUDA Agent の評価自体が、制御されたベンチマーク環境と特定の GPU 環境に依存している。

本研究が示す結論は範囲こそ狭いものの、依然として重要である。十分な実行フィードバックと専用の強化学習のサポートがあれば、言語モデルは、テストされたほとんどの GPU カーネルタスクにおいて静的コンパイラのベースラインを上回る最適化戦略を学習できるということだ。

これは、低レベルの性能エンジニアリングを、単発のコード生成ではなく、エージェントベースの学習が合理的に適用できる領域に変換することを意味する。

よくある質問

CUDA Agent とは何か？

CUDA Agent は、ByteDance Seed、清華大学智能産業研究院（AIR）、およびその共同 SIA-Lab によって開発された、大規模エージェント強化学習システムである。言語モデルを訓練し、CUDA カーネルを反復的に記述、検証、解析、最適化する。

CUDA Agent はどのモデルに基づいているか？

本研究では Seed1.6 をベースモデルとして使用している。論文では、総パラメータ数 2300 億、活性化パラメータ数 230 億の混合専門家モデルと説明されている。

KernelBench とは何か？

KernelBench は、言語モデルが PyTorch プログラムに対して正確かつ効率的な GPU カーネルを生成できるかを評価するためのオープンベンチマークである。CUDA Agent は、KernelBench のレベル 1 からレベル 3 を網羅する 250 のタスクで評価された。

CUDA Agent は torch.compile よりどのくらい高速か？

ベンチマーク全体を通して、CUDA Agent は torch.compile に対して 2.11 倍の幾何平均高速化 を達成した。生成されたカーネルは、評価タスクの 96.8% でコンパイラのベースラインより高速であった。

CUDA Agent は強化学習を使用するか？

はい。このシステムは PPO を使用し、多段階のウォームアップ、価値モデルの事前学習、堅牢な報酬設計、および長時間のマルチターンエージェント軌跡を組み合わせている。

1 回の CUDA Agent 最適化軌跡はどのくらいの長さまで実行できるか？

トレーニングのロールアウトでは最大 150 エージェントターン が許可され、評価では最大 200 ターン が許可されている。エージェントのトレーニングコンテキストウィンドウは 131,072 トークンである。

CUDA Agent はオープンソースか？

このプロジェクトは、GitHub リポジトリ、エージェント環境、SKILL.md、トレーニングレシピ、および CUDA-Agent-Ops-6K データセットを公開している。完全にトレーニングされたモデルの重みは、現在の公開リリースには含まれていない。

CUDA-Agent-Ops-6K とは何か？

CUDA-Agent-Ops-6K は、6,000 の合成演算子レベルの CUDA トレーニングタスクを含むデータセットである。大規模エージェント強化学習用に設計されており、タスクが実行可能で、決定的で、非自明であり、KernelBench 評価セットから分離されていることを保証するフィルタリングステップが含まれている。

関連ツール

• CUDA Agent：ByteDance Seed と清華大学智能産業研究院による CUDA カーネル生成システムの公式プロジェクトページ。
• CUDA Agent GitHub：エージェント環境、SKILL.md、公開プロジェクト資料を含む公式リポジトリ。
• CUDA-Agent-Ops-6K：プロジェクトとともに公開された公式の 6,000 サンプルトレーニングデータセット。
• KernelBench：大規模言語モデルが生成した GPU カーネルを評価するためのオープンベンチマーク。
• PyTorch：ディープラーニングフレームワーク。その eager モードと torch.compile 実行が研究における重要なベースラインとなる。
• NVIDIA CUDA：NVIDIA 公式の CUDA プログラミングおよび GPU カーネル開発ドキュメント。

関連リンク

• CUDA Agent プロジェクトページ：公式概要、ベンチマーク結果、プロジェクト図、論文、コード、データセットリンク。
• CUDA Agent 論文（arXiv）：完全な研究論文、「高性能 CUDA カーネル生成のための大規模エージェント強化学習」。
• CUDA Agent GitHub リポジトリ：公開された環境とプロジェクトファイルのソースコードリポジトリ。
• CUDA-Agent-Ops-6K データセット：6,000 の合成演算子タスクを含む公式 Hugging Face データセット。
• KernelBench GitHub リポジトリ：効率的な GPU カーネル生成のためのベンチマークおよび評価環境。
• PyTorchtorch.compile：研究で使用されるコンパイラベースラインの公式 PyTorch ドキュメント。
• NVIDIA CUDA C++ プログラミングガイド：NVIDIA 公式の CUDA プログラミングおよび最適化ガイド。

概要

CUDA Agent は、大規模言語モデルが生成するシステムコードにおける具体的な弱点を解決する。生成された CUDA コードは、正しいだけでなく、コンパイラが生成した代替案よりも実際に高速でなければならない。

Seed1.6 から出発し、研究者らは CUDA 専用のエージェント環境、実行フィードバック、堅牢な報酬設計、および長期的な視野での PPO トレーニングを使用して、成功率を 74.0% から 98.8% に、torch.compile よりも高速な割合を 27.2% から 96.8% に引き上げた。

このプロジェクトはまた、6,000 タスクのトレーニングデータセット、SKILL.md、エージェント環境資料、およびトレーニングレシピを公開しており、学習された GPU 最適化方向における研究者のさらなる作業のための具体的な基盤を提供している。

CUDA Agent の主な貢献は、性能エンジニアリング自体が、単発のコード生成による近似ではなく、反復的なエージェントループを通じて学習可能であることを実証したことにある。