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

MaxKernel uses a multi-agent system with real-time compiler feedback to automate high-performance TPU kernel generation.

How can agentic systems automate the generation and optimization of high-performance kernels for TPU accelerators?

Writing custom kernels for hardware accelerators like TPUs requires deep expertise in memory management and tiling, making it a bottleneck for scaling deep learning models. MaxKernel automates this by using specialized sub-agents that iteratively plan, implement, test, and profile code using real-time compiler and hardware feedback. On the JaxBench suite, this system achieves a 1.58× geometric mean speedup over standard compiler baselines and outperforms human-authored kernels on production workloads.

Paper Primer

MaxKernel functions as a modular multi-agent framework that decomposes kernel engineering into specialized tasks: planning, implementation, verification, autotuning, and profiling. The system is like a software engineering team — one agent drafts the code, another tests it for correctness, and a third profiles the hardware to identify bottlenecks, feeding that data back into the next iteration.

The framework supports three orchestration paradigms: a Human-in-the-Loop (HITL) mode for guided design, an Autonomous Loop for iterative hill-climbing, and a Graph-Based Autonomous Search that explores the design space in parallel. By modeling kernel generation as a formal search graph, the system can systematically navigate performance trade-offs and recover from local optima that trap simpler, linear optimization approaches.

MaxKernel significantly outperforms standard XLA compiler baselines on diverse TPU workloads.

Geometric mean speedup across 50 JaxBench tasks.

MaxKernel generates kernels superior to human-authored implementations for production-grade models.

Comparison against human-written Pallas kernels on eight production workloads.

Why is real-time compiler feedback necessary for this agentic approach?

Accelerator APIs are rigid and memory constraints are strict, leading to opaque low-level compiler errors that standard zero-shot LLM generation cannot resolve without empirical validation.

How does the Graph-Based Autonomous Search differ from the standard Autonomous Loop?

While the Autonomous Loop performs linear, iterative refinement, the Graph-Based Search treats the design space as a formal graph, allowing the system to branch into multiple parallel trajectories and prune suboptimal paths to avoid local optima.

The Kernel Generation Bottleneck

MaxKernel automates kernel creation by replacing manual expert tuning with LLM‑driven agents.

Designing high‑performance kernels for accelerators demands deep hardware expertise, making the process slow, brittle, and inaccessible to most developers.

Agentic Kernel Generation lets a Large Language Model iteratively refine kernel code by receiving immediate compiler feedback, turning code synthesis into a closed‑loop optimization process.

**Figure 1.** User works agent to co-develop kernels

Agentic automation replaces hand‑crafted kernel tuning with LLM‑driven iterative refinement, dramatically lowering expertise barriers.

MaxKernel System Architecture

MaxKernel orchestrates specialized agents to iteratively refine kernels via closed‑loop feedback.

MaxKernel decomposes kernel optimization into a set of focused sub‑agents and supports three execution paradigms: an interactive Human‑in‑the‑Loop mode, a fully automated Autonomous Loop, and a scalable Graph‑Based Autonomous Search.

The framework is a modular pipeline where each sub‑agent owns a single responsibility—planning, code generation, validation, testing, autotuning, or profiling—so the whole system can be orchestrated, inspected, and swapped piecewise.

The loop repeatedly feeds profiling feedback into the planner, turning a static optimization problem into a data‑driven hill‑climbing process that converges on the fastest correct kernel.

Iteration 1: Planner proposes a tiling of $1\times2$ (row‑wise) and generates Pallas code.

Compilation succeeds; test runner verifies numerical equality.

Profiling reports latency $= 12\,$ms and compute density $= 0.6$.

Planner incorporates the latency into a cost model and switches to a $2\times1$ (column‑wise) tiling for Iteration 2.

Iteration 2 compiles, passes tests, and yields latency $= 9\,$ms (improvement).

Best‑of‑N selects Iteration 2 as the final kernel because it has the lowest latency.

The loop’s ability to change tiling based on measured latency demonstrates how empirical feedback steers the search away from static heuristics.

How does the Auto Agent Loop differ from a naïve “run‑once” code generator?

The naïve generator emits code once and stops; the Auto Loop keeps the planner in the feedback loop, using real profiling numbers to iteratively improve the kernel until no further latency gains are observed.

Graph‑Based Autonomous Search treats each Auto Agent run as a node in a search graph, allowing many candidates to be explored in parallel (Parallel Search) or to focus resources on the most promising candidates (Beam Search).

