MaxKernel: Agentic Kernel Generation for TPUs

AI summary

Overview

Research area: Automated code generation for hardware accelerators — specifically, LLM-driven multi-agent systems that write, debug, and tune custom TPU kernels (JAX/Pallas).

Technical level: Advanced. The paper assumes familiarity with accelerator programming concepts (HBM vs. SRAM/VMEM memory hierarchies, DMA pipelining, tiling strategies), compiler pipelines, and empirical profiling.

Scope: The paper presents MaxKernel, a multi-agent framework with three orchestration paradigms for TPU kernel development, and evaluates it on the 50-task JaxBench suite plus eight production kernels from state-of-the-art open-source models on TPU v6e hardware.

What This Paper Is About

Writing high-performance custom kernels for accelerators like TPUs requires deep hardware expertise — manually managing memory hierarchies, orchestrating DMA pipelines, and deriving multi-dimensional tiling strategies. Even when an LLM produces functionally correct code, reaching optimal performance requires iterative empirical profiling, autotuning, and repeated debugging against rigid compiler APIs and opaque low-level errors.

The paper's goal is to build a multi-agent system that closes the loop between LLM code generation and real-time compiler/profiler feedback, so that TPU kernels can be generated and optimized with minimal human intervention. The authors evaluate whether such a system can match or exceed expert hand-tuned Pallas baselines.

Key Contributions

1. A modular multi-agent architecture. MaxKernel decomposes kernel engineering into specialized sub-agents for planning, implementation, compilation-and-fix, test synthesis, execution, autotuning, and profiling. It integrates directly with compiler pipelines and the XProf profiling tool to feed empirical hardware feedback back to the LLM.

2. Three distinct orchestration paradigms. (a) A Human-in-the-Loop (HITL) agent that halts after each sub-agent phase and returns control to the developer ("One Agent, Then Wait"); (b) an Autonomous (Auto) agent running a fully automated, closed-loop Plan → Implement → Validate → Test → Autotune → Profile cycle with best-of-N state rollback; (c) a Graph-Based Autonomous Search that wraps the Auto agent inside a formal search graph with parallel worker expansion.

3. A RAG-based knowledge store. A retrieval-augmented generation pipeline surfaces static hardware documentation (Pallas, Mosaic, XLA framework docs, memory layout guides, performance handbooks) on demand. Hand-tuned kernel code is explicitly excluded from the retrieval corpus so the evaluation measures generalization rather than memorization of human baselines.

4. Rigorous empirical evaluation on TPUs. Evaluation on JaxBench (50 tasks) and eight production kernels on TPU v6e, using compilation rate, correctness rate, geometric mean speedup, and the fast_p metric, with XProf-backed on-device timing that filters out host-side JAX dispatch and compilation overhead.

Main Findings

• Zero-shot generation fails badly on TPUs. Best-of-N with N=100 independent zero-shot completions achieved a compilation rate of only 10/50, a correctness rate of 10/50, a geometric mean speedup of 1.08×, and a fast_1 of 6/50. The authors attribute this to the absence of real-time run context.

• The iterative Auto agent closes most of the gap. MaxKernel Auto, restricted to 5 iterations, reached a median compilation rate of 49/50 (bounds 48/50 to 50/50), a median correctness rate of 48/50 (bounds 46/50 to 49/50), a median geometric mean speedup of 1.39× (bounds 1.19× to 1.42×), and a median fast_1 of 22/50 (bounds 18/50 to 28/50). The authors note that single-trajectory runs are prone to getting stuck in suboptimal compiled states, producing relatively large variance.

• Parallel search eliminates failures and gives the best benchmark speedup. MaxKernel Parallel, running 5 concurrent Auto trajectories of 5 iterations each, achieved 50/50 compilation, 50/50 correctness, a 1.58× geometric mean speedup, and a fast_1 of 34/50.

• Beam search is competitive but slightly behind. MaxKernel Beam (beam width 3, max depth 3, 2 expansion branches per node, inner loop capped at 2 iterations) achieved 50/50 compilation, 50/50 correctness, a 1.49× geometric mean speedup, and a fast_1 of 31/50.

• Parallel search dominates the fast_p trajectory. In the threshold sweep, the parallel search curve sits above all five individual Auto run curves, reaching 68% at p=1.0 and 24% at p=2.0. Beam search also consistently sits above the single runs for most thresholds, which the authors interpret as evidence that structured guided exploration prunes suboptimal compilation paths.

