EGG: An Expert-Guided Agent Framework for Kernel Generation

High-performance GPU kernels are critical for reducing the exponentially growing computational costs of large language models (LLMs), but their development heavily relies on manual tuning by domain experts. While recent advances in LLM-based approaches show promise for automating kernel generation, they still struggle to achieve both correctness and high performance. This limitation primarily arises from the lack of domain-specific optimization guidance, hindering effective exploration of the optimization space. We propose $\textbf{EGG}$, an $\underline{E}$xpert-$\underline{G}$uided Agent Framework for Kernel $\underline{G}$eneration, which incorporates expert optimization principles to guide LLMs’ decisions. Inspired by expert workflows, we decompose kernel generation into two hierarchical stages: 1) algorithmic structure design, which establishes a high-quality computational structure foundation; 2) hardware-specific tuning, which performs targeted adjustments through parallel mapping, tensor tiling, and memory optimization. This staged decomposition defines explicit optimization objectives, structuring the design space to achieve progressive refinements. To this end, a stage-aware multi-agent collaboration mechanism is designed for inter and intra-stage context management, ensuring stable optimization trajectories. Experiments on KernelBench and real-world workloads show that EGG achieves a $2.13\times$ average speedup over PyTorch, outperforming existing agent-based and RL-based approaches.

Running modern AI models efficiently requires highly optimized GPU programs, but writing and tuning these programs usually takes substantial expert engineering time. This paper presents EGG, a system that uses AI agents to help create faster GPU programs automatically. EGG guides the agents with a step-by-step workflow inspired by how human experts optimize code: first choosing a better overall computation strategy, then tuning how the work is mapped and executed on the target hardware. The system also uses separate helper agents to write code, inspect performance, and fix errors, making the optimization process more stable. Across standard benchmark tasks and real workloads, EGG consistently generates correct programs and achieves substantial speedups over common PyTorch implementations on average. The broader goal is to reduce the manual effort needed to build efficient software for machine learning systems as models and hardware continue to evolve.