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
- Vendors
From the original article
1 min read at huggingface.co
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.
By Shangkun Wang, Nina Cai, Charles Hoong · huggingface.co
A short excerpt; the article and its images belong to the publisher. Summary, key facts and assessment above are AI-written from the article and can be wrong.