• Agents beat human experts on seven of eight production kernels. Against hand-tuned Pallas references, MaxKernel Parallel achieved a 2.32× geometric mean speedup and Beam achieved 1.78×, both compared to the hand-tuned reference geometric mean of 2.02×. Notable individual results: Paged Attention, where beam search found a 6.74× kernel versus the hand-tuned 2.41×; MLA Attention, where the human baseline fell below the XLA baseline at 0.69× while Parallel reached 1.21× and Beam 1.23×; and Sparse Attention, where Parallel reached 5.03× versus 2.45× hand-tuned. The one exception was Ragged Paged Attention, where the human expert achieved 4.65× versus the agent's 1.42×.

• Periods of stability and plateau differ by search strategy. Parallel search's deep 5-iteration refinement budget raised both the lower and upper bounds of the speedup distribution, peaking above a 1.51× geometric mean. Beam search climbed during the first two depths, then plateaued at depth 3 — attributed to its shallower 2-iteration-per-node budget prematurely pruning optimizations that need longer horizons to resolve rigid memory constraints or lowering failures.

• Strong gains on production OSS workloads. On Multi-Head Latent Attention v1, MaxKernel cut latency from 3.127 ms to 2.856 ms (+8.68%) and raised throughput from 116.735 TFLOPS to 127.825 TFLOPS (+9.50%) relative to the Pallas baseline. On Qwen3-Next Gated DeltaNet, forward pass latency fell from 17.09 ms to 10.47 ms (1.63×) and the full training step from 84.51 ms to 17.99 ms (4.70×) relative to JAX. On DeepSeek-V4 Sparse Attention, speedups over JAX ranged from 2.36× for small decoding shapes (0.047 ms to 0.02 ms) up to 7.85× for user-defined prefill (27.08 ms to 3.449 ms).

• Memory-bound and robustness cases also improved. For Mamba v2 State Space Duality, keeping intermediate matrices resident in Vector Memory to avoid HBM flushes gave a 1.10× speedup (0.211 ms to 0.192 ms). For the Compressed Sparse Attention StreamIndex_topk operation, the result was 1.66× (0.415 ms to 0.250 ms). On Ragged Page Attention v3 prefill mode, MaxKernel automatically added protective ALU clamp instructions to handle left-padded inputs, preventing negative slice sizes and restoring crash-free execution with negligible overhead.

• Model and experimental setup. All LLM queries and agent interactions used Gemini 3.1 Pro; ablations with other LLMs are left as future work. Evaluation ran on TPU v6e with jnp.allclose correctness checks at atol=rtol=10⁻² for most cases, relaxed to at most 10⁻¹ for certain tasks to accommodate bf16 variance.

Methodology in Plain English

The researchers did not build a single monolithic code-writing model. Instead they built a system of specialized agents, each responsible for one stage of kernel engineering, and orchestrated them in different ways depending on how much automation and exploration they wanted.

The sub-agents. A Planning Agent produces a high-level algorithmic optimization plan from hardware specs and the reference implementation. An Implementation Agent turns that plan into Pallas code. A validation-and-fix loop repeatedly tries to compile, reads the compiler's error messages, and repairs the code before execution. A Test Synthesis Agent builds a test suite, and an Execution Agent runs it on real TPU hardware against numerical tolerances. An Autotuning Agent sweeps configuration parameters like block sizes and tile dimensions. A Profiling Agent uses XProf to extract latency, memory bandwidth utilization, and compute density, then feeds those measurements back into the next planning round.

The three orchestration modes. In the Human-in-the-Loop mode, the orchestrator runs one sub-agent and then pauses, letting a developer inspect intermediate artifacts (such as the optimization plan or a code draft) and redirect the work. In the Autonomous mode, the orchestrator chains all the sub-agents in a fixed loop with no human input: it freezes the test suite up front so implementation agents cannot alter the validation criteria, then repeatedly plans, implements, validates, tests, autotunes, and profiles. If a stage fails — for instance, too many compile-fix retries or a numerical mismatch — the pipeline short-circuits, preserves the error context, and loops back to planning with an alternative strategy. The system keeps a snapshot history of every successful iteration and rolls back to the lowest-latency valid solution at the end.

