EGG: An Expert-Guided Agent Framework for Kernel Generation
作者: Yaochen Han, Ke Fan, Hongxu Jiang, Wanqi Xu, Weiyu Xie, Runhua Zhang, Chenhui Zhu, Yixiang Zhang
分类: cs.AI
发布日期: 2026-06-25
💡 一句话要点
提出EGG框架以解决GPU内核生成中的优化指导问题
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: GPU内核生成 专家指导 多代理协作 算法结构设计 硬件特定调优 性能优化 深度学习
📋 核心要点
- 现有的内核生成方法在正确性和性能之间难以平衡,缺乏领域特定的优化指导,限制了优化空间的有效探索。
- EGG框架通过将内核生成分为算法结构设计和硬件特定调优两个阶段,结合专家优化原则,提供了明确的优化目标。
- 在KernelBench和真实工作负载上的实验表明,EGG框架实现了2.13倍的平均加速,显著优于现有的代理和强化学习方法。
📝 摘要(中文)
高性能GPU内核对于降低大型语言模型的计算成本至关重要,但其开发依赖于领域专家的手动调优。尽管基于LLM的方法在自动化内核生成方面展现出潜力,但在正确性和高性能之间仍面临挑战。为此,本文提出EGG框架,通过引入专家优化原则来指导LLM的决策,将内核生成分为算法结构设计和硬件特定调优两个阶段,从而明确优化目标,确保优化过程的稳定性。实验结果表明,EGG在KernelBench和实际工作负载上实现了2.13倍的平均加速,超越了现有的基于代理和强化学习的方法。
🔬 方法详解
问题定义:本文旨在解决GPU内核生成过程中缺乏有效优化指导的问题。现有方法依赖于手动调优,导致性能和正确性难以兼得。
核心思路:EGG框架通过引入专家优化原则,分阶段进行内核生成,首先设计高质量的算法结构,然后进行硬件特定的调优,以实现更高效的优化过程。
技术框架:EGG框架包括两个主要阶段:第一阶段是算法结构设计,建立计算结构的基础;第二阶段是硬件特定调优,通过并行映射、张量分块和内存优化进行针对性调整。
关键创新:EGG的创新在于其专家指导的多代理协作机制,确保了阶段间和阶段内的上下文管理,从而实现稳定的优化轨迹。这与现有方法的单一优化策略形成鲜明对比。
关键设计:EGG框架设计了明确的优化目标,并通过阶段感知的多代理协作机制来管理优化过程,确保了设计空间的结构化和渐进式的优化。
🖼️ 关键图片
📊 实验亮点
EGG框架在KernelBench和真实工作负载上的实验结果显示,平均加速比达到2.13倍,显著优于现有的基于代理和强化学习的方法。这一性能提升表明EGG在内核生成中的有效性和实用性。
🎯 应用场景
EGG框架在大型语言模型的GPU内核生成中具有广泛的应用潜力,能够显著降低计算成本,提高模型的运行效率。未来,该框架可扩展到其他需要高性能计算的领域,如深度学习、科学计算等,推动相关技术的发展。
📄 摘要(原文)
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.