Skip to content

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 agentic framework for GPU kernel generation; novel human-in-loop harness design with measurable compute impact

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 automates high-performance GPU Triton kernel generation by encoding expert optimization workflows as a two-stage hierarchical decomposition (algorithmic structure design → hardware-specific tuning) with a stage-aware multi-agent collaboration mechanism. It achieves 100% correctness and a 2.13× average speedup over PyTorch Eager on KernelBench, outperforming both RL-based and agent-based baselines.

Problem

LLM-based GPU kernel generation fails to jointly achieve correctness and high performance because LLMs lack domain-specific optimization guidance, leading to trial-and-error exploration over an exponentially large (\(\sim10^9\) configurations per subgraph), hardware-dependent design space. RL-based methods (e.g., AutoTriton) suffer from <50% correctness; agent-based methods (e.g., CudaForge) rely on coarse execution-time feedback without principled structuring of the search space.

Method

EGG decomposes kernel generation into two hierarchical stages:

  1. Algorithmic Structure Design: Generates \(k\) seed kernels with structurally distinct algorithmic paradigms (multi-seed search), applies semantics-preserving expert transformations to each (operator fusion, dataflow reorganization), then filters to the best-performing seed via lightweight benchmarking.

  2. Hardware-Specific Tuning: Three sequential sub-stages, each narrowing the search space for the next:

  3. Parallel mapping: maps operator dimensions \((d_1,\ldots,d_n)\) to GPU grid dimensions \((G_1,\ldots,G_n)\) to maximize SM utilization.
  4. Tensor tiling: selects tile sizes \((B_1,\ldots,B_n)\) satisfying \(G_i = \lceil d_i/B_i \rceil\), balancing register/shared-memory budget against data reuse.
  5. Memory optimization: coalesces global memory accesses and tunes multi-buffering depth for software pipelining.

Within each sub-stage, a stage-aware multi-agent loop coordinates three agents: a Profile Agent (interprets NCU profiling metrics, identifies bottlenecks, proposes fixes in structured JSON), a Code Agent (rewrites the kernel), and a Debug Agent (resolves correctness failures). Context is managed via inter-stage selective propagation (only finalized decisions carry forward, discarding exploratory outputs) and intra-stage structured JSON exchange (compressed bottleneck reports, modification plans, error diagnostics).

Key Contributions

  • Expert-guided staged decomposition that converts a joint \(10^9\)-configuration optimization into a sequence of well-scoped sub-problems aligned with human expert workflows.
  • Stage-aware multi-agent collaboration with structured JSON interfaces that prevents objective drift and context pollution across optimization rounds.
  • Selective inter-stage context propagation that accumulates improvements without regressing prior stages.
  • 100% correctness + 2.13× mean speedup on KernelBench (all three difficulty levels), with generalization validated on RTX 5090, H20, and RTX PRO 6000.

Results

  • KernelBench overall: 100% success rate, 72%/100%/94% Fast1 rate (L1/L2/L3), 1.83×/2.73×/1.52× speedup vs. PyTorch Eager.
  • vs. CudaForge (closest agent baseline): +0.40×, +0.73×, +0.22× speedup per level; CudaForge Fast1 was 56%/90%/72% vs. EGG's 72%/100%/94%.
  • vs. Torch Compile: 1.60× aggregate speedup.
  • vs. AutoTriton (RL-based): AutoTriton correctness <50%, EGG 100%.
  • vs. DeepSeek-V3.2 / ChatGPT-5.1 (vanilla LLMs): those achieve 0.84–1.13× speedup (often slower than PyTorch), EGG 2.13×.
  • Ablation (cumulative speedup across stages): multi-seed search → 1.0×, +algorithmic refinement → 2.00×, +parallel mapping → 2.37×, +tensor tiling → 2.63×, +memory optimization → 2.81×.
  • Ablation on Fast1 rate: full EGG 87.6% vs. removing multi-seed (74.0%), algo refinement (78.4%), HW tuning (60.8%), multi-agent (72.0%).

Limitations

  • Targets Triton kernels only; CUDA's larger optimization space is explicitly deferred as harder to automate.
  • Relies on GPT-5.1 (or Claude Opus 4.5) as backbone; performance and cost scale with LLM API calls.
  • Multi-seed search and multi-stage tuning incur higher latency/cost per kernel than single-pass methods (exact figures truncated in provided text).
  • Evaluation hardware is limited to a small set of NVIDIA GPUs; AMD/other architectures not addressed.
  • No handling of cross-kernel or graph-level optimization (e.g., fusion across operator boundaries at the graph level).

Relevance to Harnesses / Meta-Harnesses

EGG is a concrete domain-specific harness that operationalizes the meta-harness pattern: a fixed outer orchestration loop (stage sequencing + agent routing) wraps interchangeable inner agents (Profile/Code/Debug), with explicit inter-stage context filtering serving as the "handoff protocol" between sub-harnesses. The selective context propagation mechanism—retaining only finalized artifacts at stage boundaries—directly addresses a known failure mode of naive chaining harnesses (context pollution causing regression). Its structured JSON intra-stage exchange is a lightweight implementation of the typed message-passing interface that more general meta-harnesses use to coordinate heterogeneous agents, making it a replicable pattern for other code-optimization harnesses.