Researchers introduce MaxKernel, a multi-agent system for generating high-performance custom kernels for TPU accelerators using large language models paired with real-time compiler feedback. The system implements three paradigms: a human-in-the-loop agent for collaborative design, a fully autonomous metric-driven optimization loop, and a graph-based autonomous search for broader design-space exploration, all built on a shared pool of specialized sub-agents for planning, implementation, debugging, and hardware profiling. Evaluated on JaxBench’s 50 diverse kernel tasks and real-world workloads, MaxKernel consistently produces optimized implementations that match expert hand-tuned baselines. The system is open-sourced.
