MaxKernel: Agentic Kernel Generation for TPUs

• URL: http://arxiv.org/abs/2609.04523v1
• Date: Thu, 03 Sep 2026 22:22:37 GMT
• Title: MaxKernel: Agentic Kernel Generation for TPUs
• Abstract summary: MaxKernel is a multi-agent system that implements three distinct paradigms for TPU kernel development.<n>All three paradigms leverage a shared pool of specialized sub-agents to handle planning, implementation, self-ging, testing, and hardware profiling.<n>We demonstrate that Max Kernel consistently generates highly optimized implementations, matching expert hand-tuned baselines.
• Score: 0.1719024873384565
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• KForge: LLM-Driven Cross-Platform Kernel Generation for AI Accelerators [5.703659894138959]
  Production inference increasingly targets a heterogeneous mix of accelerators.<n>This creates a systems challenge: each pipeline now requires high-performance kernels across a growing set of hardware backends and programming models.<n>We present KForge, a cross-platform framework built around an iterative refinement loop driven by two collaborating LLM-based agents.
  arXiv  Detail & Related papers  (2026-06-01T23:48:55Z)
• AgentKernelArena: Generalization-Aware Benchmarking of GPU Kernel Optimization Agents [30.03012293068346]
  AgentArena is an open-source benchmark for measuring AI coding agents on GPU kernel optimization.<n>We find near-perfect compilation and high correctness rates on most task categories.<n>Our unseen-configuration evaluation shows that HIP-to-HIP and Triton-to-Triton optimizations largely transfer to unseen input shapes.
  arXiv  Detail & Related papers  (2026-05-16T05:25:11Z)
• KernelCraft: Benchmarking for Agentic Close-to-Metal Kernel Generation on Emerging Hardware [25.808580418841718]
  New AI accelerators often require developers to manually craft low-level kernels.<n>This prevents emerging hardware platforms from reaching the market efficiently.<n>We present KernelCraft: the first benchmark to evaluate an agent's ability to generate and optimize low-level kernels for customized accelerators.
  arXiv  Detail & Related papers  (2026-02-10T14:52:02Z)
• AKG kernel Agent: A Multi-Agent Framework for Cross-Platform Kernel Synthesis [13.239454996851771]
  Modern AI models demand high-performance computation kernels.<n>Akg kernel agent (AI-driven Kernel Generator) is designed to support multiple domain-specific languages.<n>System's modular design allows rapid integration of backend DSLs and hardware targets.<n>System achieves an average speedup of 1.46$times over PyTorch Eager baselines.
  arXiv  Detail & Related papers  (2025-12-29T12:42:05Z)
• cuPilot: A Strategy-Coordinated Multi-agent Framework for CUDA Kernel Evolution [15.701861287574296]
  cuPilot is a strategy-coordinated multi-agent framework that introduces strategy as an intermediate semantic representation for kernel evolution.<n>On the GEMM tasks, cuPilot showcases sophisticated optimizations and achieves high utilization of critical hardware units.
  arXiv  Detail & Related papers  (2025-12-18T12:34:00Z)
• xLLM Technical Report [57.13120905321185]
  We introduce xLLM, an intelligent and efficient Large Language Model (LLM) inference framework.<n>xLLM builds a novel decoupled service-engine architecture.<n>xLLM-Engine co-optimizes system and algorithm designs to fully saturate computing resources.
  arXiv  Detail & Related papers  (2025-10-16T13:53:47Z)
• MATRIX: Multimodal Agent Tuning for Robust Tool-Use Reasoning [65.200259961515]
  We develop a vision-centric agent tuning framework that automatically synthesizes multimodal trajectories.<n>We also introduce Pref-X, a set of 11K automatically generated preference pairs.<n>Across three benchmarks, Agent-X, GTA, and GAIA, MATRIX consistently surpasses both open- and closed-source VLMs.
  arXiv  Detail & Related papers  (2025-10-09T17:59:54Z)
• xLAM: A Family of Large Action Models to Empower AI Agent Systems [111.5719694445345]
  We release xLAM, a series of large action models designed for AI agent tasks. xLAM consistently delivers exceptional performance across multiple agent ability benchmarks.
  arXiv  Detail & Related papers  (2024-09-05T03:22:22Z)
• CRAB: Cross-environment Agent Benchmark for Multimodal Language Model Agents [49.68117560675367]
  Crab is the first benchmark framework designed to support cross-environment tasks.<n>Our framework supports multiple devices and can be easily extended to any environment with a Python interface.<n>The experimental results demonstrate that the single agent with GPT-4o achieves the best completion ratio of 38.01%.
  arXiv  Detail & Related papers  (2024-07-01T17:55:04Z)
• Harnessing Deep Learning and HPC Kernels via High-Level Loop and Tensor Abstractions on CPU Architectures [67.47328776279204]
  This work introduces a framework to develop efficient, portable Deep Learning and High Performance Computing kernels. We decompose the kernel development in two steps: 1) Expressing the computational core using Processing Primitives (TPPs) and 2) Expressing the logical loops around TPPs in a high-level, declarative fashion. We demonstrate the efficacy of our approach using standalone kernels and end-to-end workloads that outperform state-of-the-art implementations on diverse CPU platforms.
  arXiv  Detail & Related papers  (2023-04-25T05:04:44Z)
• PolyScientist: Automatic Loop Transformations Combined with Microkernels for Optimization of Deep Learning Primitives [55.79741270235602]
  We develop a hybrid solution to the development of deep learning kernels. We use the advanced polyhedral technology to automatically tune the outer loops for performance.
  arXiv  Detail & Related papers  (2020-02-06T08:02:34Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.