MaxKernel: Agentic Kernel Generation for TPUs
Mirrored from Hugging Face Daily Papers for archival readability. Support the source by reading on the original site.
Abstract
MaxKernel is a multi-agent system that automates TPU kernel development through collaborative, autonomous, and graph-based search paradigms, achieving expert-level performance on diverse benchmarks.
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.
Get this paper in your agent:
hf papers read 2609.04523 curl -LsSf https://hf.co/cli/install.sh | bash Models citing this paper
No model linking this paper
Datasets citing this paper
No dataset linking this paper
Spaces citing this paper
No Space linking this paper
Discussion (0)
Sign in to join the discussion. Free account, 30 seconds — email code or GitHub.
Sign in →No comments yet. Sign in and be the first to say something.