MaxKernel Uses Multiple Agents to Generate and Optimize TPU Kernels

Summary

MaxKernel is a multi-agent system for generating high-performance custom kernels for tensor processing units (TPUs), a task that normally requires deep hardware expertise. It combines large language models with real-time compiler feedback and supports three development paradigms: a human-in-the-loop agent for collaborative, step-by-step design; an autonomous agent that runs a metric- and trace-driven optimization loop; and a graph-based autonomous search that expands the automated agent’s exploration of the design space. A shared pool of specialized sub-agents handles planning, implementation, self-debugging, testing, and hardware profiling across all three modes. The system was evaluated on JaxBench, a suite of 50 diverse TPU kernel tasks, as well as complex workloads from state-of-the-art open-source models. The authors report that MaxKernel consistently produces highly optimized implementations that match expert hand-tuned baselines and deliver significant performance across the benchmark. The agent has been open-sourced in the AI-Hypercomputer accelerator-agents repository.