MaxKernel: Agentic Kernel Generation for TPUs

📄 arXiv: 2609.04523v1 📥 PDF

作者: Shangkun Wang, Nina Cai, Charles Hoong, Julian Walker, Gerson Kroiz, George Vanica, Deepak Patil, Andi Gavrilescu, Hassan Sipra, Sethu Sankaran

分类: cs.AI, cs.PF, cs.PL

发布日期: 2026-09-03

备注: 14 pages, 6 figures, 4 tables

🔗 代码/项目: GITHUB

💡 一句话要点

提出MaxKernel以简化TPU内核生成过程

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: TPU内核生成 多代理系统 大型语言模型 自动化优化 高性能计算 深度学习加速 实时编译器反馈

📋 核心要点

  1. 现有方法在设计高性能TPU内核时需要深厚的硬件知识，导致开发过程复杂且效率低下。
  2. MaxKernel通过结合人机协作、全自动优化和图搜索三种范式，简化了TPU内核的生成过程。
  3. 实验结果显示，MaxKernel生成的内核实现性能与专家手动调优相当，且在多个基准测试中表现出显著提升。

📝 摘要（中文）

设计和编写高性能自定义内核对于加速器来说是一项复杂的任务，需要深入的硬件知识。本文提出MaxKernel，一个多代理系统，结合大型语言模型（LLM）和实时编译器反馈，支持TPU内核开发的三种不同范式：人机协作代理、全自动优化代理和基于图的自主搜索。MaxKernel通过共享的专业子代理池处理规划、实现、自我调试、测试和硬件分析。我们在JaxBench上评估了MaxKernel，结果表明其生成的实现与专家手动调优的基线相匹配，并在基准测试中显著提升性能。

🔬 方法详解

问题定义：本文旨在解决高性能TPU内核生成过程中的复杂性和对硬件知识的高要求。现有方法往往依赖于专家经验，导致效率低下和可扩展性差。

核心思路：MaxKernel的核心思路是利用大型语言模型和实时编译器反馈，构建一个多代理系统，支持不同的内核开发范式，以实现更高效的内核生成。

技术框架：MaxKernel的整体架构包括三个主要模块：人机协作代理（HITL）、全自动优化代理（Auto）和基于图的自主搜索。HITL代理支持逐步设计，Auto代理执行自动化优化循环，而图搜索则扩展了Auto代理的设计空间探索能力。

关键创新：MaxKernel的创新在于其多代理系统的设计，能够灵活切换不同的开发范式，并通过共享的子代理池实现规划、实现和调试等功能。这种设计与传统的单一方法相比，显著提高了内核生成的效率和灵活性。

关键设计：在实现中，MaxKernel使用了一系列专门的子代理，负责不同的任务，如规划、实现和测试。此外，系统还集成了实时编译器反馈，以优化内核性能和调试过程。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在JaxBench的评估中，MaxKernel生成的内核实现与专家手动调优的基线相匹配，并在多个复杂的真实工作负载中表现出显著的性能提升，证明了其在TPU内核生成中的有效性和优越性。

🎯 应用场景

MaxKernel的研究成果在高性能计算、深度学习模型加速和其他需要高效内核生成的领域具有广泛的应用潜力。通过简化内核开发过程，MaxKernel可以帮助研究人员和工程师更快速地实现高效的计算解决方案，推动相关技术的进步。

📄 摘要（原文）

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.