powered by TechFeed
表示モード
主要ニュース

ByteDanceがAIでGPUカーネルを自動最適化 — PyTorchの自動コンパイラを96.8%のケースで上回り、最大73倍高速化

8月17日、MarkTechPostが「ByteDance Seed and Tsinghua AIR Introduces CUDA Agent: A Large-Scale Agentic RL System for CUDA Kernel Generation」と題した記事を公開した。この記事では、ByteDance SeedおよびTsinghua AIRが開発したCUDAカーネル自動生成エージェント「CUDA Agent」について詳しく紹介されている。

8月17日、MarkTechPostが「ByteDance Seed and Tsinghua AIR Introduces CUDA Agent: A Large-Scale Agentic RL System for CUDA Kernel Generation」と題した記事を公開した。この記事では、ByteDance SeedおよびTsinghua AIRが開発したCUDAカーネル自動生成エージェント「CUDA Agent」について詳しく紹介されている。


LLMは「正しいCUDA」は書けるが「速いCUDA」は書けなかった

GPUカーネルの手書き最適化は、AIインフラにおける数少ない「人間が機械に勝てる領域」のひとつだった。torch.compile(PyTorchの自動最適化コンパイラ)はその差を縮めてきたが、カーネルレベルの細かいチューニングには依然として専門家の手が必要とされてきた。

ByteDance SeedとTsinghua AIRが発表したCUDA Agentは、この問題を強化学習(RL)で正面から攻略する。ベースモデルとして使用したSeed1.6(アクティブパラメータ23B、総パラメータ230BのMoEモデル)は、ベンチマーク「KernelBench」の250タスクに対して74.0%の正答率を示す一方、torch.compileより速いカーネルを出せたのは**わずか27.2%**にとどまっていた。幾何平均スピードアップは0.69×、つまりコンパイラより平均的に遅い出力を生成していたことになる。

CUDA Agentはこれを大幅に改善した。トレーニング後の結果は以下の通りだ。

指標 ベースモデル(Seed1.6) CUDA Agent
パス率 74.0% 98.8%
torch.compile超え率 27.2% 96.8%
幾何平均スピードアップ 0.69× 2.11×

Level 3(最難関タスク)でのtorch.compile超え率は**90.0%**で、元記事ではClaude Opus 4.5(50.0%)やGemini 3 Pro(52.0%)との比較も示されている(※これらのモデル名・数値は元記事の記述に基づく。各モデルの評価条件の詳細は元記事を参照されたい)。


エージェントループと報酬設計が肝

CUDA Agentの特徴は、LLMを「ツールを使いながら試行錯誤するエージェント」として動かす点だ。環境はOpenHandsのツール群(Bash、ファイル読み書き、Grep、Notebook編集など)をベースに構築されており、ReActパターン(Reasoning+Actingを交互に行う推論スタイル)で動作する。

エージェントは以下のサイクルを繰り返す:

  1. PyTorchモデルをプロファイリング
  2. model_new.pyにカスタムカーネルを実装
  3. GPUサンドボックスでコンパイル・検証
  4. torch.compileより5%以上速くなるまで反復

報酬ハッキング対策も手厚い。検証スクリプトをパーミッションでロック、torch.nn.functionalへのフォールバックを禁止、5種類のランダム入力でチェック、デバイス同期とウォームアップを伴うプロファイリング、Webサーチツールの排除——計5つの対策が施されている。

報酬は生のスピードアップ比ではなく離散値(−1、1、2、3)を採用する。torch.compileとeagerモードの両方を5%以上超えた場合のみ最高スコア3が与えられる。このマイルストーン型報酬は、生のスピードアップ比を使った場合(faster-than-compile率60.4%)と比較して、36.4ポイントの差を生んだ。

また、PPOによるRL本番学習の前にRFT(Reinforcement Fine-Tuning)とバリュー事前学習を組み合わせることで、17ステップで発散していた学習軌跡を150ステップの安定した学習に変えた。アブレーション実験では、エージェントループを除去した場合にfaster-than-compile率が96.8%から**14.1%**まで落ちており、ループ構造の貢献が際立っている。


具体的な最適化事例

論文中のケーススタディが示す数字は極端だ。

  • 対角行列乗算をrow-wiseスケーリングとして書き直し:torch.compile73.31×
  • matmul→divide→sum→scaleの演算チェーンを再順序化・融合:24.04×
  • ResNetのBasicBlockでBatchNormをConvに折り込みcudnnConvolutionBiasActivationForwardを活用:3.59×

公開されたもの・されないもの

モデル重みは非公開。Seed1.6自体がクローズドな独自モデルであるため、完全な再現には128台のNVIDIA H20 GPUが必要なプロファイリングサンドボックスも含めてフロンティアラボ級のリソースが求められる。

一方で以下は公開されている:

  • **CUDA-Agent-Ops-6K**:6,000サンプルのデータセット(83.77%が2演算子合成タスク)
  • SKILL.md:エージェントへのCUDA開発手順仕様
  • 報酬設計とウォームアップのレシピ

データセットはPyTorchとHugging Face transformersライブラリから演算子を収集し、LLMが最大5クラスを組み合わせてフュージョン層を生成するパイプラインで作られている。KernelBenchとのAST類似度が0.9以上のサンプルは除外されている。


詳細はByteDance Seed and Tsinghua AIR Introduces CUDA Agent: A Large-Scale Agentic RL System for CUDA Kernel Generationを参照していただきたい。