MaxKernel：面向TPU的智能體驅動內核生成

摘要

為加速器設計與編寫高效能自訂核心，是一項需要深厚硬體層級專業知識的複雜任務。大型語言模型（LLM）可結合即時編譯器回饋，用以建構代理式系統來生成核心。在本研究中，我們提出 MaxKernel——一個多代理系統，實作三種不同的 TPU 核心開發範式：（1）人在迴圈（HITL）代理，用於協作式逐步設計；（2）自主（Auto）代理，執行全自動、由度量與追蹤資料驅動的優化迴圈；（3）基於圖的自主搜尋，將 Auto 代理擴展至設計空間的全域探索。這三種範式皆利用共享的專業子代理池，以處理規劃、實作、自我除錯、測試與硬體效能剖析。我們在 JaxBench——一個涵蓋 50 種多樣化 TPU 核心任務的全面性套件——以及來自尖端開源模型的複雜真實世界工作負載上，對 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.