arXiv:2606.26758cs.AI2026-06

用专家经验指导AI生成高性能GPU核函数,速度比PyTorch快2.13倍

EGG: An Expert-Guided Agent Framework for Kernel Generation

论文配图:EGG: An Expert-Guided Agent Framework for Kernel Generation
图 1 · 摘自论文原文
  • 分两阶段生成:先设计算法结构,再针对硬件调优
  • 在真实任务上平均提速2.13倍,优于现有自动化方法
  • 适合需要高效推理的大型语言模型开发者

高性能GPU核函数对降低大语言模型日益增长的计算成本至关重要,但其开发仍严重依赖领域专家的手动调优。尽管基于大语言模型的方法在自动化核函数生成方面展现潜力,但仍难以同时保证正确性与高性能。这主要源于缺乏领域特定优化指导,限制了优化空间的有效探索。本文提出EGG——一种专家引导的核函数生成框架,将专家工作流转化为优化原则以指导LLM决策。我们把核函数生成分解为两个层级阶段:1)算法结构设计,建立高质量的计算结构基础;2)硬件特定调优,通过并行映射、张量分块和内存优化进行针对性调整。这种分阶段设计明确了优化目标,结构化设计空间以实现渐进式优化。为此,我们设计了阶段感知的多智能体协作机制,实现跨阶段与同阶段上下文管理,确保优化轨迹稳定。在KernelBench和真实工作负载上的实验表明,EGG相比PyTorch平均提速2.13倍,优于现有基于代理和强化学习的方法。

原文摘要 · Abstract (English)

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.

GPU优化LLM加速AI生成

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。