Beam Search selects the top‑2 nodes ($+12\%$ and $+8\%$) for expansion with a budget of 3 iterations each.

Parallel Search expands all three nodes with a budget of 6 iterations each.

After expansion, Beam Search discovers a $+15\%$ node; Parallel Search discovers a $+13\%$ node.

Both strategies insert their new nodes into the graph; the frontier now contains $+15\%$, $+13\%$, and the previous $+12\%$ node.

Beam Search can jump quickly to a higher‑performing region, while Parallel Search provides deeper refinement of each candidate.

Why would one choose Beam Search over Parallel Search when both explore the same space?

Beam Search concentrates compute on the most promising candidates, reducing total runtime when the search budget is limited; Parallel Search spreads compute more evenly, which can be advantageous for very large or highly non‑convex design spaces.

**Figure 1 |** A schematic of the MaxKernel framework outlining the various modalities of operation and tools.

**Figure 3.** The closed-loop optimization workflow of the Auto Agent.

**Figure 4 |** Design of the Graph based autonomous search. The orchestrator determines the search algorithm, where each evaluation of the search opens up a new autonomous agent session, which then retrieves and returns artifacts to the main search. A database with the saved nodes and graph is also performed at the end of each search iteration.

Empirical Evaluation on JaxBench

MaxKernel delivers up to 1.58× geometric‑mean speedup over XLA while achieving full compilation and correctness.

We evaluate MaxKernel on JaxBench, a curated suite of 50 diverse hardware‑accelerated workloads that mixes 17 common operators from LLM architectures with 33 fused operators from KernelBench.

JaxBench is a benchmark collection that stresses a range of kernels—attention, dense and sparse linear algebra, and fused loss functions—so we can measure how well MaxKernel automates end‑to‑end kernel generation for real‑world workloads.

Best‑of‑N generates $N=100$ independent zero‑shot completions from the LLM and picks the fastest one that also passes the correctness check.

Why does Best‑of‑N perform so poorly despite sampling many candidates?

Without iterative feedback, most sampled kernels lack the low‑level memory‑layout fixes required to compile on TPU v6e; the few that do compile often miss subtle numeric tolerances, leading to a low correctness rate.

MaxKernel Parallel attains a 1.58× geometric‑mean speedup over the XLA baseline.

Table 1 shows MK Parallel achieving 1.58× speedup while maintaining 50/50 compilation and correctness rates.

MK Auto improves over Best‑of‑N, reaching a median speedup of 1.39× but shows larger variance because a single trajectory can get stuck in suboptimal compiled states.

MK Beam, with a beam width of 3 and depth 3, achieves a competitive 1.49× speedup and a fast1 score of 31/50, demonstrating that guided breadth‑first exploration can approach the parallel strategy’s performance.

The table compares four methods: Best-of-N, MK Auto, MK Parallel, and MK Beam across four metrics: Compilation Rate, Correctness Rate, Geometric Mean Speedup, and `fast_1`.

**Figure 5.** Comparison of the `fast_p` metric trajectories across different methods on JaxBench for varying speedup thresholds p. Parallel nodes refer to the autonomous search trajectories of individual nodes.

**Figure 6.** Progression of geometric mean speedup. Evaluated candidate nodes for Beam Search (blue) and Parallel Search (orange) are aggregated across the 50 benchmark problems. Nodes are grouped sequentially by their chronological evaluation step—denoted by Depth (D) and Iteration (I)—and sorted internally from lowest to highest performance to form the ascending segments.

**Table 3.** Performance Improvements of MaxKernel on SOTA Architectures

MaxKernel’s parallel search delivers the strongest overall performance, surpassing hand‑tuned baselines by up to 1.58× speedup and achieving full compilation and correctness on the JaxBench suite.

Conclusion and Implementation Details

Conclusion and JaxBench tolerance details.

MaxKernel shows that agentic systems can manage the intricacies of low‑level accelerator code, using a harness that synthesizes context, diagnoses failures, pinpoints bottlenecks with XProf, and applies structured heuristics to speed up hill‑climbing on TPUs. The system achieves a geometric‑mean speedup of 1.58× on JaxBench beyond XLA’s own optimizations. Looking ahead, we plan to add evolutionary and greedy search methods, maintain dynamic knowledge bases, and develop stateful hybrid optimizations to better balance exploration and exploitation.

This table lists 50 distinct workloads and their corresponding relative tolerance (rtol) and absolute tolerance (atol) values used for verification.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers