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.
- Three paradigms: human-in-the-loop design, autonomous optimization, and graph-based search
- Shared sub-agent pool handles planning, self-debugging, testing, and hardware profiling
- Matches expert hand-tuned baselines across 50 JaxBench TPU kernel tasks
- Uses LLMs combined with real-time compiler feedback for kernel generation
- Code released on GitHub under AI-Hypercomputer accelerator-agents
Full article163 words · extracted from huggingface.co · click to collapse
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.
Text extracted automatically; images, tables and formatting may be missing. Original: https://huggingface.co/papers/2609.04523