ByteDance SeedとTsinghua AIRは、コンパイラを上回るGPUカーネルを記述するためにSeed1.6大規模言語モデルを訓練するアジェンティック強化学習システムであるCUDA Agentを発表した。

  • システムは、プロファイリングサンドボックスと正確性チェック内で、131,072トークンのコンテキストを用いた150ステップのPPO訓練を使用している。
  • KernelBenchベンチマークで98.8%のパス率とtorch.compileより96.8%高速な率を達成し、幾何平均速度向上は2.11倍となった。
  • 研究チームはCUDA-Agent-Ops-6Kデータセット、SKILL.md仕様、報酬レシピを公開したが、訓練済みエージェントの重みは公開していない。
  • アブレーション実験では、エージェントループを削除するとコンパイラより高速な率が96.8%から14.1%に低下し、アジェンティックアプローチの重要性が浮き彫りになった。

本成果は、反復プロファイリングを伴う実際のCUDA開発環境にLLMを配置することで、生成されたカーネルと手書きまたはコンパイラ生成のカーネル間の性能差を大幅に縮小できることを示している。