MaxKernel: TPU向けエージェント的カーネル生成

要旨

アクセラレータ向けの高性能なカスタムカーネルの設計と作成は、深いハードウェアレベルの専門知識を必要とする複雑なタスクである。大規模言語モデル（LLM）は、リアルタイムのコンパイラフィードバックと組み合わせることで、カーネル生成のためのエージェントシステムを構築するために活用できる。本研究では、TPUカーネル開発のための3つの異なるパラダイムを実装したマルチエージェントシステムMaxKernelを提案する：（1）協調的な段階的設計を行うHuman-in-the-Loop（HITL）エージェント、（2）完全に自動化されたメトリクス／トレース駆動の最適化ループを実行する自律（Auto）エージェント、（3）Autoエージェントを拡張し、設計空間の大域探索を行うグラフベース自律探索。これら3つのパラダイムはすべて、計画、実装、自己デバッグ、テスト、ハードウェアプロファイリングを担当する専門化されたサブエージェントの共有プールを活用する。我々は、JaxBench（TPU向けの50の多様なカーネルタスクからなる包括的なベンチマークスイート）と、最先端のオープンソースモデル由来の複雑な実世界ワークロードを用いてMaxKernelを評価する。MaxKernelが、専門家による手動調整のベースラインに匹敵する高度に最適化された実装を一貫して生成し、ベンチマーク全体で顕著な性能を発揮することを実証する。本エージェントはオープンソース化されており、https://github.com/AI-Hypercomputer/accelerator-agents/tree/main/MaxKernel で入手可能である。

English

Designing and authoring high-performance custom kernels for accelerators is a complex task that requires deep hardware-level expertise. Large Language Models (LLM) can be leveraged together with real-time compiler feedback to build agentic systems for kernel generation. In this work, we present MaxKernel, a multi-agent system that implements three distinct paradigms for TPU kernel development: (1) a Human-in-the-Loop (HITL) agent for collaborative, step-by-step design; (2) an Autonomous (Auto) agent that executes a fully automated, metric/trace-driven optimization loop; and (3) a Graph-Based Autonomous Search that scales the Auto agent for global exploration of the design space. All three paradigms leverage a shared pool of specialized sub-agents to handle planning, implementation, self-debugging, testing, and hardware profiling. We evaluate MaxKernel on JaxBench, a comprehensive suite of 50 diverse kernel tasks for TPUs, alongside complex, real-world workloads from state-of-the-art open-source models. We demonstrate that MaxKernel consistently generates highly optimized implementations, matching expert hand-tuned baselines and delivering significant performance across the benchmark. Our agent is open-sourced and available https://github.com/AI-Hypercomputer/accelerator-agents/tree/main/MaxKernel.