MaxKernelとは？LLMエージェントがTPUカーネルを自動生成し専門家に匹敵

• Google の研究チームが、TPU 向けの高速カーネルを LLM エージェントに自動生成させる枠組み MaxKernel を公開しました
• 計画、実装、自己デバッグ、テスト、自動チューニング、プロファイリングを専門サブエージェントに分担させ、コンパイラとプロファイラの出力を次の反復に返します
• 50 タスクのベンチマーク JaxBench で幾何平均 1.58 倍、実モデルの 8 ワークロードでは専門家の手動実装（2.02 倍）を上回る 2.32 倍を記録しました

研究の背景

深層学習モデルの実行速度は、行列積や Attention といった演算を実際のハードウェア上でどれだけ効率よく回せるかで決まります。コンパイラ（ここでは XLA）が自動生成するコードは汎用性が高い反面、メモリ階層の使い方や処理の重ね合わせまでは詰めきれず、性能を引き出すには専用のカーネル、つまりハードウェアに密着した低レイヤの実装を人手で書く必要がありました。

この作業は TPU では特に負担が重くなります。HBM（大容量の主記憶）と VMEM（演算器に近い高速な作業領域）の間のデータ転送を DMA で先読みしながら、タイルの分割サイズを調整していく必要があり、書き手には Pallas や Mosaic といったフレームワークとハードウェア特性の両方に対する深い理解が求められます。GPU 向け CUDA カーネルを LLM に書かせる研究はすでに複数ありますが、TPU 向けのエージェント的なカーネル生成はほとんど手つかずでした。

MaxKernel は、この空白を埋めるために提案された枠組みです。単発でコードを生成させるのではなく、コンパイルエラーとハードウェアプロファイルを手がかりに、複数のエージェントが反復して実装を磨き上げます。実装は GitHub で公開されています。

MaxKernelの構成

MaxKernel の中核にあるのは、役割ごとに分けられたサブエージェントの共有プールです。一つの大きなプロンプトで全部をやらせるのではなく、工程を切り分けて各エージェントに専念させる設計になっています。

• 計画エージェント: 最適化の方針を立てる
• 実装エージェント: 方針を Pallas のコードに落とし込む
• 自己デバッグ: コンパイルエラーを解釈して修正する
• テスト合成・実行: 参照実装との数値一致を検証する
• 自動チューニング: ブロックサイズやタイル寸法の探索を担う
• プロファイリング: 実行時間やメモリ帯域の利用率を計測する

これらのエージェントは、Pallas や Mosaic、XLA の公式ドキュメントやメモリ配置のガイドを検索して参照します。ただし検索対象からは専門家が手で書いたカーネルのコードが意図的に除かれており、既存の答えを写すのではなく、エージェント自身が最適化を見つけられるかを測る設計になっています。図1は、こうしたツール群と動作モードの関係を全体像として示したものです。

3つの動作モード

MaxKernel は、開発者の関与度に応じて 3 つのモードを持ちます。Human-in-the-Loop（HITL）モードは「サブエージェントを 1 つ動かしたら止まる」方式で、計画やコードの草案を人が確認し、必要なら軌道修正してから次に進みます。専門家の直感が要る難しいカーネルでは、この介入が効きます。

自律モードでは、計画から実装、コンパイル検証、テスト、自動チューニング、プロファイリングまでが一本の閉ループとしてつながります。XProf で取得した実測メトリクスが次の反復の計画段階にそのまま渡され、データに基づいて実装が更新されていきます。テストスイートは途中で凍結され、エージェントがテストを緩めて見かけの成績を上げること（報酬ハッキング）を防ぎます。成功した反復はすべてスナップショットとして残り、最終的に最も低遅延の正しい実装へ巻き戻されます。

