MaxKernel、多エージェントで TPU カーネル生成を自動化

アクセラレーター向けの高性能カーネル開発は、通常のコード生成より難しい。演算の正しさだけでなく、メモリアクセス、並列化、データ配置、コンパイラの制約、そして特定チップの実行特性まで考慮する必要があるからだ。MaxKernel は、大規模言語モデルをコンパイラとハードウェアのフィードバックが得られる開発ループに組み込み、コードを一度生成して終わるのではなく、コンパイル、テスト、計測、修正を繰り返す。

三つの開発パラダイム

MaxKernel は複数のエージェントで構成され、計画、実装、自己デバッグ、正確性テスト、ハードウェアプロファイリングを専門のサブエージェントが分担する。その共通基盤の上に、次の三つの方式を用意している。

• Human-in-the-Loop：人間とエージェントが段階的に設計を進める。重要な方針を開発者が決められるため、探索の初期段階や共同開発に向く。
• Autonomous：性能指標とハードウェア trace を利用し、エージェントが自律的に実装を改良する。反復的な試行錯誤を自動化できる。
• グラフベース自律探索：複数の実装候補と最適化経路を探索グラフとして扱い、単一の編集経路に依存せず、より広い設計空間を調べる。

この構成は、カーネル開発が単純なコード補完ではなく、実験の連続であることを反映している。人間による方向付けが必要な場面と、ベンチマークを繰り返す作業を自動化する場面を、同じシステム内で切り替えられる点が特徴だ。

実測フィードバックを推論に取り込む

TPU カーネルでは、同じ計算結果を返す実装でも、タイル分割や配置、メモリアクセスの違いによって性能が変わる。したがって、事前学習した知識だけで最適解を決めるのは難しい。MaxKernel はコンパイル結果、テスト結果、プロファイリング情報を判断材料として扱い、次の候補を更新する。

評価には、50種類の多様な TPU カーネル課題を含む JaxBench と、先端オープンソースモデルに由来する複雑な実ワークロードが使われた。論文によれば、生成された実装は幅広い課題で高い最適化水準に達し、専門家が手作業で調整したベースラインに近い結果を示した。提供された素材には詳細な高速化倍率がないため、特定のカーネルの記録というより、開発ワークフロー全体の有効性を示す結果として読むべきだろう。

意義と残る課題

MaxKernel は、エージェントが実験を組み立て、コンパイラが即時の技術的フィードバックを返し、ハードウェアが最終的な性能を判定する分業モデルを示している。TPU に限らず、GPU など異なる AI アクセラレーターの自動チューニングにも応用できる可能性がある。

一方で、探索コスト、プロファイリングで観測できる範囲、テストの網羅性は依然として課題だ。新しいアーキテクチャーや未知の演算の組み合わせ、実運用の負荷でも安定して機能するかは、追加の検証を必要とする。MaxKernel はオープンソースとして公開されており、エージェントの役割分担や探索手法をコミュニティが検証できる。

出典：Hugging Face Daily Papers