AKG kernel Agent: A Multi-Agent Framework for Cross-Platform Kernel Synthesis

TL;DR

The paper tackles the challenge of generating high-performance, portable kernels for diverse AI workloads amid rapidly evolving hardware. It introduces AKG kernel agent, a modular multi-agent system that uses a hardware-agnostic Unified Sketch and a document-driven knowledge framework to synthesize, verify, and optimize kernels across multiple DSLs and backends. Through hierarchical retrieval and iterative island-model optimization, AKG achieves substantial speedups over PyTorch Eager baselines and high correctness on challenging benchmarks, including dynamic input shapes. The work demonstrates that decoupling optimization from implementation and leveraging structured documentation can significantly improve automation, portability, and performance in modern kernel engineering, with practical impact for accelerating AI system development and deployment.

Abstract

Modern AI models demand high-performance computation kernels. The growing complexity of LLMs, multimodal architectures, and recommendation systems, combined with techniques like sparsity and quantization, creates significant computational challenges. Moreover, frequent hardware updates and diverse chip architectures further complicate this landscape, requiring tailored kernel implementations for each platform. However, manual optimization cannot keep pace with these demands, creating a critical bottleneck in AI system development. Recent advances in LLM code generation capabilities have opened new possibilities for automating kernel development. In this work, we propose AKG kernel agent (AI-driven Kernel Generator), a multi-agent system that automates kernel generation, migration, and performance tuning. AKG kernel agent is designed to support multiple domain-specific languages (DSLs), including Triton, TileLang, CPP, and CUDA-C, enabling it to target different hardware backends while maintaining correctness and portability. The system's modular design allows rapid integration of new DSLs and hardware targets. When evaluated on KernelBench using Triton DSL across GPU and NPU backends, AKG kernel agent achieves an average speedup of 1.46$\times$ over PyTorch Eager baselines implementations, demonstrating its effectiveness in accelerating kernel development for modern AI workloads.

Figures (5)

• Figure 1: Overall architecture of the kernel agent. The Conductor serves as the central orchestrator, coordinating Designer, Coder, and Verifier agents in an iterative refinement loop. The system leverages expert knowledge, a database of optimized kernels, and multiple LLM backends to generate high-performance kernel code.
• Figure 2: Example of a Unified Sketch for the RMSNorm operator. The sketch demonstrates: (1) declarations with symbols and tensors; (2) nested for loops with @llm_hint("parallel", "coreidx") decorators for parallelization; (3) memory operations using alloc with semantic hints like "accumulator", "input_cache", "output_buffer"; (4) compute operations (mul, add, div, sqrt) for RMS normalization.
• Figure 3: Dynamic workflow orchestration by the Conductor. The Conductor analyzes verification outcomes and routes tasks to the appropriate agent with targeted feedback, enabling efficient error resolution.
• Figure 4: Hierarchical code retrieval pipeline. The system progressively narrows the candidate set through semantic search on computation logic, hard filtering on DSL/backend/operator type, and semantic matching on shape compatibility.
• Figure 5: Distribution of per-kernel speedups across three DSL-backend combinations on KernelBench Level 1, grouped by kernel category. Each category shows side-by-side comparisons of Triton-Ascend (blue), Triton-CUDA (green), and CPP-CPU (red). The box indicates the interquartile range (IQR), with the median shown as a horizontal line. The dashed red line indicates speedup = 1.0 (PyTorch Eager baseline). Note that Convolution kernels are only available for CPP-CPU as Triton backends do not support convolution operations in our current evaluation.