🤖 AI Summary
This work addresses the challenge of automatically generating high-performance GPU kernels that simultaneously achieve correctness and efficiency, a task often hindered by the absence of expert guidance. To this end, the authors propose EGG, a novel framework that decomposes kernel synthesis into two hierarchical stages: algorithmic structure design and hardware-specific tuning. EGG employs a stage-aware multi-agent collaboration architecture, where large language models make progressive decisions guided by expert-derived optimization principles. By integrating techniques such as parallel mapping, tensor tiling, and memory optimization, the framework enables co-optimization of algorithms and hardware characteristics. Experimental results demonstrate that EGG achieves an average performance 2.13× that of PyTorch on both KernelBench and real-world workloads, significantly outperforming existing agent-based and reinforcement learning approaches.
📝 Abstract
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 EGG, an Expert-Guided Agent Framework for Kernel Generation, 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 refinement. 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.13x average speedup over PyTorch, outperforming existing agent-based and RL-based approaches.