EGG: An Expert-Guided Agent Framework for Kernel Generation

🕒 Published (v1): 2026-06-25 08:44 UTC · Source: Arxiv · link

Why this paper was selected

Expert-guided agent for kernel generation; structured domain knowledge meets LLM code synthesis

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

EGG is a multi-agent framework that generates high-performance GPU kernels (in Triton) by encoding expert optimization workflows into a hierarchical two-stage pipeline: algorithmic structure design followed by hardware-specific tuning. A stage-aware multi-agent collaboration mechanism coordinates three specialized agents (code, profile, debug) with structured context management to prevent objective drift. EGG achieves a 2.13× average speedup over PyTorch Eager on KernelBench, outperforming both RL-based and prior agent-based approaches.

Problem

Existing LLM-based GPU kernel generation falls into two failure modes: (1) RL/fine-tuning approaches suffer from data scarcity and produce kernels with <50% correctness (e.g., AutoTriton); (2) agent-based systems rely on coarse-grained feedback (execution time only), lack domain-specific optimization structure, and exhibit trial-and-error behavior with marginal gains across rounds. Both categories fail to constrain the exponentially large (\(\sim10^9\) configurations for small subgraphs) hardware-dependent design space.

Method

EGG decomposes kernel generation into two hierarchical stages, each guided by explicit expert objectives:

Stage 1 — Algorithmic Structure Design: - Multi-seed search: generates a small set of structurally diverse initial kernels (e.g., direct vs. im2col convolution), then filters to the best performer after refinement. - Algorithmic refinement: applies semantics-preserving structural transformations (operator fusion, dataflow reorganization) to eliminate high-level inefficiencies (e.g., FlashAttention-style fusion eliminating \(4N^2\) memory accesses).

Stage 2 — Hardware-Specific Tuning (three sequential sub-stages): 1. Parallel mapping: maps operator dimensions \((d_1,\ldots,d_n) \to (G_1,\ldots,G_n)\) to maximize SM utilization. 2. Tensor tiling: selects tile sizes \((B_1,\ldots,B_n)\) with \(G_i = \lceil d_i/B_i \rceil\), balancing on-chip reuse against register/shared memory pressure; proposes candidates and selects via runtime measurement. 3. Memory optimization: coalesces global memory accesses and adjusts multi-buffering depth for software pipelining.

Stage-aware multi-agent collaboration: Three agents collaborate within each stage: a Profile Agent (parses NCU profiling reports, identifies bottleneck, proposes fix), a Code Agent (implements modifications), and a Debug Agent (repairs compile/runtime/numerical failures). Communication is via structured JSON to compress context. Inter-stage, only finalized decisions propagate (intermediate exploratory outputs are discarded) to prevent regression and context bloat.

Key Contributions

• Expert-guided staged decomposition that maps kernel optimization to a structured sequence of constrained sub-problems aligned with human expert workflows.
• Multi-seed search + algorithmic refinement combination as Stage 1, establishing a strong performance upper bound before hardware tuning begins.
• Stage-aware multi-agent mechanism with selective inter-stage context propagation and structured JSON intra-stage information exchange, preventing objective drift across the staged pipeline.
• 100% correctness rate on KernelBench with 2.13× average speedup over PyTorch Eager.

Results

• KernelBench (250 tasks, 3 difficulty levels, RTX 4090, GPT-5.1 backend):
• Level 1: 100% success, 72% Fast1, 1.83× speedup (vs. CudaForge 1.43×, Torch Compile 1.09×)
• Level 2: 100% success, 100% Fast1, 2.73× speedup (vs. CudaForge 2.00×, Torch Compile 1.38×)
• Level 3: 100% success, 94% Fast1, 1.52× speedup (vs. CudaForge 1.30×, Torch Compile 1.36×)
• Overall average: 2.13× over PyTorch Eager, 1.60× over Torch Compile
• Ablation (average Fast1 / speedup): Full EGG: 87.6% / 2.13×; w/o multi-seed: 74.0% / 1.67×; w/o algo refinement: 78.4% / 1.84×; w/o HW tuning: 60.8% / 1.52×; w/o multi-agent: 72.0% / 1.46×.
• Cumulative speedup across stages: initial seed 1.00×→ algo refinement 2.00×→ parallel mapping 2.37×→ tiling 2.63×→ memory optimization 2.81× (on a representative task).
• Outperforms RL-based AutoTriton (which has <50% correctness) and general LLMs (DeepSeek V3.2, ChatGPT 5.1 achieve <70% success at all levels).
• Generalization confirmed on RTX 5090, H20, RTX PRO 6000 (appendix).

Limitations

• Targets Triton exclusively; CUDA's larger optimization space is explicitly out of scope.
• Evaluation depends on a single high-capability LLM backend (GPT-5.1); cost per task is not fully reported (truncated in the provided text).
• Expert-defined stage structure may not generalize to operator classes with atypical optimization profiles not captured by the four GPU optimization principles.
• Multi-seed search adds upfront generation cost; number of seeds is a hyperparameter not fully characterized in the provided text.
• Hardware generalization beyond NVIDIA GPUs (AMD, TPUs) is not addressed.

Relevance to Agentic AI / LLM Agents

EGG is a direct case study in structuring agentic LLM pipelines around domain-expert workflows rather than relying on end-to-end reinforcement or free-form iteration — a central design challenge for LLM agents in specialized technical domains. The stage-aware context management strategy (selective inter-stage propagation, structured JSON intra-stage exchange) is a concrete solution to objective drift and context bloat, failure modes endemic to long-horizon multi-agent loops. The three-agent collaboration pattern (generator / profiler / debugger) with role-specialized prompts and structured interfaces is directly transferable to other code-generation agent architectures. The result that structured decomposition + domain priors beats both RL fine-tuning and unstructured iterative agents strongly supports the hypothesis that expert-informed workflow decomposition is a high-leverage design choice for agentic systems in constrained optimization domains.