MaxKernel: Agentic Kernel Generation for TPUs
Researchers open-source MaxKernel, a multi-agent LLM system that generates and optimizes TPU kernels matching expert hand-tuned baselines on JaxBench.
MaxKernel is a multi-agent system offering three paradigms for TPU kernel development: human-in-the-loop collaborative design, a fully autonomous metric/trace-driven optimization loop, and graph-based autonomous search for global exploration. All paradigms draw on a shared pool of specialized sub-agents for planning, implementation, self-debugging, testing, and hardware profiling. Evaluated on JaxBench's 50 diverse TPU kernel tasks and real-world workloads from open-source models, it consistently matches expert hand-tuned baselines. The system is open-sourced via the AI-Hypercomputer GitHub repository.