EGG：专家引导的内核生成智能体框架

# EGG：面向内核生成的专家引导智能体框架
来源：https://arxiv.org/html/2606.26758
Ke FanHongxu JiangWanqi XuWeiyu XieRunhua ZhangChenhui ZhuYixiang Zhang

###### 摘要

高性能 GPU 内核对于降低大语言模型（LLM）指数级增长的计算成本至关重要，但其开发严重依赖领域专家的人工调优。虽然基于 LLM 的最新方法在自动化内核生成方面展现出潜力，但它们在正确性和高性能方面仍难以兼得。这一局限主要源于缺乏领域特定的优化引导，阻碍了对优化空间的有效探索。我们提出 EGG，一个用于内核生成的专家引导智能体框架，将专家优化原则融入其中以指导 LLM 的决策。受专家工作流程启发，我们将内核生成分解为两个层级阶段：1）算法结构设计，建立高质量的计算结构基础；2）硬件特定调优，通过并行映射、张量分块和内存优化进行针对性调整。这种分阶段分解定义了明确的优化目标，构建了设计空间以实现逐步精化。为此，我们设计了一种阶段感知的多智能体协作机制，用于阶段间和阶段内的上下文管理，确保稳定的优化轨迹。在 KernelBench 和真实工作负载上的实验表明，EGG 相比 PyTorch 实现了平均 2.13 倍的加速，优于现有的基于智能体和基于强化学习的方法。机器学习，ICML

## 1 引言

