Search papers, labs, and topics across Lattice.
MaxKernel automates custom TPU kernel development through a multi-agent framework featuring human-in-the-loop collaboration, trace-driven autonomous optimization, and graph-based design space exploration. By integrating closed-loop compiler feedback with dedicated sub-agents for planning, debugging, and hardware profiling, the system targets the steep expertise barrier required for low-level accelerator programming. Across the 50-task JaxBench suite and production model workloads, the framework autonomously matches the performance of expert hand-tuned baselines.
Low-level accelerator optimization is yielding to agentic search: closed-loop LLMs driven by hardware profiling can now match human experts at generating high-performance TPU kernels.
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.
