Source-linked AI summary
MaxKernel: Agentic Kernel Generation for TPUs
Shangkun Wang, Nina Cai, Charles Hoong, Julian Walker, Gerson Kroiz, George Vanica, Deepak Patil, Andi Gavrilescu, Hassan Sipra, Sethu Sankaran
TL;DR
Developing high-performance accelerator kernels requires hardware expertise and compiler-guided iteration. MaxKernel uses specialized sub-agents across HITL, autonomous, and graph-based search workflows, achieving strong TPU benchmark and production-kernel results. Its search strategies remain bounded by the brittleness of low-level optimization and their exploration budgets.
Problem
Custom accelerator kernels require extensive hardware expertise to manage memory hierarchies, DMA pipelining, and tiling, while LLM generation benefits from real-time compiler feedback.
Method
MaxKernel combines specialized sub-agents with HITL, closed-loop Auto, and Graph-Based Autonomous Search paradigms for TPU kernel engineering.
Results
MaxKernel consistently generates optimized kernels across JAXBench and production workloads, achieving 1.58× geometric mean speedup on JAXBench and 2.32× over reference code on eight production kernels.
Takeaways & Limitations
Agentic compilation can match or exceed human-tuned kernel performance while supporting broader automated exploration of TPU optimization designs.
Takeaways & Limitations
Beam search can plateau because shallow per-depth iteration budgets may penalize multi-step optimizations requiring longer horizons to resolve hardware constraints or lowering failures.
Abstract
from arXiv · showhide
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.
1. Introduction
Custom accelerator kernels demand deep hardware expertise and iterative compiler-guided optimization. MaxKernel addresses this challenge with three orchestration paradigms and reports strong TPU performance across JAXBench and production kernels.
- Custom kernel development requires manual management of memory hierarchies, DMA pipelining, and multidimensional tiling strategies.
- Real-time compiler feedback and empirical profiling are needed because accelerator APIs, memory constraints, and low-level compiler errors limit zero-shot and few-shot LLM generation.
- MaxKernel provides HITL, Autonomous Loop, and Graph-Based Autonomous Search paradigms for TPU kernel development.
- Its modular architecture combines specialized sub-agents with compiler pipelines and profiling tools to support writing, debugging, testing, and tuning hardware-specific kernels.
- 1.58× geometric mean speedup is achieved across JAXBench workloads, while eight production kernels reach 2.32× over reference code versus 2.02× for human-written kernels.
2. System Architecture & Agent Paradigms
MaxKernel decomposes TPU kernel optimization into specialized sub-agents and supports interactive, iterative, and graph-based autonomous execution. The graph-based design expands local optimization into branching search over persistent kernel states.
- System Architecture: MaxKernel uses specialized sub-agents for planning, implementation, testing, autotuning, profiling, and hardware-specific kernel refinement.
- System Architecture: Its RAG pipeline retrieves framework documentation, memory-layout guides, and performance handbooks from a static knowledge base during execution.
- Human-in-the-Loop Agent: The HITL modality pauses after each sub-agent phase so users can inspect artifacts, provide feedback, correct the trajectory, and choose subsequent stages.
- Autonomous Loop Agent: The Auto agent performs repeated Plan Generation → Implementation → Compilation Validation → Test Execution → Autotuning → Profiling cycles without human intervention.
- Graph-Based Autonomous Search: Graph-Based Autonomous Search models kernel states as persistent graph nodes containing code, optimization plans, and empirical metrics, enabling heuristic exploration and recovery from halts.
- Graph-Based Autonomous Search: The graph framework scales beyond linear exploration by running parallel or beam-search trajectories with different iteration budgets and pruning strategies.
3. Experiments and Results
MaxKernel is evaluated on 50 diverse JaxBench workloads and complex open-source-model kernels using compilation, correctness, speedup, and fast1 metrics. Across automated search strategies, it produces correct and substantially accelerated kernels, including results competitive with or exceeding expert-written references and improvements on production workloads.
- Evaluation setup: JaxBench comprises 50 diverse workloads spanning LLM operators, fused operators, attention, linear algebra, loss functions, and other hardware-accelerated operations.
- Evaluation setup: The evaluation reports compilation rate, correctness rate, geometric mean speedup relative to standard XLA, and fast-p fraction above a specified threshold.
- JaxBench results: 10/50 Best-of-N samples compiled and were correct, with a 1.08× geometric mean speedup, whereas MaxKernel’s iterative agents improved robustness and performance.
- JaxBench results: 1.58× geometric mean speedup and 34/50 fast1 were achieved by Parallel Search with 50/50 compilation and correctness; Beam Search achieved 1.49× and 31/50 fast1 with 50/50 correctness.
- Search dynamics: Parallel Search reached 68% at p=1.0 and 24% at p=2.0 on fast-p trajectories, while Beam Search generally exceeded individual Auto runs across thresholds.
- Expert-reference comparison: On eight production kernels, Parallel Search and Beam Search achieved geometric mean speedups of 2.32 and 1.78, respectively, versus 2.02 for hand-tuned references.
4. Conclusion
MaxKernel uses context synthesis, automated diagnosis, XProf-based bottleneck identification, and structured heuristics to accelerate TPU kernel optimization. The authors report strong JaxBench performance while identifying broader search strategies as future work.
- MaxKernel combines context synthesis, automated diagnosis, XProf profiling, and structured optimization heuristics to accelerate TPU kernel hillclimbing.
- 1.58x geometric mean speedup over XLA optimizations was achieved on JaxBench.
- Evolutionary and greedy search algorithms remain directions for future exploration.The authors also identify evolving knowledge bases and stateful hybrid optimizations as possible improvements.
A. Tolerance setting for JaxBench
Table 4 reports workload-specific rtol and atol values.
- Table 4 lists workload-specific rtol values.
- Table 4 lists workload-specific atol values.
- The table organizes tolerance settings by workload.