大语言模型（LLM）推动了日益增长的训练和推理成本（Raiaan 等人，2024 (https://arxiv.org/html/2606.26758#bib.bib30)）。高性能 GPU 内核通过决定现代深度学习工作负载的吞吐量和效率，对于降低这些成本至关重要。然而，开发高效的 GPU 内核仍然严重依赖领域专家的人工调优。随着 GPU 和模型架构的演进与多样化，这一过程变得愈发昂贵且耗时（Chen 等人，2018 (https://arxiv.org/html/2606.26758#bib.bib31)）。这些限制促使自动化 GPU 内核生成成为一个关键的研究问题。

LLM 的最新进展在通用代码生成（如 Python）方面展示出强大能力（He 等人，2025 (https://arxiv.org/html/2606.26758#bib.bib37)；Zhang 等人，2024 (https://arxiv.org/html/2606.26758#bib.bib36)），使其成为自动化 GPU 内核生成的有力候选（Jiang 等人，2026 (https://arxiv.org/html/2606.26758#bib.bib29)）。然而，与通用编程不同，GPU 内核生成是一个在广阔且依赖硬件的设计空间中的严格约束优化问题。内核代码必须满足严格的语法和语义正确性要求，同时实现高性能。此外，内核性能取决于并行映射、内存层次利用及其他硬件特定特性之间的复杂交互，即使微小的代码改动也可能导致数量级的性能差异。先前的研究表明，即使是小型神经网络子图，在 GPU 上也可能暴露多达 10^9 种配置的优化空间（Zhai 等人，2024 (https://arxiv.org/html/2606.26758#bib.bib12)）。因此，直接应用 LLM 而不提供额外信息，往往会产生无效或远未达到最优的内核，这促使我们需要专门的方法来弥合这一差距。

现有的基于 LLM 的 GPU 内核生成方法可分为两类：1）使用微调和强化学习（RL）使 LLM 适应内核生成领域（Li 等人，2025c (https://arxiv.org/html/2606.26758#bib.bib5)；Woo 等人，2025 (https://arxiv.org/html/2606.26758#bib.bib6)）；2）在通用 LLM 上构建基于智能体的系统，通过迭代精化来改进生成的内核，无需额外模型训练（Zhang 等人，2025b (https://arxiv.org/html/2606.26758#bib.bib7)；Li 等人，2025a (https://arxiv.org/html/2606.26758#bib.bib18)；Zhang 等人，2025a (https://arxiv.org/html/2606.26758#bib.bib19)）。尽管存在差异，两类方法都从根本上受限于 GPU 内核优化空间的极端复杂性。1）高质量内核数据集稀缺，生成的内核必须满足严格的语法、语义及硬件特定约束。因此，RL 生成的内核可能无法编译或性能有限。例如，AutoTriton 报告的平均正确率低于 50%（Li 等人，2025c (https://arxiv.org/html/2606.26758#bib.bib5)）。2）基于智能体的方法（如 CudaForge（Zhang 等人，2025b (https://arxiv.org/html/2606.26758#bib.bib7)））通过多轮优化提高了内核的正确性和鲁棒性。然而，大多数现有智能体（Wang 等人，2025 (https://arxiv.org/html/2606.26758#bib.bib15)）缺乏领域特定的优化引导，依赖粗粒度的反馈（如执行时间），导致试错式探索，经过多轮优化后仅获得边际改进。这些局限表明，要充分释放 LLM 在 GPU 内核生成中的潜力，需要集成专家级的内核优化原则来约束设计空间并指导探索过程。有效系统不能依赖试错或端到端学习，而必须对设计空间施加反映专家推理和硬件感知约束的结构化优化目标。

为此，我们提出 EGG，一个用于自动 GPU 内核生成的专家引导智能体框架。EGG 通过显式建模专家优化工作流程，引导 LLM 做出分阶段的优化决策，在结构化探索内核优化空间的同时保持通用 LLM 的灵活性。具体而言，EGG 采用专家引导的分阶段优化，将内核生成分解为两个层级阶段：1）算法结构设计和 2）硬件特定调优。算法结构设计通过结合多种子搜索与算法精化技术，建立强大的性能上界。硬件特定调优随后通过三个具有明确目标的连续子阶段（并行映射、张量分块和内存优化）系统地实现这一潜力。通过将复杂优化问题分解为与专家工作流一致的约束子问题，EGG 有效引导基于 LLM 的智能体走向高性能内核实现。为支持这一过程，我们设计了一种阶段感知的多智能体协作机制，包含代码智能体、分析智能体和调试智能体，并在阶段间和阶段内进行上下文管理。对于阶段间，选择性上下文传播策略仅保留最终确定的优化决策，丢弃中间输出。对于阶段内，这些智能体通过结构化信息交换（包括瓶颈识别和优化建议）围绕明确定义的目标进行协作。这种设计使得优化过程中能够实现稳定、累积的性能提升。

我们在 KernelBench（Ouyang 等人，2025 (https://arxiv.org/html/2606.26758#bib.bib4)）和真实工作负载上对 EGG 进行了系统评估。实验结果表明，即使在具有挑战性的任务上，EGG 也能持续生成正确且具有竞争力的内核。EGG 相比 PyTorch 实现了平均 2.13 倍的加速，优于现有的基于智能体和基于强化学习的方法。本文的主要贡献如下：

- •我们提出了一种用于 GPU 内核生成的专家引导智能体框架，将专家内核优化原则融入其中以指导 LLM 的决策。
- •我们根据专家工作流程将内核生成分解为*算法结构设计*和*硬件特定调优*阶段，构建了优化空间以实现更高性能。
- •我们设计了一种*阶段感知的多智能体协作*机制，能够在所有优化阶段实现稳定且累积的性能提升，在 KernelBench 和真实工作负载上实现 100% 正确率和平均 2.13 倍加速。

参见图注

图 1：EGG 概览。EGG 采用专家引导的分阶段优化策略，包含算法结构设计和硬件特定调优（紫色区域）。在每个阶段内，采用阶段感知的多智能体协作机制以确保稳定的优化轨迹（蓝色区域）。黄色区域展示了代表性 Triton 代码片段，具体说明了每个优化阶段的效果。最佳彩色视图。

## 2 相关工作

### 2.1 GPU 内核优化

GPU 内核优化依赖一组广泛适用的技术：1）算子融合通过合并多个操作减少内核启动开销和冗余内存访问（Jia 等人，2019 (https://arxiv.org/html/2606.26758#bib.bib20)）；2）并行映射决定计算如何在 GPU 执行网格上分布，直接影响流式多处理器（SM）的利用率（Osama 等人，2023 (https://arxiv.org/html/2606.26758#bib.bib21)）；3）分块将计算划分为适合寄存器和缓存容量的块，以最大化数据重用（Cai 等人，2023 (https://arxiv.org/html/2606.26758#bib.bib23)）；4）流水线将计算与数据移动重叠，以隐藏内存访问延迟（Cheng 等人，2025 (https://arxiv.org/html/2606.26758#bib.bib22)）。虽然这些优化原则在 GPU 平台上大致共享，但其具体实现高度依赖硬件：SM 配置、寄存器文件大小、片上内存层次结构和内存带宽都关键地影响最优设计选择。如今的高性能内核主要来自专家优化的库（如 cuBLAS（NVIDIA, 2026a (https://arxiv.org/html/2606.26758#bib.bib17)），cuDNN（NVIDIA, 2026b (https://arxiv.org/html/2606.26758#bib.bib16)））或专用实现（如 FlashAttention（Dao 等人，2022 (https://arxiv.org/html/2606.26758#bib.bib14)）），实现接近峰值的性能，但开发成本高且对新型算子结构的适应性有限。

### 2.2 编译器与 Triton

编译器和领域特定语言（DSL）通过提供更高级抽象同时将底层细节委托给编译器，部分缓解了这些挑战（Wang 等人，2026 (https://arxiv.org/html/2606.26758#bib.bib24)；Spector 等人，2025 (https://arxiv.org/html/2606.26758#bib.bib25)）。Triton（Tillet 等人，2019 (https://arxiv.org/html/2606.26758#bib.bib13)）是一种基于分块编程抽象的 Python 嵌入式 DSL，已被广泛用于现代 LLM 系统中实现性能关键的内核（如 SGLang（Zheng 等人，2024 (https://arxiv.org/html/2606.26758#bib.bib2)），vLLM（Kwon 等人，2023 (https://arxiv.org/html/2606.26758#bib.bib1)））。然而，使用 Triton 实现高性能仍需要仔细手动调优并行性、分块和流水线决策。次优配置可能导致显著性能损失。虽然 CUDA 提供了更细粒度的控制，但其更大的优化空间使得自动化优化更难控制。因此，本文聚焦于 Triton，因其结构化抽象为 LLM 驱动的优化提供了更易处理的空间。

### 2.3 基于 LLM 的内核生成

**基于强化学习的方法** 预训练 LLM 由于对硬件特性和优化原则的理解有限，难以生成高性能 GPU 内核。为解决此问题，近期工作（Baronio 等人，2026 (https://arxiv.org/html/2606.26758#bib.bib3)；Li 等人，2025c (https://arxiv.org/html/2606.26758#bib.bib5)；Woo 等人，2025 (https://arxiv.org/html/2606.26758#bib.bib6)；Li 等人，2026 (https://arxiv.org/html/2606.26758#bib.bib33)；Su 等人，2025 (https://arxiv.org/html/2606.26758#bib.bib32)；Fisches 等人，2025 (https://arxiv.org/html/2606.26758#bib.bib34)）采用任务特定的微调和强化学习（RL）。Kevin-32B（Baronio 等人，2026 (https://arxiv.org/html/2606.26758#bib.bib3)）通过多轮 RL 训练提高了内核的正确性和性能。AutoTriton（Li 等人，2025c (https://arxiv.org/html/2606.26758#bib.bib5)）和 TritonRL（Woo 等人，2025 (https://arxiv.org/html/2606.26758#bib.bib6)）应用监督微调学习 Triton 语法，然后使用 RL 进一步优化性能。然而，这些方法需要大规模领域特定数据集和大量计算资源，可能限制其实用性。

**基于智能体的方法** 多智能体系统（Zhang 等人，2025b (https://arxiv.org/html/2606.26758#bib.bib7)；Lange 等人，2025 (https://arxiv.org/html/2606.26758#bib.bib8)；Wang 等人，2025 (https://arxiv.org/html/2606.26758#bib.bib15)；Li 等人，2025a (https://arxiv.org/html/2606.26758#bib.bib18)；Sereda 等人，2025 (https://arxiv.org/html/2606.26758#bib.bib35)）提供了一种避免昂贵训练成本的替代方法。CudaForge（Zhang 等人，2025b (https://arxiv.org/html/2606.26758#bib.bib7)）采用由编码器和评判器组成的硬件反馈驱动框架，使用 NVIDIA Nsight Compute (NCU) 指标指导优化。AI CUDA Engineer（Lange 等人，2025 (https://arxiv.org/html/2606.26758#bib.bib8)）使用随机变异和迭代搜索探索内核设计。然而，这些方法缺乏对内核优化原则的整体理解，依赖粗粒度的性能反馈来指导调整，导致行为不稳定且改进有限。

## 3 方法

### 3.1 概述

在本节中，我们介绍 EGG，一个用于自动生成高性能 GPU 内核的专家引导多智能体框架。如图 1 (https://arxiv.org/html/2606.26758#S1.F1) 所示，EGG 包含两个核心组件：1）专家引导的分阶段优化策略，将内核优化空间分解为定义明确的子问题；2）阶段感知的多智能体协作机制，在每个阶段内协调专用智能体，并在阶段间传播精化后的决策。我们首先描述专家引导的分阶段优化策略，然后针对每个优化阶段设计阶段感知的多智能体协作机制。

### 3.2 专家引导的分阶段优化

GPU 内核优化涉及广阔设计空间中的复杂优化决策，使得端到端优化难以控制。为应对这一挑战，我们引入专家优化工作流程，将 GPU 内核生成分解为两个层级阶段：*算法结构设计*和*硬件特定调优*。算法结构决策决定内核的计算结构和数据流组织，从根本上奠定可达性能的上界。针对目标设备的硬件特定调优聚焦于三个子阶段：并行映射、张量分块和内存优化，直接影响运行时性能。这种层级分解构建了优化空间，并实现渐进式的阶段级精化。

#### 3.2.1 算法结构设计

算法结构设计建立高质量算法基础，决定性能上界。一个关键挑战