用AI自动生成TPU高性能计算核,三类智能体协同优化性能。
MaxKernel: Agentic Kernel Generation for TPUs

- 三类智能体协作:人工参与、全自动优化、图搜索全局探索
- 在50个任务上达到专家手调水平,性能显著提升
- 适合芯片优化、编译器研究者使用
为加速器设计高性能定制内核是一项复杂任务,需深入硬件知识。本文提出MaxKernel,一个面向TPU内核开发的多智能体系统,包含三种范式:(1) 人机协同的逐步设计(HITL);(2) 完全自动、基于指标与追踪驱动的优化循环(Auto);(3) 基于图的全局搜索,扩展Auto能力。三个范式共享一组专业子智能体,负责规划、实现、自调试、测试与硬件分析。我们在JaxBench(包含50个多样化内核任务)及开源先进模型的真实工作负载上评估,结果表明MaxKernel持续生成高度优化代码,性能媲美专家手调基线,并在基准测试中实现显著加速。该系统已开源,地址:https://github.com/AI-Hypercomputer/accelerator-agents/tree/main/MaxKernel。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。