ARGUS：データフロー不変条件に導かれるエージェント型GPU最適化

要旨: LLMベースのコーディング・エージェントは機能的に正しいGPUカーネルを生成できますが、行列乗算、注意（attention）、およびMixture-of-Experts（MoE）といった重要な計算において、その性能は手で最適化されたライブラリに比べて依然として大きく劣っています。ピーク時のGPU性能を引き出すには、タイル化、共有メモリへの段取り（staging）、ソフトウェア・パイプライン化、命令スケジューリングを含む、密に結合した最適化群にわたる協調的な推論が必要です。一方で既存のエージェントは疎な通過/不通過（pass/fail）のフィードバックに依存しているため、グローバルな制約違反を診断できません。
本研究では、データフロー不変条件（data-flow invariants）によってこれを解決するエージェント的フレームワークArgusを提案します。Argusは、カーネル実行全体でデータをどのように取り回す（choreograph）べきかをエンコードするコンパイル時仕様を導入します。Argusは、低レベル表現を隠しつつ、タイルベースのPython的DSLによってハードウェア命令とコンパイラ方針を公開します。このDSLは、データおよび制御フローを通じて記号的な注釈（symbolic annotations）を伝播させるためのタグ関数と、利用箇所（use sites）での関係制約を強制するタグアサーションを提供します。違反が起きた場合、コンパイラはスレッド、データ要素、プログラム上の位置（program point）を特定した具体的な反例を返し、対象を絞った修正のための密で構造化されたフィードバックを可能にします。不変条件は、レイアウト代数（layout algebra）上での抽象解釈とSMTソルビングによりコンパイル時に検証され、実行時オーバーヘッドはゼロです。状況（in-context）の強化学習プランナが最適化の選択と、有効な不変条件の合成を学習し、GPU最適化技術のキュレーションされた知識ベースによって支えられます。
GEMM、フラッシュ注意（flash attention）、およびMoEカーネルにおいて、LLM推論でGPU時間の90%以上を占めるケースを対象に、AMD MI300X GPU上でArgusを評価します。生成されたカーネルは、最先端の手最適化アセンブリのスループットに対して99-104%を達成し、既存のエージェント的システムより2-1543倍高速です。さらにArgusは200のKernelBenchタスクに一般化し、レベル1を100%、レベル2を90%の問題で解決します。