Google's MaxKernel: agents match expert-tuned TPU kernels

Google's cloud team has released MaxKernel, a multi-agent system for writing Pallas kernels for TPUs. The accompanying paper reports the agents hit a geometric-mean speedup of 1.58x over XLA baselines across JaxBench, a suite of 50 diverse kernel tasks for TPUs.

The system runs in three modes: a human-in-the-loop design partner, a fully autonomous agent that iterates on metrics and hardware traces, and a graph-based search that explores the design space more widely. All three share a pool of sub-agents that split up planning, coding, debugging, validation and profiling.

The most concrete real-world result comes from a Deepseek MLA kernel on a v5p TPU. According to a Google Developer forum post from Group Product Manager Deepak Patil and Senior ML software developer Nina Cai, MaxKernel delivered an "8.7% speedup (latency dropped from 3.12ms to 2.856 ms)" against a human-optimized baseline, with throughput a "9% increase (jumped from 116.73 TFLOPS to 127.82 TFLOPS)."

The code is available in Google's Accelerator Agents OSS GitHub repository. The paper's own framing is careful: it says the system produces code "comparable to expert hand-crafted solutions" rather than beating them outright. Neither the paper abstract nor the blog post breaks out per-task wins and losses on the 50 JaxBench kernels, so how uniform the 1.58x geo-mean is across the suite is not shown.