MaxKernel: Agentic Kernel Generation for TPUs
Shangkun Wang, Nina Cai, Charles Hoong, Julian Walker, Gerson Kroiz, George Vanica, Deepak Patil, Andi Gavrilescu, et al.
cs.AI · cs.PF · cs.PL
Quick verdict
MaxKernel agents auto-write TPU Pallas kernels, claiming 1.58× geomean on JaxBench; strong pipeline but uncosted, single-LLM, and a Table 2/text mismatch.
Summary
MaxKernel is a multi-agent system from Google that generates and optimizes custom TPU kernels written in JAX/Pallas. It offers three orchestration paradigms sharing a common pool of sub-agents (planning, implementation, compile-fix, testing, autotuning, XProf-based profiling): a Human-in-the-Loop mode, a fully autonomous hill-climbing loop, and a graph-based search (parallel or beam) that scales the autonomous loop. A RAG store provides static framework documentation while deliberately excluding hand-tuned kernels. Evaluated on TPU v6e with Gemini 3.1 Pro across JaxBench's 50 tasks, the autonomous agent reaches a 1.39× median geometric-mean speedup over XLA baselines (range 1.19–1.42× across five runs); parallel search reaches 1.58× with 50/50 compile and correctness; beam search reaches 1.49×. On eight production workloads, parallel search attains 2.32× geomean versus 2.02× for expert hand-tuned references, though it loses on Flash Attention and Ragged Paged Attention; the text's claim of winning 'seven of eight' conflicts with Table 2. Real-model kernels (MLA, Qwen3-Next GDN, DeepSeek-V4 sparse attention) show 1.10–7.85× gains. Limitations: no statistical testing or compute-cost accounting, speedups floored at 1.0×, a single closed-source LLM, and proprietary-hardware dependence limiting independent reproduction despite open-sourced code.
Review complete
Assessment at a glance
5
risk flags identified
● high 1
● medium 3
● low 1
5 strengths · 5 weaknesses · 5 author questions
AI text signal = Binoculars zero-shot detector score over 20 sampled passages (Fast-DetectGPT co-reported) — absence does not prove human authorship, and a flag does not imply misconduct.
Levels reflect verification complexity — not a quality score. Evolve never rates your paper.
Key claims under review
- A multi-agent loop with compiler feedback and XProf profiling lifts JaxBench geometric-mean speedup over the XLA baseline from 1.08× (Best-of-100 zero-shot) to 1.58× (parallel search) with 50/50 compile and correctness on TPU v6e.
- On eight production kernels, MaxKernel parallel search achieves a 2.32× geometric-mean speedup versus 2.02× for expert hand-tuned Pallas references.
- Graph-based search (beam) avoids the local optima that trap single-pass Auto runs, though it underperforms parallel search at the tested budgets (1.49× vs 1.58×).
- MaxKernel generalizes to SOTA open-source model kernels, e.g., 4.70× faster Qwen3-Next GDN training step and up to 7.85× on DeepSeek-V4 sparse-attention prefill versus JAX baselines.
- With hand-tuned kernels excluded from the retrieval corpus, agents can discover TPU-specific optimizations independently of human-engineered baselines.
Strengths
Anti-reward-hacking evaluation design
The test suite is synthesized and frozen before the optimization loop so implementation agents cannot alter validation criteria (§2.4); compile rate and correctness rate are reported separately (Table 1); latency uses XProf on-device timing that explicitly filters host-side JAX dispatch/compile overhead (§3.1).
Large, variance-disclosed gains over zero-shot
Table 1: Best-of-100 zero-shot compiles only 10/50 at 1.08× geomean, while Auto reaches 1.39× median [1.19–1.42] with 49/50 compile / 48/50 correct, Parallel 1.58× with 50/50 and fast1 34/50, Beam 1.49×.
Human-expert comparison with losing cases shown
Real-world transfer beyond the benchmark
Reproducibility scaffolding
Weaknesses
Internal contradiction on the headline human-comparison claim
No compute or budget accounting
Thin statistical treatment
Favorably biased aggregate metrics
Single model, single hardware, figure-only evidence
Full Review Report
$1.99 · one-timeUnlock the complete preliminary analysis Evolve produced for this paper.
Report preview
Scope & Claims
MaxKernel (Google; arXiv 2609.04523) is a systems/method paper presenting a multi-agent LLM framework that generates and optimizes TPU kernels in JAX/Pallas. Three orchestration paradigms share one pool of sub-agents (planning, implementation with a compile-fix loop, test synthesis/execution, autotuning, XProf profiling): a Human-in-the-Loop router following the 'One Agent, Then Wait' pattern (§2.3), a fully autonomous closed-loop Auto agent (§2.4) with a frozen pre-generated test suite to block reward hacking, best-of-N state rollback, and orchestrator-driven path management, and a Graph-Based Search (§2.5) wrapping Auto workers in parallel or beam search over a persistent SearchGraph. A RAG store (§2.2) supplies only static documentation (Pallas, Mosaic, XLA) and deliberately excludes hand-tuned kernel code, so agents must discover optimizations independently. Central claims: agentic loops with compiler/profiler feedback beat zero-shot generation by a wide margin; parallel search reaches 1.58× geometric-mean speedup on JaxBench over XLA (Table 1); on eight production kernels the agents (2.32×) beat expert hand-tuned references (2.02×) (Table 2); and the system transfers to SOTA OSS model kernels with gains up to 7.85× (Table 3).
The remaining sections cover Methodology Assessment, Statistical Evidence, Results & Reproducibility, Limitations & Risks, and Verdict & Recommendations.
- All 5 strengths & 5 weaknesses, analyzed in detail
- 5 risk flags with severity, section & reviewer questions
- Methods & statistical evidence assessment
- 5 questions a peer reviewer would ask
- Full six-dimension evaluation report & recommendations
One-time purchase · this paper only
One-time purchase, this paper only. AI-generated preliminary analysis — not peer review.
Frequently asked questions
What speedup does MaxKernel achieve on TPU kernels?
On the 50-task JaxBench suite running on TPU v6e, MaxKernel's parallel-search variant reaches a 1.58× geometric-mean speedup over the standard XLA-compiled baseline with 50/50 compile and correctness rates; the single-trajectory Auto agent reaches a median 1.39× and beam search 1.49×. On eight production kernels it reports 2.32× geometric mean versus 2.02× for expert hand-tuned Pallas code.
Is MaxKernel's code open source?
Yes — the agent is open-sourced in the AI-Hypercomputer 'accelerator-agents' GitHub repository under MaxKernel. However, reproducing the reported numbers requires TPU v6e hardware and the Gemini 3.1 Pro model used in all experiments, so independent verification is not trivial.
Can AI agents replace human TPU kernel engineers?
The paper shows agents matching or exceeding expert hand-tuned code on most, but not all, of eight production workloads: parallel search clearly beats the human references on 5 of 8 (with one tie) and loses notably on Ragged Paged Attention (1.42× vs the human 4.65×). The abstract's framing is 'matching' expert baselines; the text's stronger 'seven of eight' claim conflicts with the paper's own Table 2.
Which LLM does MaxKernel use?
All experiments in the paper use Gemini 3.1 Pro. Ablations with other LLMs are explicitly deferred to future work, so the system's sensitivity to the choice of base model is currently untested.
About this review
This is an AI-generated preliminary review by Evolve.Pub, produced from the openly available text of the paper. It is not peer review and does not represent the authors' views. The canonical record of the paper is at the publisher link above.
- Figures 1–6 are available only via captions; figure-internal data (e.g., the fast-p curves in Figure 5 and the search-progression plateaus in Figure 6) cannot be verified from the extracted text, so judgments about them rely on the prose.
- Table 2's column alignment relies on extraction order; its values are cross-corroborated by prose (MLA 0.69×, Paged 6.74× vs 2.41×), so the 'seven of eight' discrepancy is judged an inconsistency in the original text, not an extraction loss.
Go deeper on this paper
Wrote a paper? Stress-test it before reviewers do.
Upload your manuscript for the same structured scrutiny — reasoning, statistics, evidence, figures, and reproducibility.