The graph search. For broader exploration, each Auto agent session becomes a node in a search graph storing the kernel source, its optimization plan, and its measured correctness and speedup. The orchestrator selects promising frontier nodes, dispatches expansion tasks to parallel workers, folds the results back into the graph, and repeats until a termination criterion (max depth or target latency) is met. Two algorithms are implemented: Parallel Search, where independent trajectories run concurrently with a larger per-worker iteration budget and no competition, and Beam Search, which keeps only the top-k candidates and aggressively prunes the rest.

The knowledge base. Rather than stuffing full hardware specifications into the context window, the system uses retrieval-augmented generation to pull relevant documentation on demand. Hand-tuned kernels were deliberately left out, so any gains must come from the agent's own reasoning and measurement.

How they evaluated it. They compared four methods — zero-shot Best-of-N, Auto, Parallel, and Beam — on JaxBench, measuring compilability, numerical correctness, geometric mean speedup relative to the XLA compiler baseline (with regressions floored at 1.0×), and the fraction of tasks that are both correct and exceed a given speedup threshold. Timing came from XProf so that host-side dispatch and compilation overhead would not contaminate the numbers.

Why This Matters

Impact on research. The paper shows that agentic scaffolding — not just a better model — is what determines whether LLM-generated accelerator code compiles and runs fast. The 10/50 zero-shot compilation rate versus 50/50 for parallel search is a dramatic illustration that real-time compiler and profiler feedback is essential for rigid low-level APIs. It also provides evidence that automated agents can match or exceed expert human kernel engineers on seven of eight production workloads, which reframes hand-tuning as a benchmark to be surpassed rather than a ceiling to be approached.

Real-world applications:

• LLM training and inference infrastructure. The evaluated kernels come from actual production architectures — Multi-Head Latent Attention, DeepSeek-V4 Sparse Attention, Qwen3-Next Gated DeltaNet, Mamba v2 SSD — where latency and throughput gains translate directly into cheaper serving and faster training steps.

• Long-context and attention-heavy serving. The largest reported wins (up to 7.85× for prefill, 6.74× on paged attention) target the attention and KV-cache machinery that dominates memory traffic in long-context inference.

• Robustness and correctness debugging. The Ragged Page Attention v3 case shows the harness can repair crash and deadlock conditions in existing kernels, not just generate new ones — useful for maintaining kernels across compiler and hardware generations.

• Reducing accelerator specialization costs. Because the agentic principles are described as broadly applicable to accelerator programming, teams working on GPUs, Trainium/NKI, and MTIA could adopt similar scaffolding without deep TPU-specific expertise.

Industry relevance. The work comes from Google and is open-sourced (the code is linked in the abstract at the AI-Hypercomputer/accelerator-agents repository under MaxKernel, released under CC BY 4.0). Reducing the requirement for scarce hardware-level kernel engineers is directly valuable to any organization running large-scale accelerator fleets, and the benchmarks used (JaxBench, plus KernelBench-derived fused operators) give the industry a shared yardstick for automated kernel generation.

Future Directions

• Other search algorithms. The authors explicitly name evolutionary search and greedy search via an "Auto Research" approach as directions they want to explore, beyond the currently implemented Parallel and Beam search.

• Dynamic and evolving knowledge bases. The paper suggests that knowledge bases built from past experiments — rather than the current static corpus of framework documentation, memory layout guides, and performance handbooks — could yield further improvements.

• Stateful hybrid optimization. The authors propose that stateful hybrid optimizations could better manage the exploration-exploitation tradeoff across the different optimization methods, since Parallel search favors depth and Beam search favors breadth.

• Alternate LLM ablations. All experiments used Gemini 3.1 Pro, and the authors state that ablations with other LLMs are left as future work — leaving open the question of how much of the result depends on the specific model.

• Unresolved variance in single-trajectory runs. The Auto agent's wide speedup bounds (1.19× to 1.42×) point to a stability problem that the graph search partially addresses but does not fully solve.

Target Audience

This paper is most useful to compiler and kernel engineers working on TPU or other accelerator backends, ML systems researchers studying agentic code generation and LLM-guided program synthesis, and infrastructure teams at organizations that train or serve large models and are weighing whether to invest in custom kernels. Readers evaluating automated approaches against human hand-tuning will find the direct head-to-head comparison on eight production kernels particularly relevant. A background in accelerator memory hierarchies or JAX/Pallas is helpful but not strictly required to follow the benchmark results and the architectural design.

Authors’ abstract

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.