EGG: 전문가 지향 에이전트 프레임워크를 활용한 커널 생성

EGG: An Expert-Guided Agent Framework for Kernel Generation

고성능 GPU 커널은 대규모 언어 모델(LLM)의 기하급수적으로 증가하는 계산 비용을 줄이는 데 매우 중요하지만, 그 개발은 주로 해당 분야 전문가의 수동적인 튜닝에 크게 의존합니다. 최근 LLM 기반 접근 방식에서 자동화된 커널 생성 가능성이 제시되었지만, 여전히 정확성과 높은 성능을 동시에 달성하는 데 어려움을 겪고 있습니다. 이러한 한계는 주로 도메인 특유의 최적화 지침 부족으로 인해 발생하며, 이는 효과적인 최적화 공간 탐색을 방해합니다. 본 연구에서는 전문가의 최적화 원칙을 LLM의 의사 결정에 반영하여 커널 생성을 위한 전문가 지향 에이전트 프레임워크인 EGG를 제안합니다. 전문가 워크플로우에서 영감을 얻어, 커널 생성 과정을 두 가지 계층 구조로 분해했습니다. 첫째, 고품질의 계산 구조 토대를 구축하는 알고리즘 구조 설계 단계이며, 둘째는 병렬 매핑, 텐서 타일링 및 메모리 최적화를 통해 목표 지향적인 조정을 수행하는 하드웨어 특화 튜닝 단계입니다. 이러한 계층 구조 분해를 통해 명확한 최적화 목표를 정의하고, 점진적인 개선을 달성할 수 있도록 설계 공간을 구조화합니다. 또한, 각 단계 간 및 내부의 컨텍스트 관리를 위해 스테이지 인지 멀티 에이전트 협업 메커니즘을 설계하여 안정적인 최적화 경로를 보장합니다. KernelBench 및 실제 워크로드에 대한 실험 결과, EGG는 PyTorch보다 평균 2.13배 빠른 속도를 달성했으며, 기존의 에이전트 기반 및 강화 학습 기반 접근 방식보다 뛰어난 성능을 보여주었습니다.

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.