Creating high‑performance custom kernels for TPUs traditionally demands deep hardware expertise. By coupling large language models (LLM) with real‑time compiler feedback, we built MaxKernel, a multi‑agent system. It supports three development paradigms: a Human‑in‑the‑Loop agent that collaborates step by step; an Autonomous agent that runs a metric‑ and trace‑driven optimization loop; and a Graph‑Based Autonomous Search that scales the Auto agent for global design‑space exploration. All paradigms share a pool of specialized sub‑agents handling planning, implementation, self‑debugging, testing, and hardware profiling. Evaluation on JaxBench—a suite of 50 diverse TPU kernel tasks—and on complex real‑world workloads from state‑of‑the‑art open‑source models shows MaxKernel consistently produces implementations on par with expert‑tuned baselines and delivers substantial speedups. The agent is open‑sourced at https://github.com/AI-Hypercomputer/accelerator-agents/tree/main/MaxKernel.
Review