3 つ目はグラフ探索モードです。カーネル生成を探索問題として定式化し、各ノードにソースコード、最適化方針、実測メトリクスをまとめて保持します。図4のように、オーケストレータが探索アルゴリズムを決め、各ノードの評価は独立したエージェントセッションとして開かれます。これにより、長い試行錯誤で文脈が溢れる問題を回避しつつ、成果物だけを本体の探索に返せます。探索方式は 2 種類が評価されており、5 本の独立した軌跡を並行して深く走らせる並列探索と、有望な上位候補だけを残して枝を刈るビーム探索（幅 3、深さ 3）が比較されました。

JaxBenchでの結果

評価には、LLM で広く使われる 17 種の演算子と、KernelBench から取り込んだ 33 種の融合演算子を合わせた 50 タスクのベンチマーク JaxBench が使われました。ハードウェアは TPU v6e、生成に使った言語モデルは Gemini 3.1 Pro です。比較対象は XLA コンパイラの出力で、幾何平均の速度向上と、しきい値を満たしたタスク数を表す fast_p 指標で評価しています。

差が最も大きく出たのは、速度よりもむしろ正しく動くかどうかの部分です。単発生成ではコンパイルが通った時点で 50 件中 10 件にとどまったのに対し、コンパイラのエラーを読んで修正する反復を挟むだけで、ほぼ全タスクが動作するところまで届いています。低レイヤのコード生成では、フィードバックループの有無が成否を分けると言えます。

探索方式の比較では、並列探索が最も高い速度向上を示しました。枝刈りをせずに深い試行を続けられるため、複雑な変更のデバッグに時間を割ける点が有利に働いています。一方でビーム探索は実行時間が短く、計算資源と性能のバランスを取りたい場面に向きます。

実モデルでの成果

実運用に近い 8 種類のワークロードでは、専門家が手でチューニングした Pallas 実装が基準になりました。専門家の実装が幾何平均 2.02 倍だったのに対し、並列探索の MaxKernel は 2.32 倍を記録し、8 件中 7 件で人手の実装に匹敵するか上回っています。Paged Attention では専門家の 2.41 倍に対して 6.74 倍、Sparse Attention では 2.45 倍に対して 5.03 倍と、大きく引き離した例もありました。

ただし全勝ではありません。Ragged Paged Attention では専門家の 4.65 倍に対してエージェントは 1.42 倍にとどまり、細かなハードウェア知識が要る場面では人手に及ばないことが示されています。

最新のアーキテクチャへの適用も報告されています。Multi-Head Latent Attention では遅延を 8.68%、スループットを 9.50% 改善し、Qwen3-Next の Gated DeltaNet は順伝播で 1.63 倍、学習全体では 4.70 倍まで伸びました。DeepSeek-V4 の Sparse Attention では prefill が 7.85 倍、decode が 2.36 倍です。こうした低レイヤの最適化は、Uno のようなアルゴリズム側の推論高速化と組み合わせられる性質のもので、両輪で効いてくる領域だと言えます。

まとめと今後の展望

MaxKernel が示したのは、LLM に低レイヤの最適化を任せる際、モデルの賢さそのものよりも、コンパイラとプロファイラという客観的なフィードバックを何回返せるかが効くという点です。単発生成と反復探索の間にある正答率の落差は、その裏づけになっています。

課題も残ります。論文では Gemini 3.1 Pro のみで評価しており、他のモデルでの比較検証は今後の作業とされています。探索アルゴリズムも並列探索とビーム探索の 2 種類にとどまり、進化的探索や貪欲探索、蓄積した知見を再利用する動的な知識ベースの導入が展望として挙げられました。また、探索の実行にはそれなりの計算資源と時間がかかるため、どこまで自動化に任せるかは対象カーネルの性質次第という面もあります。実装が公開されている点は、こうした検証を第三者が進めるうえで意味があります。

論文情報: "MaxKernel: Agentic Kernel Generation for TPUs"（Shangkun Wang et al., 2026） arXiv:2609.04523, CC BY 4.0

本記事の図（図1、図3、図4）とサムネイルは上記論文より引用しています（縮小・形式変換あり）。