EGG: An Expert-Guided Agent Framework for Kernel Generation

# EGG: An Expert-Guided Agent Framework for Kernel Generation
Source: [https://arxiv.org/html/2606.26758](https://arxiv.org/html/2606.26758)
Ke FanHongxu JiangWanqi XuWeiyu XieRunhua ZhangChenhui ZhuYixiang Zhang

###### Abstract

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 proposeEGG, anExpert\-Guided Agent Framework for KernelGeneration, 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 a2\.13×2\.13\\timesaverage speedup over PyTorch, outperforming existing agent\-based and RL\-based approaches\.

Machine Learning, ICML

## 1Introduction

Large language models \(LLMs\) drive increasing training and inference costs\(Raiaanet al\.,[2024](https://arxiv.org/html/2606.26758#bib.bib30)\)\. High\-performance GPU kernels are critical to reducing these costs by determining the throughput and efficiency of modern deep learning workloads\. However, developing efficient GPU kernels still relies heavily on manual tuning by domain experts\. As GPU and model architectures evolve and diversify, this process becomes increasingly expensive and time\-consuming\(Chenet al\.,[2018](https://arxiv.org/html/2606.26758#bib.bib31)\)\. These limitations motivate automated GPU kernel generation as a critical research problem\.

Recent advances in LLMs have demonstrated strong capabilities in general\-purpose code generation \(e\.g\., Python\)\(Heet al\.,[2025](https://arxiv.org/html/2606.26758#bib.bib37); Zhanget al\.,[2024](https://arxiv.org/html/2606.26758#bib.bib36)\), making them promising candidates for automating GPU kernel generation\(Jianget al\.,[2026](https://arxiv.org/html/2606.26758#bib.bib29)\)\. However, unlike general\-purpose programming, GPU kernel generation is a tightly constrained optimization problem over a vast, hardware\-dependent design space\. Kernel code must satisfy strict syntactic and semantic correctness requirements while simultaneously achieving high performance\. Additionally, kernel performance depends on intricate interactions among parallel mapping, memory hierarchy utilization, and other hardware\-specific features, where even minor code changes can lead to orders\-of\-magnitude performance differences\. Prior work shows that even a small neural network subgraph can expose an optimization space of up to10910^\{9\}configurations on GPUs\(Zhaiet al\.,[2024](https://arxiv.org/html/2606.26758#bib.bib12)\)\. As a result, directly applying LLMs without additional information often produces kernels that are invalid or far from optimal, motivating specialized methodology to bridge this gap\.

Existing LLM\-based GPU kernel generation methods can be divided into two categories: 1\) using fine\-tuning and reinforcement learning \(RL\) to adapt LLMs to the kernel generation domain\(Liet al\.,[2025c](https://arxiv.org/html/2606.26758#bib.bib5); Wooet al\.,[2025](https://arxiv.org/html/2606.26758#bib.bib6)\); 2\) building agent\-based systems on general\-purpose LLMs, which leverage iterative refinement to improve generated kernels without additional model training\(Zhanget al\.,[2025b](https://arxiv.org/html/2606.26758#bib.bib7); Liet al\.,[2025a](https://arxiv.org/html/2606.26758#bib.bib18); Zhanget al\.,[2025a](https://arxiv.org/html/2606.26758#bib.bib19)\)\.

Despite their differences, both classes of approaches are fundamentally constrained by the extreme complexity of the GPU kernel optimization space\. 1\) High\-quality kernel datasets are scarce, and generated kernels must satisfy strict syntactic, semantic, and hardware\-specific constraints\. As a result, RL\-generated kernels may fail to compile or deliver limited performance\. For example, AutoTriton reports an average correctness rate below 50%\(Liet al\.,[2025c](https://arxiv.org/html/2606.26758#bib.bib5)\)\. 2\) Agent\-based approaches, such as CudaForge\(Zhanget al\.,[2025b](https://arxiv.org/html/2606.26758#bib.bib7)\), improve kernel correctness and robustness via multi\-turn optimization\. However, most existing agents\(Wanget al\.,[2025](https://arxiv.org/html/2606.26758#bib.bib15)\)lack domain\-specific optimization guidance and rely on coarse\-grained feedback \(e\.g\., execution time\), resulting in trial\-and\-error exploration and only marginal improvement after multiple optimization rounds\.

These limitations suggest that fully unlocking the potential of LLMs for GPU kernel generation requires integrating expert\-level kernel optimization principles to constrain the design space and guide the exploration process\. Rather than relying on trial\-and\-error or end\-to\-end learning, effective systems must impose structured optimization objectives on the design space that reflect expert reasoning and hardware\-aware constraints\.

To this end, we propose EGG, an expert\-guided agent framework for automatic GPU kernel generation\. EGG guides LLMs to make stage\-wise optimization decisions by explicitly modeling expert optimization workflows, enabling structured exploration of the kernel optimization space while preserving the flexibility of general\-purpose LLMs\.

Specifically, EGG employsexpert\-guided staged optimization, decomposing kernel generation into two hierarchical stages:1\) algorithmic structure designand2\) hardware\-specific tuning\. The algorithmic structure design establishes a strong performance upper bound by combining multi\-seed search with algorithmic refinement techniques\. The hardware\-specific tuning then systematically realizes this potential through three sequential sub\-stages with explicit objectives: parallel mapping, tensor tiling, and memory optimization\. By decomposing the complex optimization problem into constrained sub\-problems aligned with expert workflows, EGG effectively guides LLM\-based agents toward high\-performance kernel implementations\.

To support this process, we propose astage\-aware multi\-agent collaboration mechanism, which consists of a code agent, a profile agent, and a debug agent, with context management during inter and intra\-stages\. For inter\-stage, a selective context propagation strategy retains only finalized optimization decisions while discarding intermediate outputs\. For intra\-stage, these agents coordinate around clearly defined objectives via structured information exchange, including bottleneck identification and optimization proposals\. This design enables stable, cumulative performance improvements throughout the optimization process\.

We systematically evaluate EGG on KernelBench\(Ouyanget al\.,[2025](https://arxiv.org/html/2606.26758#bib.bib4)\)and real\-world workloads\. Experimental results show that EGG consistently generates correct and competitive kernels even for challenging tasks\. EGG achieves a2\.13×2\.13\\timesaverage speedup over PyTorch, outperforming existing agent\-based and RL\-based approaches\.

The key contributions of this paper are as follows:

- •We propose an expert\-guided agent framework for GPU kernel generation, which incorporates expert kernel optimization principles to guide LLMs’ decisions\.
- •We decompose kernel generation into*algorithmic structure design*and*hardware\-specific tuning*stages according to expert workflows, which structures the optimization space to achieve higher performance\.
- •We design a*stage\-aware multi\-agent collaboration*mechanism that enables stable and cumulative performance improvements across all optimization stages, achieving 100% correctness and a2\.13×2\.13\\timesaverage speedup on KernelBench and real\-world workloads\.

![Refer to caption](https://arxiv.org/html/2606.26758v1/x1.png)

Figure 1:Overview of EGG\. EGG adopts anexpert\-guided staged optimizationstrategy that consists ofAlgorithmic Structure DesignandHardware\-Specific Tuning\(purple region\)\. Within each stage, astage\-aware multi\-agent collaborationmechanism is employed to ensure stable optimization trajectories \(blue region\)\. The yellow region presents representative Triton code snippets that concretely illustrate the effect of each optimization stage\. Best viewed in color\.
## 2Related Work

### 2\.1GPU Kernel Optimization

GPU kernel optimization relies on a set of widely applicable techniques: 1\)operator fusionreduces kernel launch overhead and redundant memory accesses by merging multiple operations\(Jiaet al\.,[2019](https://arxiv.org/html/2606.26758#bib.bib20)\); 2\)parallel mappingdetermines how computation is distributed across the GPU execution grid, directly affecting the utilization of streaming multiprocessors \(SMs\)\(Osamaet al\.,[2023](https://arxiv.org/html/2606.26758#bib.bib21)\); 3\)tilingpartitions computation into tiles that fit register and cache capacities to maximize data reuse\(Caiet al\.,[2023](https://arxiv.org/html/2606.26758#bib.bib23)\); 4\)pipeliningoverlaps computation with data movement to hide memory access latency\(Chenget al\.,[2025](https://arxiv.org/html/2606.26758#bib.bib22)\)\.

While these optimization principles are largely shared across GPU platforms, their concrete implementations are highly hardware\-dependent: SM configuration, register file size, on\-chip memory hierarchy, and memory bandwidth all critically influence optimal design choices\. High\-performance kernels today predominantly come from expert\-optimized libraries \(e\.g\., cuBLAS\(NVIDIA,[2026a](https://arxiv.org/html/2606.26758#bib.bib17)\), cuDNN\(NVIDIA,[2026b](https://arxiv.org/html/2606.26758#bib.bib16)\)\) or specialized implementations \(e\.g\., FlashAttention\(Daoet al\.,[2022](https://arxiv.org/html/2606.26758#bib.bib14)\)\), achieving near\-peak performance but at high development cost and limited adaptability to new operator structures\.

### 2\.2Compilers and Triton

Compilers and domain\-specific languages \(DSLs\) partially alleviate these challenges by providing higher\-level abstractions while delegating low\-level details to the compiler\(Wanget al\.,[2026](https://arxiv.org/html/2606.26758#bib.bib24); Spectoret al\.,[2025](https://arxiv.org/html/2606.26758#bib.bib25)\)\. Triton\(Tilletet al\.,[2019](https://arxiv.org/html/2606.26758#bib.bib13)\), a Python\-embedded DSL with tile\-based programming abstraction, has been widely adopted in modern LLM systems for implementing performance\-critical kernels \(e\.g\., SGLang\(Zhenget al\.,[2024](https://arxiv.org/html/2606.26758#bib.bib2)\), vLLM\(Kwonet al\.,[2023](https://arxiv.org/html/2606.26758#bib.bib1)\)\)\. However, achieving high performance with Triton still requires careful manual tuning of parallelism, tiling, and pipelining decisions\. Suboptimal configurations can lead to substantial performance loss\. While CUDA offers finer\-grained control, its larger optimization space makes automated optimization harder to control\. As a result, this work focuses on Triton, as its structured abstraction provides a more tractable space for LLM\-driven optimization\.

### 2\.3LLM\-based Kernel Generation

Reinforcement Learning–based MethodsPretrained LLMs struggle to generate high\-performance GPU kernels due to limited understanding of hardware characteristics and optimization principles\. To address this, recent works\(Baronioet al\.,[2026](https://arxiv.org/html/2606.26758#bib.bib3); Liet al\.,[2025c](https://arxiv.org/html/2606.26758#bib.bib5); Wooet al\.,[2025](https://arxiv.org/html/2606.26758#bib.bib6); Liet al\.,[2026](https://arxiv.org/html/2606.26758#bib.bib33); Suet al\.,[2025](https://arxiv.org/html/2606.26758#bib.bib32); Fischeset al\.,[2025](https://arxiv.org/html/2606.26758#bib.bib34)\)apply task\-specific fine\-tuning and reinforcement learning \(RL\)\. Kevin\-32B\(Baronioet al\.,[2026](https://arxiv.org/html/2606.26758#bib.bib3)\)improves kernel correctness and performance through multi\-round RL training\. AutoTriton\(Liet al\.,[2025c](https://arxiv.org/html/2606.26758#bib.bib5)\)and TritonRL\(Wooet al\.,[2025](https://arxiv.org/html/2606.26758#bib.bib6)\)apply supervised fine\-tuning to learn Triton syntax, then use RL to further refine performance\. However, these approaches require large\-scale domain\-specific datasets and substantial computational resources, which may limit their practical applicability\.

Agent\-based MethodsMulti\-agent systems\(Zhanget al\.,[2025b](https://arxiv.org/html/2606.26758#bib.bib7); Langeet al\.,[2025](https://arxiv.org/html/2606.26758#bib.bib8); Wanget al\.,[2025](https://arxiv.org/html/2606.26758#bib.bib15); Liet al\.,[2025a](https://arxiv.org/html/2606.26758#bib.bib18); Seredaet al\.,[2025](https://arxiv.org/html/2606.26758#bib.bib35)\)offer an alternative approach that avoids expensive training costs\. CudaForge\(Zhanget al\.,[2025b](https://arxiv.org/html/2606.26758#bib.bib7)\)employs a hardware\-feedback\-driven framework composed of a Coder and a Judger, using NVIDIA Nsight Compute \(NCU\) metrics to guide optimization\. AI CUDA Engineer\(Langeet al\.,[2025](https://arxiv.org/html/2606.26758#bib.bib8)\)uses stochastic mutation and iterative search to explore kernel designs\. However, these methods lack a holistic understanding of kernel optimization principles and rely on coarse\-grained performance feedback to guide adjustments, leading to unstable behavior and limited improvements\.

## 3Method

### 3\.1Overview

In this section, we introduceEGG, an expert\-guided multi\-agent framework for automatic generation of high\-performance GPU kernels\. As illustrated in Figure[1](https://arxiv.org/html/2606.26758#S1.F1), EGG consists of two core components: 1\) an*expert\-guided staged optimization*strategy that decomposes the kernel optimization space into well\-defined sub\-problems, and 2\) a*stage\-aware multi\-agent collaboration*mechanism that coordinates specialized agents within each stage while propagating refined decisions across stages\. We first describe the expert\-guided staged optimization strategy, followed by the design of the stage\-aware multi\-agent collaboration mechanism for each optimization stage\.

### 3\.2Expert\-Guided Staged Optimization

GPU kernel optimization involves complex optimization decisions over a vast design space, making end\-to\-end optimization difficult to control\. To address this challenge, we introduce expert optimization workflows to decompose GPU kernel generation into two hierarchical stages:*algorithmic structure design*and*hardware\-specific tuning*\. Algorithmic structure design determines the computational structure and dataflow organization of the kernel, fundamentally establishing the upper bound of achievable performance\. Hardware\-specific tuning for the target device focuses on three sub\-stages: parallel mapping, tensor tiling, and memory optimization, directly affecting runtime performance\. This hierarchical decomposition structures the optimization space and enables progressive, stage\-wise refinement\.

#### 3\.2\.1Algorithmic Structure Design

Algorithmic structure design establishes a high\-quality algorithmic foundation that determines the performance upper bound\. A key challenge is that the optimal algorithmic approach for a given operator is often non\-obvious, and LLM\-generated initial implementations frequently contain structural inefficiencies\. We address this challenge by combining two complementary strategies:*multi\-seed search*, which explores diverse algorithmic paradigms, and*algorithmic refinement*, which optimizes each paradigm to eliminate inefficiencies\.

Multi\-Seed Search\.For a given operator, fundamentally different algorithmic paradigms may exist\. For example, convolution can be implemented via a direct nested\-loop computation or through anim2col\-based matrix multiplication formulation\. Since our structured workflow constrains the optimization space at each subsequent stage, the initial choice of algorithmic paradigm has a significant impact on final performance\. Accordingly, as shown in Figure[1](https://arxiv.org/html/2606.26758#S1.F1), we generate a small set of initial kernel seeds with distinct algorithmic structures to provide diverse starting points\. After applying algorithmic refinement to each seed, a lightweight performance filter evaluates the refined seeds and retains only the best\-performing candidate for subsequent stages\.

Algorithmic Refinement\.While multi\-seed search provides diverse algorithmic directions, each seed’s initial implementation typically contains structural inefficiencies\. Algorithmic refinement applies expert\-guided, semantics\-preserving transformations to optimize the computational structure within each seed\. Given a seed kernel, the LLM analyzes its operator composition and dataflow structure, and then performs structural optimizations, such as operator fusion, algorithm reformulation, or dataflow reorganization\. These refinements reduce redundant computation and unnecessary memory accesses at the algorithmic level\.

![Refer to caption](https://arxiv.org/html/2606.26758v1/x2.png)Figure 2:Algorithmic refinementfor the Attention operator\. FlashAttention eliminates4​N24N^\{2\}memory accesses via kernel fusion\.Figure[2](https://arxiv.org/html/2606.26758#S3.F2)illustrates the importance of algorithmic refinement using the Attention operator as a representative example\. Traditional implementations decompose Attention into multiple independent kernels that computeQ​KTQK^\{T\}, apply softmax normalization, and multiply byVV\. For a sequence of lengthNN, this approach materializes theN×NN\\times Nattention matrix in global memory, incurring4​N24N^\{2\}additional memory accesses\. FlashAttention restructures this computation by introducing an online softmax mechanism that maintains normalization statistics during traversal ofKKandVV\. This design fuses attention score computation, normalization, and weighted accumulation into a single kernel, eliminating the memory bottleneck\.

Overall, the combination of multi\-seed search and algorithmic refinement balances coarse\-grained paradigm exploration with fine\-grained structural optimization, establishing a strong algorithmic foundation for subsequent hardware\-specific tuning\. We leave the prompt details in Appendix\.[E](https://arxiv.org/html/2606.26758#A5)\.

#### 3\.2\.2Hardware\-Specific Tuning

After establishing the high\-level algorithmic structure, we perform fine\-grained hardware\-specific tuning through three progressive stages:*parallel mapping*, which determines the parallel strategy to fully utilize GPU cores;*tensor tiling*, which selects tile\-level granularity to maximize on\-chip data reuse; and*memory optimization*, which optimizes data access patterns and execution pipelines\. Each stage constrains the optimization space of subsequent stages\. Parallel mapping fixes the global parallel structure; tensor tiling operates within each GPU execution grid under this structure; and memory optimization further refines execution under fixed parallelization and tiling\. By decomposing the joint optimization problem into a sequence of constrained sub\-problems, we provide the LLM with well\-scoped objectives at each stage\.

Parallel Mapping\.The parallel mapping stage determines how computational tasks are distributed across the GPU execution grid\. The objective is to identify parallelizable dimensions of the operator \(e\.g\., batch, head, or expert dimensions\) and map them to the GPU execution grid to fully utilize parallelism\. Formally, we define parallel mapping as a dimension\-to\-grid mapping:

\(d1,d2,…,dn\)→\(G1,G2,…,Gn\),\(d\_\{1\},d\_\{2\},\\ldots,d\_\{n\}\)\\rightarrow\(G\_\{1\},G\_\{2\},\\ldots,G\_\{n\}\),\(1\)wheredid\_\{i\}denotes the size of theii\-th operator dimension andGiG\_\{i\}denotes the size of theii\-th grid dimension\.

Tensor Tiling\.After determining the global parallel structure, the tensor tiling stage focuses on determining the size of tensor tiles processed within each execution grid\. We formalize tensor tiling as a dimension\-to\-tile mapping:

\(d1,d2,…,dn\)→\(B1,B2,…,Bn\),\(d\_\{1\},d\_\{2\},\\ldots,d\_\{n\}\)\\rightarrow\(B\_\{1\},B\_\{2\},\\ldots,B\_\{n\}\),\(2\)whereBiB\_\{i\}denotes the corresponding tile size\. The relationship between parallel mapping and tiling is given byGi=⌈di/Bi⌉G\_\{i\}=\\lceil d\_\{i\}/B\_\{i\}\\rceil\.

Tensor tiling balances computation efficiency and memory behavior\. Larger tiles improve on\-chip data reuse but incur higher register and shared memory consumption, potentially limiting parallelism\. In contrast, smaller tiles reduce per\-instance resource usage but often fail to exploit data locality, leading to increased memory access overhead\. As a result, tile size selection is inherently a constraint\-aware design problem: the choice of\(B1,…,Bn\)\(B\_\{1\},\\ldots,B\_\{n\}\)must respect hardware resource limits while achieving high computational throughput\. In EGG, the LLM is guided to propose a small set of tiling candidates, from which the final configuration is selected based on runtime performance measurements\.

Memory Optimization\.Under fixed parallelization and tiling, the stage\-specific prompt directs the agent to optimize memory access patterns by organizing global memory accesses into coalesced loads\. It also improves software pipelining by adjusting multi\-buffering depth to overlap data movement with computation and hide memory access latency\.

### 3\.3Stage\-Aware Multi\-Agent Collaboration

After decomposing the optimization process into distinct stages, we introduce a stage\-aware multi\-agent collaboration mechanism to perform optimization within each stage\. A single agent responsible for code generation, performance analysis, and debugging across all stages often suffers from*objective drift*, where accumulated context obscures the current optimization goal and leads to redundant revisions or regression of prior optimizations\. To mitigate this issue, EGG adopts a collaborative multi\-agent design with structured context management, enabling cumulative and stable optimization throughout the staged workflow\.

#### 3\.3\.1Multi\-Agent Design

![Refer to caption](https://arxiv.org/html/2606.26758v1/x3.png)Figure 3:Example of intra\-stage multi\-agent information exchange\. The profile and debug agents pass feedback to the code agent via structured JSON outputs\.The multi\-agent system decomposes each optimization stage into functionally distinct roles\. As shown in Figure[1](https://arxiv.org/html/2606.26758#S1.F1), three specialized agents form a closed\-loop collaboration around a unified optimization objective:

- •Profile Agent:analyzes runtime profiling metrics \(e\.g\., NVIDIA Nsight Compute reports\) together with the current kernel code and stage objective, identifies the primary performance bottleneck \(e\.g\., compute\-bound, memory\-bound\), and proposes targeted modification plans to the code agent\.
- •Code Agent:generates the revised kernel according to the modification plan from the profile or debug agent, given the current kernel code and stage objective\. Based on runtime results, this kernel is dispatched either to the debug agent for further repair or to the profile agent in the next stage for optimization\.
- •Debug Agent:diagnoses compilation errors, runtime exceptions, or numerical inconsistencies when failures occur, based on error logs and the current kernel code, and outputs targeted fixes to the code agent\.

#### 3\.3\.2Structured Context Management

Context management governs how optimization context is organized, filtered, and shared to ensure effective collaboration among agents\. We introduce two complementary mechanisms:*inter\-stage context propagation*, which manages context flow across stages to support cumulative optimization, and*intra\-stage information exchange*, which coordinates agent interactions within each stage to ensure stable optimization behavior\.

Table 1:Performance comparison on KernelBench\. We report success rate,Fast1\\mathrm\{Fast\}\_\{1\}rate, and mean speedup over PyTorch Eager across three difficulty levels\.MethodLevel 1Level 2Level 3Success\(%\)Fast1\\mathrm\{Fast\}\_\{1\}\(%\)SpeedupSuccessFast1\\mathrm\{Fast\}\_\{1\}SpeedupSuccessFast1\\mathrm\{Fast\}\_\{1\}SpeedupTorch Compile100%72%1\.09×\\times100%84%1\.38×\\times100%92%1\.36×\\timesDeepseek V3\.234%11%0\.99×\\times40%17%0\.88×\\times36%14%0\.84×\\timesChatGPT 5\.160%18%0\.90×\\times60%32%1\.13×\\times66%24%0\.91×\\timesAutoTriton36%9%1\.20×\\times55%22%0\.96×\\times56%26%0\.83×\\timesCudaForge100%56%1\.43×\\times100%90%2\.00×\\times100%72%1\.30×\\timesOurs100%72%1\.83×\\times100%100%2\.73×\\times100%94%1\.52×\\times

Inter\-Stage Context Propagation\.At stage boundaries, we filter and reorganize context to avoid cross\-stage interference\. When transitioning from stagetttot\+1t\+1, the system retains only finalized decisions and kernel implementations from stagett, discards intermediate exploratory outputs, and constructs a new context view that explicitly defines the optimization objective for staget\+1t\+1\. This mechanism establishes a cumulative optimization trajectory without regressing prior improvements\.

Intra\-Stage Information Exchange\.Intra\-stage information exchange coordinates efficient collaboration among agents through structured JSON interfaces that compress and isolate context\. The profile agent provides the bottleneck analysis and a modification plan\. When failures occur, the debug agent reports critical issues and the corresponding required fixes\. The code agent consumes structured feedback along with the current kernel code and the stage\-specific objective to generate code modifications\. Figure[3](https://arxiv.org/html/2606.26758#S3.F3)illustrates this context flow with an example\. This structured interface design enables tight collaboration loops across different agents\.

Overall, combining multi\-agent collaboration with structured context flow management enables cumulative optimization across stages while maintaining stable, focused exploration within each stage\.

## 4Experiments

In this section, we comprehensively evaluateEGGthrough systematic experiments on Triton kernel generation to analyze its effectiveness, robustness, and optimization behavior\.

### 4\.1Experimental Setup

Hardware Platforms\.We show the results performing on NVIDIA GeForce RTX 4090 with 24 GB GDDR6X memory, Ada Lovelace architecture, 128 SMs, and 16,384 CUDA cores; To validate generality across different hardware, we report additional results on RTX 5090, H20, and RTX PRO 6000 GPUs in Appendix[A](https://arxiv.org/html/2606.26758#A1)\.

Software\.All experiments are conducted using CUDA 13\.0, PyTorch 2\.9\.1 and Triton 3\.5\.1\. LLM inference is primarily supported by GPT\-5\.1\. Claude Opus 4\.5 results are reported in the Appendix[A](https://arxiv.org/html/2606.26758#A1)\.

Benchmark\.We adopt KernelBench\(Ouyanget al\.,[2025](https://arxiv.org/html/2606.26758#bib.bib4)\)as the primary evaluation benchmark, which consists of 250 kernel tasks spanning three difficulty levels: basic operators, fused operators, and complete networks\. Details are reported in the Appendix[C](https://arxiv.org/html/2606.26758#A3)\.

Baselines\.We compare EGG against the following baselines: 1\) PyTorch Eager, the default execution mode invoking vendor\-optimized libraries \(e\.g\., cuBLAS, cuDNN\); 2\) Torch Compile\(Anselet al\.,[2024](https://arxiv.org/html/2606.26758#bib.bib28)\)\(default mode\), PyTorch’s graph compilation framework that applies operator fusion over pre\-built kernel libraries; 3\) ChatGPT\-5\.1\(OpenAI,[2025](https://arxiv.org/html/2606.26758#bib.bib26)\)and 4\) DeepSeek\-V3\.2\(Liuet al\.,[2025](https://arxiv.org/html/2606.26758#bib.bib27)\), state\-of\-the\-art general\-purpose LLMs; 5\) AutoTriton\(Liet al\.,[2025c](https://arxiv.org/html/2606.26758#bib.bib5)\), a RL\-fine\-tuned LLM for Triton kernel generation; and 6\) CudaForge\(Zhanget al\.,[2025b](https://arxiv.org/html/2606.26758#bib.bib7)\), a multi\-agent framework with performance\-feedback\-driven iterative refinement\. For a fair comparison, we re\-evaluate CudaForge under the same setup as EGG\. We provide additional comparisons with compiler\-based baselines, including Torch Compile in max\-autotune mode and TVM Relax\(Fenget al\.,[2023](https://arxiv.org/html/2606.26758#bib.bib38)\), in the Appendix[B](https://arxiv.org/html/2606.26758#A2)\.

Metrics\.Following prior work\(Ouyanget al\.,[2025](https://arxiv.org/html/2606.26758#bib.bib4); Liet al\.,[2025b](https://arxiv.org/html/2606.26758#bib.bib9); Baronioet al\.,[2026](https://arxiv.org/html/2606.26758#bib.bib3)\)and the standard evaluation protocol of KernelBench, we evaluate all methods using three metrics: 1\)*Success Rate*: the fraction of tasks that successfully compile and pass correctness verification; 2\)*Fast1\\mathrm\{Fast\}\_\{1\}Rate*: the fraction of tasks for which the generated kernels are correct and outperform PyTorch Eager \(speedup\>1\.0×\>1\.0\\times\); 3\)*Speedup*: the average execution\-time improvement over PyTorch Eager, computed over correct kernels\.

### 4\.2Overall Performance

Table[1](https://arxiv.org/html/2606.26758#S3.T1)presents a comprehensive comparison across the three difficulty levels of KernelBench\. Overall, EGG consistently outperforms all baselines as task complexity increases\. It achieves a 100% success rate while delivering substantial performance improvements, with an average speedup of2\.13×2\.13\\timesover PyTorch Eager and1\.60×1\.60\\timesover the auto\-compilation baseline Torch Compile\.

Table 2:Ablation study on multi\-seed search, algorithmic refinement \(Algo Refine\), hardware\-specific tuning \(HW Tune\), and multi\-agent collaboration\.Comparison with PyTorch Eager\.On Level 1 basic operators, EGG achieves a1\.83×1\.83\\timesspeedup with a 72%Fast1\\mathrm\{Fast\}\_\{1\}rate and a 100% success rate\. As task complexity increases to Level 2 fused operators, the advantage of EGG becomes more pronounced: EGG achieves a2\.73×2\.73\\timesspeedup with a 100%Fast1\\mathrm\{Fast\}\_\{1\}rate, where every generated kernel outperforms PyTorch Eager\. This result highlights the effectiveness of expert\-guided staged optimization for handling fused operators\. On the most challenging Level 3 end\-to\-end models, EGG maintains a 100% success rate while achieving a1\.52×1\.52\\timesspeedup and a 94%Fast1\\mathrm\{Fast\}\_\{1\}rate, demonstrating robust performance under complex optimization scenarios\.

Comparison with Other Baselines\.Torch Compile, which leverages graph\-level optimizations and pre\-built kernel libraries, achieves1\.091\.09–1\.38×1\.38\\timesspeedup\. However, its reliance on existing kernel implementations limits its effectiveness on novel or highly fused operator patterns\. In contrast, EGG synthesizes custom kernels tailored to specific workload patterns, delivering further performance improvements across all difficulty levels\.

General\-purpose LLMs \(DeepSeek and ChatGPT\) achieve only 11\-32%Fast1\\mathrm\{Fast\}\_\{1\}rates with 34–66% success rates due to the lack of domain expertise\. The RL\-based AutoTriton attempts to address this limitation by fine\-tuning the model with execution\-time reward signals\. However, due to the scarcity of high\-quality kernels, effective exploration of the optimization space remains challenging\.

CudaForge employs multi\-agent iterative refinement driven by hardware feedback and achieves1\.301\.30\-2\.00×2\.00\\timesspeedup\. However, without expert guidance on the optimization direction, it relies on trial\-and\-error exploration, limiting its effectiveness on complex tasks \(e\.g\., 72%Fast1\\mathrm\{Fast\}\_\{1\}rate on Level 3\)\. In contrast, EGG guides kernel generation with expert optimization principles, enabling higher performance\.

Cost Efficiency\.Under the same GPT\-5\.1 and RTX 4090 setup, EGG completes kernel generation for a single task in approximately 20 minutes, consuming around 50,000 output tokens per kernel\. In comparison, CudaForge takes approximately 30 minutes and consumes around 110,000 output tokens per kernel under the same setup\. This result demonstrates that decomposing kernel optimization into stage\-wise objectives improves not only final kernel quality but also search efficiency\.

### 4\.3Ablation Study

![Refer to caption](https://arxiv.org/html/2606.26758v1/x4.png)Figure 4:Average cumulative speedup across four expert\-guided optimization stages\.As shown in Table[2](https://arxiv.org/html/2606.26758#S4.T2), we systematically analyze the contribution of individual components inEGGthrough a set of ablation studies\.

Effect of Multi\-Seed Search\.Introducing multi\-seed search improves the average speedup from1\.84×1\.84\\timesto2\.13×2\.13\\timesand increases theFast1\\mathrm\{Fast\}\_\{1\}rate from 78\.4% to 87\.6%\. The gains vary with task complexity: Level 2 and Level 3 achieve additional speedups of1\.7×1\.7\\timesand2\.0×2\.0\\times, respectively, whereas Level 1 shows only a modest improvement of1\.09×1\.09\\times\. This trend arises because complex tasks expose larger algorithmic design spaces, allowing broader exploration, while simple operators tend to converge to similar implementations\. Overall, multi\-seed search mitigates early convergence to suboptimal structures, with benefits increasing as task complexity grows\.

Stage\-Specific Contributions\.Enabling only the algorithmic refinement stage yields a1\.52×1\.52\\timesspeedup but achieves a moderateFast1\\mathrm\{Fast\}\_\{1\}rate of 60\.8%\. This suggests that while algorithmic refinement can unlock high performance potential, it does not consistently deliver speedups across kernels without hardware\-specific optimization\. Conversely, enabling only hardware\-specific tuning achieves a1\.46×1\.46\\timesspeedup with a higherFast1\\mathrm\{Fast\}\_\{1\}rate of 72%, providing more stable acceleration across tasks but limited peak performance due to suboptimal initial algorithmic choices\. Combining both components achieves the best overall performance, demonstrating their complementary roles: algorithmic refinement establishes the performance ceiling, while hardware\-specific tuning ensures reliable realization of that potential\.

Multi\-Agent Contributions\.Using only a single code agent limits the system’s ability to identify performance bottlenecks and correctness issues, leading to degraded performance and lower success rates\. This result highlights the importance of stage\-aware multi\-agent collaboration, where coordination among specialized agents enables consistent and stable performance improvements\.

Cumulative Effects Across Stages\.Figure[4](https://arxiv.org/html/2606.26758#S4.F4)illustrates the average cumulative performance improvement across four expert\-guided optimization stages\. The algorithmic refinement stage achieves an initial2\.0×2\.0\\timesspeedup over the raw seed implementation, effectively establishing a strong performance baseline\. Subsequent hardware\-specific tuning stages contribute an additional1\.7×1\.7\\timesimprovement through parallel mapping \(1\.33×1\.33\\times\), tensor tiling \(1\.08×1\.08\\times\), and memory optimization \(1\.13×1\.13\\times\)\. These results indicate that EGG achieves high performance through the accumulation of improvements across different expert\-guided stages\.

### 4\.4Case Study: 3D Transposed Convolution

Table 3:Optimization trajectory for grouped 3D transposed convolution: latency and speedup relative to PyTorch Eager\.We demonstrate the optimization workflow ofEGGusing a grouped 3D transposed convolution operator with batch size 16, 32 channels, kernel size\(3,5,7\)\(3,5,7\), stride\(2,2,2\)\(2,2,2\), and 4 groups\. As shown in Table[3](https://arxiv.org/html/2606.26758#S4.T3), the PyTorch Eager baseline executes this operator in 20\.78 ms\.

Algorithmic Structure Design\.*Multi\-Seed Search\.*The multi\-seed search generates two initial implementations with distinct algorithmic structures\. Seed 1 adopts a backward\-mapping strategy, achieving a latency of 29\.62 ms \(0\.7×0\.7\\timesspeedup\)\. Seed 2 employs nested iteration over input elements and kernel dimensions, resulting in a significantly higher latency of 76\.04 ms \(0\.27×0\.27\\timesspeedup\)\. Both initial implementations underperform the PyTorch Eager baseline\.

*Algorithmic Refinement\.*For Seed 1, the profile agent identifies that backward mapping requires expensive division and modulo operations when enumerating theCI×KD×KH×KWC\_\{I\}\\times K\_\{D\}\\times K\_\{H\}\\times K\_\{W\}combinations, leading to control\-flow overhead and warp divergence\. The code agent restructures the kernel to a forward\-mapping formulation with in\-block accumulation, reducing the latency to 19\.41 ms \(1\.07×1\.07\\times\)\. For Seed 2, the profile agent diagnoses low arithmetic intensity as the primary bottleneck\. By transforming the transposed convolution into anim2col\-based matrix multiplication, the refined kernel reduces latency to 12\.38 ms \(1\.68×1\.68\\times\)\. This refined implementation is selected for subsequent hardware\-specific tuning\.

Hardware\-Specific Tuning\.*Parallel mapping\.*NCU profiling reveals low streaming multiprocessor \(SM\) utilization\. Guided by the profile agent, the code agent remaps output pixels, output channels, and groups across threeprogram\_idaxes, improving workload distribution across SMs and reducing latency to 10\.78 ms \(1\.93×1\.93\\times\)\.

*Tensor tiling\.*Further analysis of register and shared memory usage guides autotuning overBLOCKM/N/Kconfigurations\. Among\(32,32,32\)\(32,32,32\),\(64,32,32\)\(64,32,32\), and\(32,64,32\)\(32,64,32\), the optimal configuration achieves a latency of 9\.20 ms, corresponding to a2\.26×2\.26\\timesspeedup over PyTorch Eager\.

Table[3](https://arxiv.org/html/2606.26758#S4.T3)summarizes the complete optimization trajectory\. The algorithmic structure design delivers a1\.68×1\.68\\timesimprovement over PyTorch Eager, where the algorithmic refinement achieves a 6\.2×\\timesimprovement over the worst seed, establishing a strong foundation\. Subsequent hardware\-specific tuning contributes an additional cumulative1\.35×1\.35\\timesgain through parallel mapping \(1\.15×1\.15\\times\) and tensor tiling \(1\.17×1\.17\\times\), achieving a final2\.26×2\.26\\timesspeedup over PyTorch Eager and an8\.4×8\.4\\timesimprovement over the initial seed\.

### 4\.5Practical Application Verification

Table 4:Execution time comparison for representative real\-world Triton workloads\.To evaluate practical deployment effectiveness, we assess EGG on representative operators from TritonBench\(Liet al\.,[2025b](https://arxiv.org/html/2606.26758#bib.bib9)\), as shown in Table[4](https://arxiv.org/html/2606.26758#S4.T4)\. These operators are derived from production Triton kernels in GitHub repositories and commonly used in real\-world LLM workloads\.

We compare our generated kernels against the original hand\-written Triton implementations\. EGG achieves substantial speedups: 1\.24×\\timesfor Flash Attention, 1\.63×\\timesfor RoPE Embedding, and 1\.08×\\timesfor INT8 Dequant MatMul\. These results demonstrate that expert manual tuning does not always achieve optimal performance for complex operators\. EGG effectively harnesses LLMs’ exploration capabilities to discover implementations that surpass hand\-tuned production kernels, validating the practical value of our approach\.

## 5Conclusion

In this work, we propose an expert\-guided GPU kernel generation agent framework that decomposes optimization into algorithmic structure design and hardware\-specific tuning, guiding LLM decisions with expert optimization workflows\. A stage\-aware multi\-agent collaboration mechanism coordinates code, profile, and debug agents to achieve stable and cumulative improvements\. Experiments on KernelBench and real\-world workloads demonstrate an average2\.13×2\.13\\timesspeedup over PyTorch, outperforming existing agent\-based and RL\-based approaches\.

## Impact Statement

This paper presents work whose goal is to advance the field of Machine Learning\. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here\.

## References

- J\. Ansel, E\. Yang, H\. He, N\. Gimelshein, A\. Jain, M\. Voznesensky, B\. Bao, P\. Bell, D\. Berard, E\. Burovski,et al\.\(2024\)PyTorch 2: faster machine learning through dynamic Python bytecode transformation and graph compilation\.InProceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2,pp\. 929–947\.Cited by:[§4\.1](https://arxiv.org/html/2606.26758#S4.SS1.p4.1)\.
- C\. Baronio, P\. Marsella, B\. Pan, S\. Guo, and S\. Alberti \(2026\)Kevin: multi\-turn RL for generating CUDA kernels\.InThe Fourteenth International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=xu1XwVZtDi)Cited by:[§2\.3](https://arxiv.org/html/2606.26758#S2.SS3.p1.1),[§4\.1](https://arxiv.org/html/2606.26758#S4.SS1.p5.2)\.
- J\. Cai, Y\. Wei, Z\. Wu, S\. Peng, and K\. Ma \(2023\)Inter\-Layer scheduling space definition and exploration for tiled accelerators\.InProceedings of the 50th Annual International Symposium on Computer Architecture,pp\. 1–17\.Cited by:[§2\.1](https://arxiv.org/html/2606.26758#S2.SS1.p1.1)\.
- T\. Chen, T\. Moreau, Z\. Jiang, L\. Zheng, E\. Yan, H\. Shen, M\. Cowan, L\. Wang, Y\. Hu, L\. Ceze,et al\.\(2018\)TVM: an automated End\-to\-End optimizing compiler for deep learning\.In13th USENIX Symposium on Operating Systems Design and Implementation \(OSDI 18\),pp\. 578–594\.Cited by:[§1](https://arxiv.org/html/2606.26758#S1.p1.1)\.
- Y\. Cheng, L\. Wang, Y\. Shi, Y\. Xia, L\. Ma, J\. Xue, Y\. Wang, Z\. Mo, F\. Chen, F\. Yang,et al\.\(2025\)PipeThreader: software\-defined pipelining for efficient DNN execution\.In19th USENIX Symposium on Operating Systems Design and Implementation \(OSDI 25\),Cited by:[§2\.1](https://arxiv.org/html/2606.26758#S2.SS1.p1.1)\.
- T\. Dao, D\. Fu, S\. Ermon, A\. Rudra, and C\. Ré \(2022\)FlashAttention: fast and memory\-efficient exact attention with IO\-awareness\.Advances in Neural Information Processing Systems35,pp\. 16344–16359\.Cited by:[§2\.1](https://arxiv.org/html/2606.26758#S2.SS1.p2.1)\.
- S\. Feng, B\. Hou, H\. Jin, W\. Lin, J\. Shao, R\. Lai, Z\. Ye, L\. Zheng, C\. H\. Yu, Y\. Yu,et al\.\(2023\)TensorIR: an abstraction for automatic tensorized program optimization\.InProceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2,pp\. 804–817\.Cited by:[§4\.1](https://arxiv.org/html/2606.26758#S4.SS1.p4.1)\.
- Z\. V\. Fisches, S\. Paliskara, S\. Guo, A\. Zhang, J\. Spisak, C\. Cummins, H\. Leather, G\. Synnaeve, J\. Isaacson, A\. Markosyan,et al\.\(2025\)KernelLLM: making kernel development more accessible\.Note:[https://huggingface\.co/facebook/KernelLLM](https://huggingface.co/facebook/KernelLLM)Cited by:[§2\.3](https://arxiv.org/html/2606.26758#S2.SS3.p1.1)\.
- J\. He, C\. Treude, and D\. Lo \(2025\)LLM\-based multi\-agent systems for software engineering: literature review, vision, and the road ahead\.ACM Transactions on Software Engineering and Methodology34\(5\),pp\. 1–30\.Cited by:[§1](https://arxiv.org/html/2606.26758#S1.p2.1)\.
- Z\. Jia, O\. Padon, J\. Thomas, T\. Warszawski, M\. Zaharia, and A\. Aiken \(2019\)TASO: optimizing deep learning computation with automatic generation of graph substitutions\.InProceedings of the 27th ACM Symposium on Operating Systems Principles,pp\. 47–62\.Cited by:[§2\.1](https://arxiv.org/html/2606.26758#S2.SS1.p1.1)\.
- J\. Jiang, F\. Wang, J\. Shen, S\. Kim, and S\. Kim \(2026\)A survey on large language models for code generation\.ACM Transactions on Software Engineering and Methodology35\(2\),pp\. 1–72\.Cited by:[§1](https://arxiv.org/html/2606.26758#S1.p2.1)\.
- W\. Kwon, Z\. Li, S\. Zhuang, Y\. Sheng, L\. Zheng, C\. H\. Yu, J\. E\. Gonzalez, H\. Zhang, and I\. Stoica \(2023\)Efficient memory management for large language model serving with PagedAttention\.InProceedings of the 29th Symposium on Operating Systems Principles,pp\. 611–626\.Cited by:[§2\.2](https://arxiv.org/html/2606.26758#S2.SS2.p1.1)\.
- R\. T\. Lange, Q\. Sun, A\. Prasad, M\. Faldor, Y\. Tang, and D\. Ha \(2025\)Towards robust agentic CUDA kernel benchmarking, verification, and optimization\.arXiv preprint arXiv:2509\.14279\.Cited by:[§2\.3](https://arxiv.org/html/2606.26758#S2.SS3.p2.1)\.
- H\. Li, K\. Man, P\. Kanuparthy, H\. Chen, W\. Sun, S\. Tallam, C\. Zhu, K\. Zhu, and Z\. Qian \(2025a\)TritonForge: profiling\-guided framework for automated Triton kernel optimization\.arXiv preprint arXiv:2512\.09196\.Cited by:[§1](https://arxiv.org/html/2606.26758#S1.p3.1),[§2\.3](https://arxiv.org/html/2606.26758#S2.SS3.p2.1)\.
- J\. Li, S\. Li, Z\. Gao, Q\. Shi, Y\. Li, Z\. Wang, J\. Huang, W\. WangHaojie, J\. Wang, X\. Han,et al\.\(2025b\)TritonBench: benchmarking large language model capabilities for generating Triton operators\.InFindings of the Association for Computational Linguistics: ACL 2025,pp\. 23053–23066\.Cited by:[§4\.1](https://arxiv.org/html/2606.26758#S4.SS1.p5.2),[§4\.5](https://arxiv.org/html/2606.26758#S4.SS5.p1.1)\.
- S\. Li, Z\. Wang, Y\. He, Y\. Li, Q\. Shi, J\. Li, Y\. Hu, W\. Che, X\. Han, Z\. Liu,et al\.\(2025c\)AutoTriton: automatic Triton programming with reinforcement learning in LLMs\.arXiv preprint arXiv:2507\.05687\.Cited by:[§1](https://arxiv.org/html/2606.26758#S1.p3.1),[§1](https://arxiv.org/html/2606.26758#S1.p4.1),[§2\.3](https://arxiv.org/html/2606.26758#S2.SS3.p1.1),[§4\.1](https://arxiv.org/html/2606.26758#S4.SS1.p4.1)\.
- X\. Li, X\. Sun, A\. Wang, J\. Li, and C\. Shum \(2026\)CUDA\-L1: improving CUDA optimization via contrastive reinforcement learning\.InThe Fourteenth International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=igZItUbY6n)Cited by:[§2\.3](https://arxiv.org/html/2606.26758#S2.SS3.p1.1)\.
- A\. Liu, A\. Mei, B\. Lin, B\. Xue, B\. Wang, B\. Xu, B\. Wu, B\. Zhang, C\. Lin, C\. Dong,et al\.\(2025\)DeepSeek\-V3\.2: pushing the frontier of open large language models\.arXiv preprint arXiv:2512\.02556\.Cited by:[§4\.1](https://arxiv.org/html/2606.26758#S4.SS1.p4.1)\.
- NVIDIA \(2026a\)NVIDIA cuBLAS\.Note:[https://developer\.nvidia\.com/cublas](https://developer.nvidia.com/cublas)Cited by:[§2\.1](https://arxiv.org/html/2606.26758#S2.SS1.p2.1)\.
- NVIDIA \(2026b\)NVIDIA cuDNN\.Note:[https://developer\.nvidia\.com/cudnn](https://developer.nvidia.com/cudnn)Cited by:[§2\.1](https://arxiv.org/html/2606.26758#S2.SS1.p2.1)\.
- OpenAI \(2025\)GPT\-5\.1: a smarter, more conversational ChatGPT\.Note:[https://openai\.com/index/gpt\-5\-1/](https://openai.com/index/gpt-5-1/)Cited by:[§4\.1](https://arxiv.org/html/2606.26758#S4.SS1.p4.1)\.
- M\. Osama, D\. Merrill, C\. Cecka, M\. Garland, and J\. D\. Owens \(2023\)Stream\-K: work\-centric parallel decomposition for dense matrix\-matrix multiplication on the GPU\.InProceedings of the 28th ACM SIGPLAN Annual Symposium on Principles and Practice of Parallel Programming,pp\. 429–431\.Cited by:[§2\.1](https://arxiv.org/html/2606.26758#S2.SS1.p1.1)\.
- A\. Ouyang, S\. Guo, S\. Arora, A\. L\. Zhang, W\. Hu, C\. Re, and A\. Mirhoseini \(2025\)KernelBench: can LLMs write efficient GPU kernels?\.InForty\-second International Conference on Machine Learning,External Links:[Link](https://openreview.net/forum?id=yeoN1iQT1x)Cited by:[§1](https://arxiv.org/html/2606.26758#S1.p9.1),[§4\.1](https://arxiv.org/html/2606.26758#S4.SS1.p3.1),[§4\.1](https://arxiv.org/html/2606.26758#S4.SS1.p5.2)\.
- M\. A\. K\. Raiaan, M\. S\. H\. Mukta, K\. Fatema, N\. M\. Fahad, S\. Sakib, M\. M\. J\. Mim, J\. Ahmad, M\. E\. Ali, and S\. Azam \(2024\)A review on large language models: architectures, applications, taxonomies, open issues and challenges\.IEEE Access12,pp\. 26839–26874\.Cited by:[§1](https://arxiv.org/html/2606.26758#S1.p1.1)\.
- T\. Sereda, T\. S\. John, B\. Bartan, N\. Serrino, S\. Katti, and Z\. Asgar \(2025\)KForge: program synthesis for diverse AI hardware accelerators\.arXiv preprint arXiv:2511\.13274\.Cited by:[§2\.3](https://arxiv.org/html/2606.26758#S2.SS3.p2.1)\.
- B\. F\. Spector, S\. Arora, A\. Singhal, A\. Parthasarathy, D\. Y\. Fu, and C\. Re \(2025\)ThunderKittens: simple, fast, and adorable kernels\.InThe Thirteenth International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=0fJfVOSUra)Cited by:[§2\.2](https://arxiv.org/html/2606.26758#S2.SS2.p1.1)\.
- S\. Su, X\. Sun, X\. Li, A\. Wang, J\. Li, and C\. Shum \(2025\)CUDA\-L2: surpassing cuBLAS performance for matrix multiplication through reinforcement learning\.arXiv preprint arXiv:2512\.02551\.Cited by:[§2\.3](https://arxiv.org/html/2606.26758#S2.SS3.p1.1)\.
- P\. Tillet, H\. Kung, and D\. Cox \(2019\)Triton: an intermediate language and compiler for tiled neural network computations\.InProceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages,pp\. 10–19\.Cited by:[§2\.2](https://arxiv.org/html/2606.26758#S2.SS2.p1.1)\.
- J\. Wang, V\. Joshi, S\. Majumder, X\. Chao, B\. Ding, Z\. Liu, P\. P\. Brahma, D\. Li, Z\. Liu, and E\. Barsoum \(2025\)GEAK: introducing Triton kernel AI agent & evaluation benchmarks\.arXiv preprint arXiv:2507\.23194\.Cited by:[§1](https://arxiv.org/html/2606.26758#S1.p4.1),[§2\.3](https://arxiv.org/html/2606.26758#S2.SS3.p2.1)\.
- L\. Wang, Y\. Cheng, Y\. Shi, Z\. Mo, Z\. Tang, W\. Xie, T\. Wu, L\. Ma, Y\. Xia, J\. Xue, F\. Yang, and Z\. Yang \(2026\)TileLang: bridge programmability and performance in modern neural kernels\.InThe Fourteenth International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=Jb1WkNSfUB)Cited by:[§2\.2](https://arxiv.org/html/2606.26758#S2.SS2.p1.1)\.
- J\. Woo, S\. Zhu, A\. Nie, Z\. Jia, Y\. Wang, and Y\. Park \(2025\)TritonRL: training LLMs to think and code Triton without cheating\.arXiv preprint arXiv:2510\.17891\.Cited by:[§1](https://arxiv.org/html/2606.26758#S1.p3.1),[§2\.3](https://arxiv.org/html/2606.26758#S2.SS3.p1.1)\.
- Y\. Zhai, S\. Yang, K\. Pan, R\. Zhang, S\. Liu, C\. Liu, Z\. Ye, J\. Ji, J\. Zhao, Y\. Zhang,et al\.\(2024\)Enabling tensor language model to assist in generating High\-Performance tensor programs for deep learning\.In18th USENIX Symposium on Operating Systems Design and Implementation \(OSDI 24\),pp\. 289–305\.Cited by:[§1](https://arxiv.org/html/2606.26758#S1.p2.1)\.
- K\. Zhang, J\. Li, G\. Li, X\. Shi, and Z\. Jin \(2024\)CodeAgent: enhancing code generation with tool\-integrated agent systems for real\-world repo\-level coding challenges\.InProceedings of the 62nd Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\),pp\. 13643–13658\.Cited by:[§1](https://arxiv.org/html/2606.26758#S1.p2.1)\.
- X\. Zhang, S\. Peng, Q\. Zhou, Y\. Wen, Q\. Guo, R\. Chen, X\. Zhu, W\. Xiong, H\. Chen, C\. Ma,et al\.\(2025a\)QiMeng\-TensorOp: one\-line prompt is enough for high\-performance tensor operator generation with hardware primitives\.InProceedings of the Thirty\-Fourth International Joint Conference on Artificial Intelligence,pp\. 7038–7046\.Cited by:[§1](https://arxiv.org/html/2606.26758#S1.p3.1)\.
- Z\. Zhang, R\. Wang, Y\. Luo, S\. Li, M\. Hong, and C\. Ding \(2025b\)CudaForge: an agent framework with hardware feedback for CUDA kernel optimization\.arXiv preprint arXiv:2511\.01884\.Cited by:[§1](https://arxiv.org/html/2606.26758#S1.p3.1),[§1](https://arxiv.org/html/2606.26758#S1.p4.1),[§2\.3](https://arxiv.org/html/2606.26758#S2.SS3.p2.1),[§4\.1](https://arxiv.org/html/2606.26758#S4.SS1.p4.1)\.
- L\. Zheng, L\. Yin, Z\. Xie, C\. Sun, J\. Huang, C\. H\. Yu, S\. Cao, C\. Kozyrakis, I\. Stoica, J\. E\. Gonzalez, C\. Barrett, and Y\. Sheng \(2024\)SGLang: efficient execution of structured language model programs\.InAdvances in Neural Information Processing Systems 37,Vol\.37,pp\. 62557–62583\.Cited by:[§2\.2](https://arxiv.org/html/2606.26758#S2.SS2.p1.1)\.

## Overview

This appendix provides additional experimental evidence and implementation details to complement the main paper\. It is organized as follows:

- •Additional Hardware and LLM Results \(Section[A](https://arxiv.org/html/2606.26758#A1)\):We report EGG’s performance on additional GPU platforms, including NVIDIA RTX 5090, H20, and RTX PRO 6000, and evaluate Claude Opus 4\.5 to demonstrate model\-agnostic effectiveness\.
- •Additional Baseline Results \(Section[B](https://arxiv.org/html/2606.26758#A2)\):We provide additional comparisons withtorch\.compileinmax\-autotunemode and TVM Relax to further validate EGG against compiler\-based baselines\.
- •Benchmark Details \(Section[C](https://arxiv.org/html/2606.26758#A3)\):We summarize the task structure and interface conventions of KernelBench, and describe the construction of TritonBench\.
- •Reproducibility Details \(Section[D](https://arxiv.org/html/2606.26758#A4)\):We summarize the fixed interaction budget, stopping criteria, deterministic failure handling, and pseudocode of the EGG optimization pipeline\.
- •Prompt Details \(Section[E](https://arxiv.org/html/2606.26758#A5)\):We provide the full prompts used in our framework, including the seed prompt, stage\-specific system prompts, and the profile/code/debug agent prompts used for optimization and repair\.
- •Nsight Compute Profiling Metrics \(Section[F](https://arxiv.org/html/2606.26758#A6)\):We list the NCU metrics collected during profiling, and explain how they guide targeted optimization decisions in our framework\.
- •Limitations \(Section[G](https://arxiv.org/html/2606.26758#A7)\):We list some limitations in our current work, which we treat as our future work for further improvement\.

## Appendix AAdditional Hardware and LLM Results

### A\.1Overall Performance Across Hardware Platforms

To validate the generality of EGG across different GPU architectures, we conduct additional experiments on NVIDIA GeForce RTX 5090, NVIDIA H20, and NVIDIA RTX PRO 6000 GPUs\. The RTX 5090 features the Blackwell architecture with 170 SMs, 21,760 CUDA cores, and 32GB GDDR7 memory\. The H20 is a Hopper\-based data\-center GPU with 96GB HBM3 memory, while the RTX PRO 6000 is a Blackwell\-based professional workstation GPU with 96GB GDDR7 memory\. Table[5](https://arxiv.org/html/2606.26758#A1.T5)summarizes the RTX 5090 results across all three KernelBench levels, while Table[6](https://arxiv.org/html/2606.26758#A1.T6)summarizes the H20 and RTX PRO 6000 results on the Level 2 fused operators\. Across these evaluations, EGG maintains high correctness and consistently achieves strong speedups over PyTorch Eager, demonstrating the architectural robustness of our framework beyond the RTX 4090 setting used in the main text\.

Table 5:Performance summary on RTX 5090 across KernelBench difficulty levels\.Table 6:Performance summary on additional GPU platforms for KernelBench Level 2 fused operators\.
### A\.2Results with Opus 4\.5

To assess the model\-agnostic nature of our agent framework, we conduct additional experiments using Claude Opus 4\.5 on the medium\-difficulty Level 2 tasks \(fused operators\) of KernelBench\. As summarized in Table[7](https://arxiv.org/html/2606.26758#A1.T7), our framework achieves consistently strong performance\. These results indicate that the proposed agent framework generalizes well across different LLMs\.

Table 7:Performance results using Opus 4\.5 on KernelBench\.

## Appendix BAdditional Baseline Results

This section provides additional baseline comparisons to further validate the effectiveness of EGG\. We includetorch\.compileinmax\-autotunemode and TVM Relax as compiler\-based baselines, complementing the main results\.

### B\.1Comparison withtorch\.compileinmax\-autotunemode

In the main experiments,torch\.compileis evaluated under itsdefaultmode\. We further evaluatetorch\.compilewithmax\-autotune, which enables more aggressive autotuning and typically provides stronger performance\. As shown in Table[8](https://arxiv.org/html/2606.26758#A2.T8), EGG still achieves clear performance gains across all KernelBench levels\.

Table 8:Comparison withtorch\.compileinmax\-autotunemode on KernelBench\.
### B\.2Comparison with TVM Relax

We also compare EGG with TVM Relax, a traditional machine learning compiler baseline\. TVM Relax applies compiler\-defined graph transformations and schedule optimizations to improve deep learning workloads\. We evaluate TVM Relax on 25 representative KernelBench tasks selected from the three difficulty levels\. EGG achieves a 1\.92×\\timesgeometric mean speedup over PyTorch Eager and a 1\.56×\\timesspeedup over TVM Relax on the same task subset\. The performance gap is smaller on Level 3, where TVM Relax benefits from graph\-level optimization\. In contrast, EGG shows clearer advantages on Level 1 and Level 2 tasks, where its gains mainly come from expert\-guided computation restructuring and workload\-specific kernel synthesis\.

Table 9:Comparison with TVM Relax on 25 representative KernelBench tasks\.

## Appendix CBenchmark Details

This section provides detailed specifications of the benchmarks used in our evaluation, including representative task examples with their PyTorch reference implementations\.

### C\.1KernelBench

KernelBench is a recently proposed benchmark specifically designed for evaluating LLM\-based GPU kernel optimization frameworks\. The benchmark comprises 250 carefully curated tasks organized into three difficulty tiers:

- •Level 1 \(Basic\):100 fundamental operators including arithmetic operations \(matrix multiplication, convolution\), element\-wise operations \(ReLU, GELU, sigmoid\), and basic reductions \(softmax, layer normalization\)\.
- •Level 2 \(Medium\):100 medium\-difficulty tasks that fuse multiple primitive operations, such as GEMM\+ReLU\+Add and Conv2D\+BatchNorm\+ReLU\.
- •Level 3 \(Hard\):50 challenging tasks implementing complete neural network architectures and complex computational patterns, including AlexNet, ResNet and Vision Transformer\.

All tasks adhere to a standardized interface design that facilitates automated evaluation\. Each task provides a reference PyTorch implementation following a consistent structure: aModelclass inheriting fromnn\.Modulewith aforward\(\)method defining the computation, aget\_inputs\(\)function generating runtime inputs, and aget\_init\_inputs\(\)function providing model initialization parameters\. This uniform interface enables reliable correctness verification through numerical comparison and consistent performance benchmarking across different implementations\.

Below we present one representative example from each difficulty level to illustrate the benchmark’s task structure and interface conventions\.

Level 1: Square Matrix Multiplication[⬇](data:text/plain;base64,aW1wb3J0IHRvcmNoCmltcG9ydCB0b3JjaC5ubiBhcyBubgoKY2xhc3MgTW9kZWwobm4uTW9kdWxlKToKICAgICIiIgogICAgU2ltcGxlIG1vZGVsIHRoYXQgcGVyZm9ybXMgYSBzaW5nbGUgc3F1YXJlIG1hdHJpeCBtdWx0aXBsaWNhdGlvbiAoQyA9IEEgKiBCKQogICAgIiIiCiAgICBkZWYgX19pbml0X18oc2VsZik6CiAgICAgICAgc3VwZXIoTW9kZWwsIHNlbGYpLl9faW5pdF9fKCkKCiAgICBkZWYgZm9yd2FyZChzZWxmLCBBOiB0b3JjaC5UZW5zb3IsIEI6IHRvcmNoLlRlbnNvcikgLT4gdG9yY2guVGVuc29yOgogICAgICAgICIiIgogICAgICAgIFBlcmZvcm1zIHRoZSBtYXRyaXggbXVsdGlwbGljYXRpb24uCgogICAgICAgIEFyZ3M6CiAgICAgICAgICAgIEEgKHRvcmNoLlRlbnNvcik6IElucHV0IG1hdHJpeCBBIG9mIHNoYXBlIChOLCBOKS4KICAgICAgICAgICAgQiAodG9yY2guVGVuc29yKTogSW5wdXQgbWF0cml4IEIgb2Ygc2hhcGUgKE4sIE4pLgoKICAgICAgICBSZXR1cm5zOgogICAgICAgICAgICB0b3JjaC5UZW5zb3I6IE91dHB1dCBtYXRyaXggQyBvZiBzaGFwZSAoTiwgTikuCiAgICAgICAgIiIiCiAgICAgICAgcmV0dXJuIHRvcmNoLm1hdG11bChBLCBCKQoKTiA9IDIwNDggKiAyCgpkZWYgZ2V0X2lucHV0cygpOgogICAgQSA9IHRvcmNoLnJhbmQoTiwgTikKICAgIEIgPSB0b3JjaC5yYW5kKE4sIE4pCiAgICByZXR1cm4gW0EsIEJdCgpkZWYgZ2V0X2luaXRfaW5wdXRzKCk6CiAgICByZXR1cm4gW10gICMgTm8gc3BlY2lhbCBpbml0aWFsaXphdGlvbiBpbnB1dHMgbmVlZGVk)importtorchimporttorch\.nnasnnclassModel\(nn\.Module\):"""Simplemodelthatperformsasinglesquarematrixmultiplication\(C=A\*B\)"""def\_\_init\_\_\(self\):super\(Model,self\)\.\_\_init\_\_\(\)defforward\(self,A:torch\.Tensor,B:torch\.Tensor\)\-\>torch\.Tensor:"""Performsthematrixmultiplication\.Args:A\(torch\.Tensor\):InputmatrixAofshape\(N,N\)\.B\(torch\.Tensor\):InputmatrixBofshape\(N,N\)\.Returns:torch\.Tensor:OutputmatrixCofshape\(N,N\)\."""returntorch\.matmul\(A,B\)N=2048\*2defget\_inputs\(\):A=torch\.rand\(N,N\)B=torch\.rand\(N,N\)return\[A,B\]defget\_init\_inputs\(\):return\[\]\#Nospecialinitializationinputsneeded

Level 2: Conv2D \+ ReLU \+ BiasAdd[⬇](data:text/plain;base64,aW1wb3J0IHRvcmNoCmltcG9ydCB0b3JjaC5ubiBhcyBubgoKY2xhc3MgTW9kZWwobm4uTW9kdWxlKToKICAgICIiIgogICAgU2ltcGxlIG1vZGVsIHRoYXQgcGVyZm9ybXMgYSBjb252b2x1dGlvbiwgYXBwbGllcyBSZUxVLCBhbmQgYWRkcyBhIGJpYXMgdGVybS4KICAgICIiIgogICAgZGVmIF9faW5pdF9fKHNlbGYsIGluX2NoYW5uZWxzLCBvdXRfY2hhbm5lbHMsIGtlcm5lbF9zaXplLCBiaWFzX3NoYXBlKToKICAgICAgICBzdXBlcihNb2RlbCwgc2VsZikuX19pbml0X18oKQogICAgICAgIHNlbGYuY29udiA9IG5uLkNvbnYyZChpbl9jaGFubmVscywgb3V0X2NoYW5uZWxzLCBrZXJuZWxfc2l6ZSkKICAgICAgICBzZWxmLmJpYXMgPSBubi5QYXJhbWV0ZXIodG9yY2gucmFuZG4oYmlhc19zaGFwZSkpCgogICAgZGVmIGZvcndhcmQoc2VsZiwgeCk6CiAgICAgICAgeCA9IHNlbGYuY29udih4KQogICAgICAgIHggPSB0b3JjaC5yZWx1KHgpCiAgICAgICAgeCA9IHggKyBzZWxmLmJpYXMKICAgICAgICByZXR1cm4geAoKYmF0Y2hfc2l6ZSA9IDEyOAppbl9jaGFubmVscyAgPSA2NApvdXRfY2hhbm5lbHMgPSAxMjgKaGVpZ2h0ID0gd2lkdGggPSAxMjgKa2VybmVsX3NpemUgPSAzCmJpYXNfc2hhcGUgPSAob3V0X2NoYW5uZWxzLCAxLCAxKQoKZGVmIGdldF9pbnB1dHMoKToKICAgIHJldHVybiBbdG9yY2gucmFuZChiYXRjaF9zaXplLCBpbl9jaGFubmVscywgaGVpZ2h0LCB3aWR0aCldCgpkZWYgZ2V0X2luaXRfaW5wdXRzKCk6CiAgICByZXR1cm4gW2luX2NoYW5uZWxzLCBvdXRfY2hhbm5lbHMsIGtlcm5lbF9zaXplLCBiaWFzX3NoYXBlXQ==)importtorchimporttorch\.nnasnnclassModel\(nn\.Module\):"""Simplemodelthatperformsaconvolution,appliesReLU,andaddsabiasterm\."""def\_\_init\_\_\(self,in\_channels,out\_channels,kernel\_size,bias\_shape\):super\(Model,self\)\.\_\_init\_\_\(\)self\.conv=nn\.Conv2d\(in\_channels,out\_channels,kernel\_size\)self\.bias=nn\.Parameter\(torch\.randn\(bias\_shape\)\)defforward\(self,x\):x=self\.conv\(x\)x=torch\.relu\(x\)x=x\+self\.biasreturnxbatch\_size=128in\_channels=64out\_channels=128height=width=128kernel\_size=3bias\_shape=\(out\_channels,1,1\)defget\_inputs\(\):return\[torch\.rand\(batch\_size,in\_channels,height,width\)\]defget\_init\_inputs\(\):return\[in\_channels,out\_channels,kernel\_size,bias\_shape\]

Level 3: Multi\-Layer Perceptron \(MLP\)[⬇](data:text/plain;base64,aW1wb3J0IHRvcmNoCmltcG9ydCB0b3JjaC5ubiBhcyBubgppbXBvcnQgdG9yY2gubm4uZnVuY3Rpb25hbCBhcyBGCgpjbGFzcyBNb2RlbChubi5Nb2R1bGUpOgogICAgZGVmIF9faW5pdF9fKHNlbGYsIGlucHV0X3NpemUsIGxheWVyX3NpemVzLCBvdXRwdXRfc2l6ZSk6CiAgICAgICAgIiIiCiAgICAgICAgOnBhcmFtIGlucHV0X3NpemU6IFRoZSBudW1iZXIgb2YgaW5wdXQgZmVhdHVyZXMKICAgICAgICA6cGFyYW0gbGF5ZXJfc2l6ZXM6IEEgbGlzdCBvZiBpbnRzIGNvbnRhaW5pbmcgdGhlIHNpemVzIG9mIGVhY2ggaGlkZGVuIGxheWVyCiAgICAgICAgOnBhcmFtIG91dHB1dF9zaXplOiBUaGUgbnVtYmVyIG9mIG91dHB1dCBmZWF0dXJlcwogICAgICAgICIiIgogICAgICAgIHN1cGVyKE1vZGVsLCBzZWxmKS5fX2luaXRfXygpCgogICAgICAgIGxheWVycyA9IFtdCiAgICAgICAgY3VycmVudF9pbnB1dF9zaXplID0gaW5wdXRfc2l6ZQoKICAgICAgICBmb3IgbGF5ZXJfc2l6ZSBpbiBsYXllcl9zaXplczoKICAgICAgICAgICAgbGF5ZXJzLmFwcGVuZChubi5MaW5lYXIoY3VycmVudF9pbnB1dF9zaXplLCBsYXllcl9zaXplKSkKICAgICAgICAgICAgbGF5ZXJzLmFwcGVuZChubi5SZUxVKCkpCiAgICAgICAgICAgIGN1cnJlbnRfaW5wdXRfc2l6ZSA9IGxheWVyX3NpemUKCiAgICAgICAgbGF5ZXJzLmFwcGVuZChubi5MaW5lYXIoY3VycmVudF9pbnB1dF9zaXplLCBvdXRwdXRfc2l6ZSkpCgogICAgICAgIHNlbGYubmV0d29yayA9IG5uLlNlcXVlbnRpYWwoKmxheWVycykKCiAgICBkZWYgZm9yd2FyZChzZWxmLCB4KToKICAgICAgICAiIiIKICAgICAgICA6cGFyYW0geDogVGhlIGlucHV0IHRlbnNvciwgc2hhcGUgKGJhdGNoX3NpemUsIGlucHV0X3NpemUpCiAgICAgICAgOnJldHVybjogVGhlIG91dHB1dCB0ZW5zb3IsIHNoYXBlIChiYXRjaF9zaXplLCBvdXRwdXRfc2l6ZSkKICAgICAgICAiIiIKICAgICAgICByZXR1cm4gc2VsZi5uZXR3b3JrKHgpCgojIFRlc3QgY29kZQpiYXRjaF9zaXplID0gMTI4CmlucHV0X3NpemUgPSAxNjM4NApsYXllcl9zaXplcyA9IFsxNjM4NCwgMTYzODRdCm91dHB1dF9zaXplID0gODE5MgoKZGVmIGdldF9pbnB1dHMoKToKICAgIHJldHVybiBbdG9yY2gucmFuZChiYXRjaF9zaXplLCBpbnB1dF9zaXplKV0KCmRlZiBnZXRfaW5pdF9pbnB1dHMoKToKICAgIHJldHVybiBbaW5wdXRfc2l6ZSwgbGF5ZXJfc2l6ZXMsIG91dHB1dF9zaXplXQ==)importtorchimporttorch\.nnasnnimporttorch\.nn\.functionalasFclassModel\(nn\.Module\):def\_\_init\_\_\(self,input\_size,layer\_sizes,output\_size\):""":paraminput\_size:Thenumberofinputfeatures:paramlayer\_sizes:Alistofintscontainingthesizesofeachhiddenlayer:paramoutput\_size:Thenumberofoutputfeatures"""super\(Model,self\)\.\_\_init\_\_\(\)layers=\[\]current\_input\_size=input\_sizeforlayer\_sizeinlayer\_sizes:layers\.append\(nn\.Linear\(current\_input\_size,layer\_size\)\)layers\.append\(nn\.ReLU\(\)\)current\_input\_size=layer\_sizelayers\.append\(nn\.Linear\(current\_input\_size,output\_size\)\)self\.network=nn\.Sequential\(\*layers\)defforward\(self,x\):""":paramx:Theinputtensor,shape\(batch\_size,input\_size\):return:Theoutputtensor,shape\(batch\_size,output\_size\)"""returnself\.network\(x\)\#Testcodebatch\_size=128input\_size=16384layer\_sizes=\[16384,16384\]output\_size=8192defget\_inputs\(\):return\[torch\.rand\(batch\_size,input\_size\)\]defget\_init\_inputs\(\):return\[input\_size,layer\_sizes,output\_size\]

### C\.2TritonBench

TritonBench features 184 real\-world operators collected from GitHub repositories \(\>100 stars\), providing a diverse collection of Triton kernels spanning various computational patterns and optimization levels\. The dataset includes kernels for matrix operations, convolution operations, attention mechanisms, and custom computational kernels, each with corresponding unit tests\. It provides natural language descriptions as input, we convert them to PyTorch reference implementations to align with practical kernel optimization scenarios\.

Below we show an example Rotary Position Embedding \(RoPE\) operator, presenting both the original Triton implementation from GitHub \(https://github\.com/turbo\-llm/turbo\-alignment\) and our converted PyTorch reference\.

Original Triton Implementation:[⬇](data:text/plain;base64,CmltcG9ydCB0b3JjaAppbXBvcnQgdHJpdG9uCmltcG9ydCB0cml0b24ubGFuZ3VhZ2UgYXMgdGwKCkB0cml0b24uaml0CmRlZiBfdHJpdG9uX3JvcGUoCiAgICBxX3B0ciwKICAgIHFfcm93X3N0cmlkZSwKICAgIGtfcHRyLAogICAga19yb3dfc3RyaWRlLAogICAgY29zLAogICAgY29zX3Jvd19zdHJpZGUsCiAgICBzaW4sCiAgICBzaW5fcm93X3N0cmlkZSwKICAgIHNsLAogICAgYnM6IHRsLmNvbnN0ZXhwciwKICAgIG5fcWg6IHRsLmNvbnN0ZXhwciwKICAgIG5fa2g6IHRsLmNvbnN0ZXhwciwKICAgIGhkOiB0bC5jb25zdGV4cHIsCiAgICBwYWRfbl9xaDogdGwuY29uc3RleHByLAogICAgcGFkX25fa2g6IHRsLmNvbnN0ZXhwciwKICAgIHBhZF9oZDogdGwuY29uc3RleHByLAogICAgQkxPQ0tfU0laRTogdGwuY29uc3RleHByLAogICAgQkFDS1dBUkRfUEFTUzogdGwuY29uc3RleHByID0gRmFsc2UsCik6CiAgICBwaWQgPSB0bC5wcm9ncmFtX2lkKDApCgogICAgcV9wdHIgPSBxX3B0ciArIHBpZCAqIHFfcm93X3N0cmlkZQogICAga19wdHIgPSBrX3B0ciArIHBpZCAqIGtfcm93X3N0cmlkZQoKICAgIGNvc19yb3dfaWR4ID0gcGlkICUgKHNsKQogICAgY29zID0gY29zICsgY29zX3Jvd19pZHggKiBjb3Nfcm93X3N0cmlkZQogICAgc2luID0gc2luICsgY29zX3Jvd19pZHggKiBzaW5fcm93X3N0cmlkZQogICAgY29zX29mZnNldHMgPSB0bC5hcmFuZ2UoMCwgcGFkX2hkIC8vIDIpCiAgICBjb3NfbWFzayA9IGNvc19vZmZzZXRzIDwgaGQgLy8gMgogICAgY29zX3JvdyA9IHRsLmxvYWQoY29zICsgY29zX29mZnNldHMsIG1hc2s9Y29zX21hc2ssIG90aGVyPTApCiAgICBzaW5fcm93ID0gdGwubG9hZChzaW4gKyBjb3Nfb2Zmc2V0cywgbWFzaz1jb3NfbWFzaywgb3RoZXI9MCkKCiAgICBmaXJzdF9oYWxmX3Ffb2Zmc2V0cyA9IHRsLmFyYW5nZSgwLCBwYWRfbl9xaClbOiwgTm9uZV0gKiBoZCArIHRsLmFyYW5nZSgwLCBwYWRfaGQgLy8gMilbTm9uZSwgOl0KICAgIGZpcnN0X2hhbGZfa19vZmZzZXRzID0gdGwuYXJhbmdlKDAsIHBhZF9uX2toKVs6LCBOb25lXSAqIGhkICsgdGwuYXJhbmdlKDAsIHBhZF9oZCAvLyAyKVtOb25lLCA6XQogICAgZmlyc3RfcV9tYXNrID0gKHRsLmFyYW5nZSgwLCBwYWRfbl9xaClbOiwgTm9uZV0gPCBuX3FoKSAmICh0bC5hcmFuZ2UoMCwgcGFkX2hkIC8vIDIpW05vbmUsIDpdIDwgaGQgLy8gMikKICAgIGZpcnN0X2tfbWFzayA9ICh0bC5hcmFuZ2UoMCwgcGFkX25fa2gpWzosIE5vbmVdIDwgbl9raCkgJiAodGwuYXJhbmdlKDAsIHBhZF9oZCAvLyAyKVtOb25lLCA6XSA8IGhkIC8vIDIpCiAgICBxX3RpbGVfMSA9IHRsLmxvYWQocV9wdHIgKyBmaXJzdF9oYWxmX3Ffb2Zmc2V0cywgbWFzaz1maXJzdF9xX21hc2ssIG90aGVyPTApLnRvKHNpbl9yb3cuZHR5cGUpCiAgICBrX3RpbGVfMSA9IHRsLmxvYWQoa19wdHIgKyBmaXJzdF9oYWxmX2tfb2Zmc2V0cywgbWFzaz1maXJzdF9rX21hc2ssIG90aGVyPTApLnRvKHNpbl9yb3cuZHR5cGUpCgogICAgc2Vjb25kX2hhbGZfcV9vZmZzZXRzID0gZmlyc3RfaGFsZl9xX29mZnNldHMgKyAoaGQgLy8gMikKICAgIHNlY29uZF9oYWxmX2tfb2Zmc2V0cyA9IGZpcnN0X2hhbGZfa19vZmZzZXRzICsgKGhkIC8vIDIpCiAgICBzZWNvbmRfcV9tYXNrID0gZmlyc3RfcV9tYXNrCiAgICBzZWNvbmRfa19tYXNrID0gZmlyc3Rfa19tYXNrCiAgICBxX3RpbGVfMiA9IHRsLmxvYWQocV9wdHIgKyBzZWNvbmRfaGFsZl9xX29mZnNldHMsIG1hc2s9c2Vjb25kX3FfbWFzaywgb3RoZXI9MCkudG8oc2luX3Jvdy5kdHlwZSkKICAgIGtfdGlsZV8yID0gdGwubG9hZChrX3B0ciArIHNlY29uZF9oYWxmX2tfb2Zmc2V0cywgbWFzaz1zZWNvbmRfa19tYXNrLCBvdGhlcj0wKS50byhzaW5fcm93LmR0eXBlKQoKICAgIGlmIG5vdCBCQUNLV0FSRF9QQVNTOgogICAgICAgIG5ld19xX3RpbGVfMSA9IHFfdGlsZV8xICogY29zX3JvdyAtIHFfdGlsZV8yICogc2luX3JvdwogICAgICAgIHRsLnN0b3JlKHFfcHRyICsgZmlyc3RfaGFsZl9xX29mZnNldHMsIG5ld19xX3RpbGVfMSwgbWFzaz1maXJzdF9xX21hc2spCiAgICAgICAgbmV3X3FfdGlsZV8yID0gcV90aWxlXzIgKiBjb3Nfcm93ICsgcV90aWxlXzEgKiBzaW5fcm93CiAgICAgICAgdGwuc3RvcmUocV9wdHIgKyBzZWNvbmRfaGFsZl9xX29mZnNldHMsIG5ld19xX3RpbGVfMiwgbWFzaz1zZWNvbmRfcV9tYXNrKQoKICAgICAgICBuZXdfa190aWxlXzEgPSBrX3RpbGVfMSAqIGNvc19yb3cgLSBrX3RpbGVfMiAqIHNpbl9yb3cKICAgICAgICB0bC5zdG9yZShrX3B0ciArIGZpcnN0X2hhbGZfa19vZmZzZXRzLCBuZXdfa190aWxlXzEsIG1hc2s9Zmlyc3Rfa19tYXNrKQogICAgICAgIG5ld19rX3RpbGVfMiA9IGtfdGlsZV8yICogY29zX3JvdyArIGtfdGlsZV8xICogc2luX3JvdwogICAgICAgIHRsLnN0b3JlKGtfcHRyICsgc2Vjb25kX2hhbGZfa19vZmZzZXRzLCBuZXdfa190aWxlXzIsIG1hc2s9c2Vjb25kX2tfbWFzaykKICAgIGVsc2U6CiAgICAgICAgbmV3X3FfdGlsZV8xID0gcV90aWxlXzEgKiBjb3Nfcm93ICsgcV90aWxlXzIgKiBzaW5fcm93CiAgICAgICAgdGwuc3RvcmUocV9wdHIgKyBmaXJzdF9oYWxmX3Ffb2Zmc2V0cywgbmV3X3FfdGlsZV8xLCBtYXNrPWZpcnN0X3FfbWFzaykKICAgICAgICBuZXdfcV90aWxlXzIgPSBxX3RpbGVfMiAqIGNvc19yb3cgLSBxX3RpbGVfMSAqIHNpbl9yb3cKICAgICAgICB0bC5zdG9yZShxX3B0ciArIHNlY29uZF9oYWxmX3Ffb2Zmc2V0cywgbmV3X3FfdGlsZV8yLCBtYXNrPXNlY29uZF9xX21hc2spCgogICAgICAgIG5ld19rX3RpbGVfMSA9IGtfdGlsZV8xICogY29zX3JvdyArIGtfdGlsZV8yICogc2luX3JvdwogICAgICAgIHRsLnN0b3JlKGtfcHRyICsgZmlyc3RfaGFsZl9rX29mZnNldHMsIG5ld19rX3RpbGVfMSwgbWFzaz1maXJzdF9rX21hc2spCiAgICAgICAgbmV3X2tfdGlsZV8yID0ga190aWxlXzIgKiBjb3Nfcm93IC0ga190aWxlXzEgKiBzaW5fcm93CiAgICAgICAgdGwuc3RvcmUoa19wdHIgKyBzZWNvbmRfaGFsZl9rX29mZnNldHMsIG5ld19rX3RpbGVfMiwgbWFzaz1zZWNvbmRfa19tYXNrKQoKCmRlZiByb3BlX2ZvcndhcmQocSwgaywgY29zLCBzaW4pOgogICAgcSA9IHEudHJhbnNwb3NlKDEsIDIpCiAgICBrID0gay50cmFuc3Bvc2UoMSwgMikKCiAgICBiYXRjaF9zaXplLCBzZXFfbGVuLCBuX3FfaGVhZCwgaGVhZF9kaW0gPSBxLnNoYXBlCiAgICBuX2t2X2hlYWQgPSBrLnNoYXBlWzJdCiAgICBwYWRfaGQgPSB0cml0b24ubmV4dF9wb3dlcl9vZl8yKGhlYWRfZGltKQogICAgcGFkX25fcV9oZWFkID0gdHJpdG9uLm5leHRfcG93ZXJfb2ZfMihuX3FfaGVhZCkKICAgIHBhZF9uX2t2X2hlYWQgPSB0cml0b24ubmV4dF9wb3dlcl9vZl8yKG5fa3ZfaGVhZCkKICAgIEJMT0NLX1NJWkUgPSBtYXgocGFkX25fcV9oZWFkLCBwYWRfbl9rdl9oZWFkKQoKICAgIG5fcm93ID0gYmF0Y2hfc2l6ZSAqIHNlcV9sZW4KCiAgICBxID0gcS5jb250aWd1b3VzKCkKICAgIGsgPSBrLmNvbnRpZ3VvdXMoKQogICAgY29zID0gY29zLmNvbnRpZ3VvdXMoKQogICAgc2luID0gc2luLmNvbnRpZ3VvdXMoKQoKICAgIF90cml0b25fcm9wZVsobl9yb3csKV0oCiAgICAgICAgcSwKICAgICAgICBxLnN0cmlkZSgxKSwKICAgICAgICBrLAogICAgICAgIGsuc3RyaWRlKDEpLAogICAgICAgIGNvcywKICAgICAgICBjb3Muc3RyaWRlKC0yKSwKICAgICAgICBzaW4sCiAgICAgICAgc2luLnN0cmlkZSgtMiksCiAgICAgICAgc2VxX2xlbiwKICAgICAgICBiYXRjaF9zaXplLAogICAgICAgIG5fcV9oZWFkLAogICAgICAgIG5fa3ZfaGVhZCwKICAgICAgICBoZWFkX2RpbSwKICAgICAgICBwYWRfbl9xX2hlYWQsCiAgICAgICAgcGFkX25fa3ZfaGVhZCwKICAgICAgICBwYWRfaGQsCiAgICAgICAgQkxPQ0tfU0laRT1CTE9DS19TSVpFLAogICAgICAgIEJBQ0tXQVJEX1BBU1M9RmFsc2UsCiAgICApCiAgICByZXR1cm4gcS50cmFuc3Bvc2UoMSwgMiksIGsudHJhbnNwb3NlKDEsIDIpLCBjb3MsIHNpbgoKaW1wb3J0IHRvcmNoCgpkZWYgdGVzdF9yb3BlX2ZvcndhcmQoKToKICAgICMgRGVmaW5lIHRoZSB0ZXN0IHBhcmFtZXRlcnMKICAgIGJhdGNoX3NpemUgPSAyCiAgICBzZXFfbGVuID0gNAogICAgbl9xX2hlYWQgPSA4CiAgICBuX2t2X2hlYWQgPSA4CiAgICBoZWFkX2RpbSA9IDE2CgogICAgIyBDcmVhdGUgcmFuZG9tIGlucHV0IHRlbnNvcnMKICAgIHEgPSB0b3JjaC5yYW5kbihiYXRjaF9zaXplLCBuX3FfaGVhZCwgc2VxX2xlbiwgaGVhZF9kaW0sIGR0eXBlPXRvcmNoLmZsb2F0MzIsIGRldmljZT0nY3VkYScpCiAgICBrID0gdG9yY2gucmFuZG4oYmF0Y2hfc2l6ZSwgbl9rdl9oZWFkLCBzZXFfbGVuLCBoZWFkX2RpbSwgZHR5cGU9dG9yY2guZmxvYXQzMiwgZGV2aWNlPSdjdWRhJykKICAgIGNvcyA9IHRvcmNoLnJhbmRuKHNlcV9sZW4sIGhlYWRfZGltIC8vIDIsIGR0eXBlPXRvcmNoLmZsb2F0MzIsIGRldmljZT0nY3VkYScpCiAgICBzaW4gPSB0b3JjaC5yYW5kbihzZXFfbGVuLCBoZWFkX2RpbSAvLyAyLCBkdHlwZT10b3JjaC5mbG9hdDMyLCBkZXZpY2U9J2N1ZGEnKQoKICAgICMgRGljdGlvbmFyeSB0byBzdG9yZSByZXN1bHRzIGZvciBlYWNoIHRlc3QgY2FzZQogICAgcmVzdWx0cyA9IHt9CgogICAgIyBUZXN0IGNhc2UgMTogRm9yd2FyZCBwYXNzCiAgICBxX291dF8xLCBrX291dF8xLCBjb3Nfb3V0XzEsIHNpbl9vdXRfMSA9IHJvcGVfZm9yd2FyZChxLCBrLCBjb3MsIHNpbikKICAgIHJlc3VsdHNbJ3Rlc3RfY2FzZV8xJ10gPSAocV9vdXRfMSwga19vdXRfMSwgY29zX291dF8xLCBzaW5fb3V0XzEpCgogICAgIyBUZXN0IGNhc2UgMjogQmFja3dhcmQgcGFzcwogICAgcV9vdXRfMiwga19vdXRfMiwgY29zX291dF8yLCBzaW5fb3V0XzIgPSByb3BlX2ZvcndhcmQocSwgaywgY29zLCBzaW4pCiAgICByZXN1bHRzWyd0ZXN0X2Nhc2VfMiddID0gKHFfb3V0XzIsIGtfb3V0XzIsIGNvc19vdXRfMiwgc2luX291dF8yKQoKICAgIHJldHVybiByZXN1bHRzCgpyZXN1bHRfZ29sZCA9IHRlc3Rfcm9wZV9mb3J3YXJkKCkK)importtorchimporttritonimporttriton\.languageastl@triton\.jitdef\_triton\_rope\(q\_ptr,q\_row\_stride,k\_ptr,k\_row\_stride,cos,cos\_row\_stride,sin,sin\_row\_stride,sl,bs:tl\.constexpr,n\_qh:tl\.constexpr,n\_kh:tl\.constexpr,hd:tl\.constexpr,pad\_n\_qh:tl\.constexpr,pad\_n\_kh:tl\.constexpr,pad\_hd:tl\.constexpr,BLOCK\_SIZE:tl\.constexpr,BACKWARD\_PASS:tl\.constexpr=False,\):pid=tl\.program\_id\(0\)q\_ptr=q\_ptr\+pid\*q\_row\_stridek\_ptr=k\_ptr\+pid\*k\_row\_stridecos\_row\_idx=pid%\(sl\)cos=cos\+cos\_row\_idx\*cos\_row\_stridesin=sin\+cos\_row\_idx\*sin\_row\_stridecos\_offsets=tl\.arange\(0,pad\_hd//2\)cos\_mask=cos\_offsets<hd//2cos\_row=tl\.load\(cos\+cos\_offsets,mask=cos\_mask,other=0\)sin\_row=tl\.load\(sin\+cos\_offsets,mask=cos\_mask,other=0\)first\_half\_q\_offsets=tl\.arange\(0,pad\_n\_qh\)\[:,None\]\*hd\+tl\.arange\(0,pad\_hd//2\)\[None,:\]first\_half\_k\_offsets=tl\.arange\(0,pad\_n\_kh\)\[:,None\]\*hd\+tl\.arange\(0,pad\_hd//2\)\[None,:\]first\_q\_mask=\(tl\.arange\(0,pad\_n\_qh\)\[:,None\]<n\_qh\)&\(tl\.arange\(0,pad\_hd//2\)\[None,:\]<hd//2\)first\_k\_mask=\(tl\.arange\(0,pad\_n\_kh\)\[:,None\]<n\_kh\)&\(tl\.arange\(0,pad\_hd//2\)\[None,:\]<hd//2\)q\_tile\_1=tl\.load\(q\_ptr\+first\_half\_q\_offsets,mask=first\_q\_mask,other=0\)\.to\(sin\_row\.dtype\)k\_tile\_1=tl\.load\(k\_ptr\+first\_half\_k\_offsets,mask=first\_k\_mask,other=0\)\.to\(sin\_row\.dtype\)second\_half\_q\_offsets=first\_half\_q\_offsets\+\(hd//2\)second\_half\_k\_offsets=first\_half\_k\_offsets\+\(hd//2\)second\_q\_mask=first\_q\_masksecond\_k\_mask=first\_k\_maskq\_tile\_2=tl\.load\(q\_ptr\+second\_half\_q\_offsets,mask=second\_q\_mask,other=0\)\.to\(sin\_row\.dtype\)k\_tile\_2=tl\.load\(k\_ptr\+second\_half\_k\_offsets,mask=second\_k\_mask,other=0\)\.to\(sin\_row\.dtype\)ifnotBACKWARD\_PASS:new\_q\_tile\_1=q\_tile\_1\*cos\_row\-q\_tile\_2\*sin\_rowtl\.store\(q\_ptr\+first\_half\_q\_offsets,new\_q\_tile\_1,mask=first\_q\_mask\)new\_q\_tile\_2=q\_tile\_2\*cos\_row\+q\_tile\_1\*sin\_rowtl\.store\(q\_ptr\+second\_half\_q\_offsets,new\_q\_tile\_2,mask=second\_q\_mask\)new\_k\_tile\_1=k\_tile\_1\*cos\_row\-k\_tile\_2\*sin\_rowtl\.store\(k\_ptr\+first\_half\_k\_offsets,new\_k\_tile\_1,mask=first\_k\_mask\)new\_k\_tile\_2=k\_tile\_2\*cos\_row\+k\_tile\_1\*sin\_rowtl\.store\(k\_ptr\+second\_half\_k\_offsets,new\_k\_tile\_2,mask=second\_k\_mask\)else:new\_q\_tile\_1=q\_tile\_1\*cos\_row\+q\_tile\_2\*sin\_rowtl\.store\(q\_ptr\+first\_half\_q\_offsets,new\_q\_tile\_1,mask=first\_q\_mask\)new\_q\_tile\_2=q\_tile\_2\*cos\_row\-q\_tile\_1\*sin\_rowtl\.store\(q\_ptr\+second\_half\_q\_offsets,new\_q\_tile\_2,mask=second\_q\_mask\)new\_k\_tile\_1=k\_tile\_1\*cos\_row\+k\_tile\_2\*sin\_rowtl\.store\(k\_ptr\+first\_half\_k\_offsets,new\_k\_tile\_1,mask=first\_k\_mask\)new\_k\_tile\_2=k\_tile\_2\*cos\_row\-k\_tile\_1\*sin\_rowtl\.store\(k\_ptr\+second\_half\_k\_offsets,new\_k\_tile\_2,mask=second\_k\_mask\)defrope\_forward\(q,k,cos,sin\):q=q\.transpose\(1,2\)k=k\.transpose\(1,2\)batch\_size,seq\_len,n\_q\_head,head\_dim=q\.shapen\_kv\_head=k\.shape\[2\]pad\_hd=triton\.next\_power\_of\_2\(head\_dim\)pad\_n\_q\_head=triton\.next\_power\_of\_2\(n\_q\_head\)pad\_n\_kv\_head=triton\.next\_power\_of\_2\(n\_kv\_head\)BLOCK\_SIZE=max\(pad\_n\_q\_head,pad\_n\_kv\_head\)n\_row=batch\_size\*seq\_lenq=q\.contiguous\(\)k=k\.contiguous\(\)cos=cos\.contiguous\(\)sin=sin\.contiguous\(\)\_triton\_rope\[\(n\_row,\)\]\(q,q\.stride\(1\),k,k\.stride\(1\),cos,cos\.stride\(\-2\),sin,sin\.stride\(\-2\),seq\_len,batch\_size,n\_q\_head,n\_kv\_head,head\_dim,pad\_n\_q\_head,pad\_n\_kv\_head,pad\_hd,BLOCK\_SIZE=BLOCK\_SIZE,BACKWARD\_PASS=False,\)returnq\.transpose\(1,2\),k\.transpose\(1,2\),cos,sinimporttorchdeftest\_rope\_forward\(\):\#Definethetestparametersbatch\_size=2seq\_len=4n\_q\_head=8n\_kv\_head=8head\_dim=16\#Createrandominputtensorsq=torch\.randn\(batch\_size,n\_q\_head,seq\_len,head\_dim,dtype=torch\.float32,device=’cuda’\)k=torch\.randn\(batch\_size,n\_kv\_head,seq\_len,head\_dim,dtype=torch\.float32,device=’cuda’\)cos=torch\.randn\(seq\_len,head\_dim//2,dtype=torch\.float32,device=’cuda’\)sin=torch\.randn\(seq\_len,head\_dim//2,dtype=torch\.float32,device=’cuda’\)\#Dictionarytostoreresultsforeachtestcaseresults=\{\}\#Testcase1:Forwardpassq\_out\_1,k\_out\_1,cos\_out\_1,sin\_out\_1=rope\_forward\(q,k,cos,sin\)results\[’test\_case\_1’\]=\(q\_out\_1,k\_out\_1,cos\_out\_1,sin\_out\_1\)\#Testcase2:Backwardpassq\_out\_2,k\_out\_2,cos\_out\_2,sin\_out\_2=rope\_forward\(q,k,cos,sin\)results\[’test\_case\_2’\]=\(q\_out\_2,k\_out\_2,cos\_out\_2,sin\_out\_2\)returnresultsresult\_gold=test\_rope\_forward\(\)

Our Converted PyTorch Reference Implementation:[⬇](data:text/plain;base64,aW1wb3J0IHRvcmNoCmltcG9ydCB0b3JjaC5ubiBhcyBubgoKY2xhc3MgTW9kZWwobm4uTW9kdWxlKToKICAgICIiIgogICAgUm9QRSAoUm90YXJ5IFBvc2l0aW9uIEVtYmVkZGluZykgLSBQeVRvcmNoIFJlZmVyZW5jZSBJbXBsZW1lbnRhdGlvbgoKICAgIFJvdGFyeSBQb3NpdGlvbiBFbWJlZGRpbmcgYXBwbGllcyBhIHJvdGF0aW9uIHRvIHRoZSBxdWVyeSBhbmQga2V5IHZlY3RvcnMKICAgIGJhc2VkIG9uIHRoZWlyIHBvc2l0aW9uIGluIHRoZSBzZXF1ZW5jZS4gVGhpcyBhbGxvd3MgdGhlIG1vZGVsIHRvIG5hdHVyYWxseQogICAgZW5jb2RlIHJlbGF0aXZlIHBvc2l0aW9ucy4KCiAgICBGb3JtdWxhOgogICAgICAgIEZvciBlYWNoIHBvc2l0aW9uLCBzcGxpdCB0aGUgZW1iZWRkaW5nIGludG8gdHdvIGhhbHZlcyBbeDEsIHgyXQogICAgICAgIEFwcGx5IHJvdGF0aW9uOiBbeDEqY29zIC0geDIqc2luLCB4Mipjb3MgKyB4MSpzaW5dCgogICAgVXNlZCBpbjogTExhTUEsIEdQVC1KLCBHUFQtTmVvWCwgUGFMTSwgYW5kIG1hbnkgbW9kZXJuIExMTXMKICAgICIiIgogICAgZGVmIF9faW5pdF9fKHNlbGYpOgogICAgICAgIHN1cGVyKE1vZGVsLCBzZWxmKS5fX2luaXRfXygpCgogICAgZGVmIGZvcndhcmQoc2VsZiwgcSwgaywgY29zLCBzaW4pOgogICAgICAgICIiIgogICAgICAgIEFwcGx5IFJvUEUgdG8gcXVlcnkgYW5kIGtleSB0ZW5zb3JzLgoKICAgICAgICBBcmdzOgogICAgICAgICAgICBxICh0b3JjaC5UZW5zb3IpOiBRdWVyeSB0ZW5zb3Igb2Ygc2hhcGUgKGJhdGNoLCBuX2hlYWRzLCBzZXFfbGVuLCBoZWFkX2RpbSkKICAgICAgICAgICAgayAodG9yY2guVGVuc29yKTogS2V5IHRlbnNvciBvZiBzaGFwZSAoYmF0Y2gsIG5faGVhZHMsIHNlcV9sZW4sIGhlYWRfZGltKQogICAgICAgICAgICBjb3MgKHRvcmNoLlRlbnNvcik6IENvc2luZSB2YWx1ZXMgb2Ygc2hhcGUgKHNlcV9sZW4sIGhlYWRfZGltLy8yKQogICAgICAgICAgICBzaW4gKHRvcmNoLlRlbnNvcik6IFNpbmUgdmFsdWVzIG9mIHNoYXBlIChzZXFfbGVuLCBoZWFkX2RpbS8vMikKCiAgICAgICAgUmV0dXJuczoKICAgICAgICAgICAgVHVwbGVbdG9yY2guVGVuc29yLCB0b3JjaC5UZW5zb3IsIHRvcmNoLlRlbnNvciwgdG9yY2guVGVuc29yXToKICAgICAgICAgICAgICAgIC0gcV9yb3RhdGVkOiBSb3RhdGVkIHF1ZXJ5IChiYXRjaCwgbl9oZWFkcywgc2VxX2xlbiwgaGVhZF9kaW0pCiAgICAgICAgICAgICAgICAtIGtfcm90YXRlZDogUm90YXRlZCBrZXkgKGJhdGNoLCBuX2hlYWRzLCBzZXFfbGVuLCBoZWFkX2RpbSkKICAgICAgICAgICAgICAgIC0gY29zOiBDb3NpbmUgdmFsdWVzICh1bmNoYW5nZWQpCiAgICAgICAgICAgICAgICAtIHNpbjogU2luZSB2YWx1ZXMgKHVuY2hhbmdlZCkKICAgICAgICAiIiIKICAgICAgICAjIFRyYW5zcG9zZSB0byAoYmF0Y2gsIHNlcV9sZW4sIG5faGVhZHMsIGhlYWRfZGltKSBmb3IgZWFzaWVyIHBvc2l0aW9uLXdpc2Ugb3BlcmF0aW9uCiAgICAgICAgcSA9IHEudHJhbnNwb3NlKDEsIDIpCiAgICAgICAgayA9IGsudHJhbnNwb3NlKDEsIDIpCgogICAgICAgIGJhdGNoX3NpemUsIHNlcV9sZW4sIG5faGVhZHMsIGhlYWRfZGltID0gcS5zaGFwZQogICAgICAgIGhhbGZfZGltID0gaGVhZF9kaW0gLy8gMgoKICAgICAgICAjIFNwbGl0IGludG8gdHdvIGhhbHZlcyBhbG9uZyBoZWFkX2RpbQogICAgICAgIHExID0gcVsuLi4sIDpoYWxmX2RpbV0gICMgRmlyc3QgaGFsZgogICAgICAgIHEyID0gcVsuLi4sIGhhbGZfZGltOl0gICMgU2Vjb25kIGhhbGYKICAgICAgICBrMSA9IGtbLi4uLCA6aGFsZl9kaW1dCiAgICAgICAgazIgPSBrWy4uLiwgaGFsZl9kaW06XQoKICAgICAgICAjIFJlc2hhcGUgY29zL3NpbiBmb3IgYnJvYWRjYXN0aW5nOiAoc2VxX2xlbiwgaGVhZF9kaW0vLzIpIC0+ICgxLCBzZXFfbGVuLCAxLCBoZWFkX2RpbS8vMikKICAgICAgICBjb3MgPSBjb3NbOnNlcV9sZW4sIDpdLnVuc3F1ZWV6ZSgwKS51bnNxdWVlemUoMikKICAgICAgICBzaW4gPSBzaW5bOnNlcV9sZW4sIDpdLnVuc3F1ZWV6ZSgwKS51bnNxdWVlemUoMikKCiAgICAgICAgIyBBcHBseSByb3RhdGlvbiB0cmFuc2Zvcm1hdGlvbgogICAgICAgICMgUm9QRSBmb3JtdWxhOiByb3RhdGVfaGFsZihbeDEsIHgyXSkgPSBbeDEqY29zIC0geDIqc2luLCB4Mipjb3MgKyB4MSpzaW5dCiAgICAgICAgcV9yb3RhdGVkID0gdG9yY2guY2F0KFsKICAgICAgICAgICAgcTEgKiBjb3MgLSBxMiAqIHNpbiwgICMgTmV3IGZpcnN0IGhhbGYKICAgICAgICAgICAgcTIgKiBjb3MgKyBxMSAqIHNpbiAgICMgTmV3IHNlY29uZCBoYWxmCiAgICAgICAgXSwgZGltPS0xKQoKICAgICAgICBrX3JvdGF0ZWQgPSB0b3JjaC5jYXQoWwogICAgICAgICAgICBrMSAqIGNvcyAtIGsyICogc2luLAogICAgICAgICAgICBrMiAqIGNvcyArIGsxICogc2luCiAgICAgICAgXSwgZGltPS0xKQoKICAgICAgICAjIFRyYW5zcG9zZSBiYWNrIHRvIChiYXRjaCwgbl9oZWFkcywgc2VxX2xlbiwgaGVhZF9kaW0pCiAgICAgICAgcV9yb3RhdGVkID0gcV9yb3RhdGVkLnRyYW5zcG9zZSgxLCAyKQogICAgICAgIGtfcm90YXRlZCA9IGtfcm90YXRlZC50cmFuc3Bvc2UoMSwgMikKCiAgICAgICAgIyBSZXR1cm4gY29zL3NpbiBhcyB3ZWxsIHRvIG1hdGNoIFRyaXRvbiBpbnRlcmZhY2UKICAgICAgICByZXR1cm4gcV9yb3RhdGVkLCBrX3JvdGF0ZWQsIGNvcy5zcXVlZXplKDApLnNxdWVlemUoMSksIHNpbi5zcXVlZXplKDApLnNxdWVlemUoMSkKCgpCQVRDSF9TSVpFID0gMgpOX0hFQURTID0gOApTRVFfTEVOID0gNApIRUFEX0RJTSA9IDE2CgpkZWYgZ2V0X2lucHV0cygpOgogICAgIiIiCiAgICBHZW5lcmF0ZSB0ZXN0IGlucHV0cyBmb3IgUm9QRS4KCiAgICBSZXR1cm5zOgogICAgICAgIExpc3QgY29udGFpbmluZyBbcSwgaywgY29zLCBzaW5dOgogICAgICAgICAgICAtIHE6IFF1ZXJ5IHRlbnNvciAoYmF0Y2gsIG5faGVhZHMsIHNlcV9sZW4sIGhlYWRfZGltKQogICAgICAgICAgICAtIGs6IEtleSB0ZW5zb3IgKGJhdGNoLCBuX2hlYWRzLCBzZXFfbGVuLCBoZWFkX2RpbSkKICAgICAgICAgICAgLSBjb3M6IENvc2luZSB2YWx1ZXMgKHNlcV9sZW4sIGhlYWRfZGltLy8yKQogICAgICAgICAgICAtIHNpbjogU2luZSB2YWx1ZXMgKHNlcV9sZW4sIGhlYWRfZGltLy8yKQogICAgIiIiCiAgICBxID0gdG9yY2gucmFuZG4oQkFUQ0hfU0laRSwgTl9IRUFEUywgU0VRX0xFTiwgSEVBRF9ESU0sIGR0eXBlPXRvcmNoLmZsb2F0MzIpCiAgICBrID0gdG9yY2gucmFuZG4oQkFUQ0hfU0laRSwgTl9IRUFEUywgU0VRX0xFTiwgSEVBRF9ESU0sIGR0eXBlPXRvcmNoLmZsb2F0MzIpCiAgICBjb3MgPSB0b3JjaC5yYW5kbihTRVFfTEVOLCBIRUFEX0RJTSAvLyAyLCBkdHlwZT10b3JjaC5mbG9hdDMyKQogICAgc2luID0gdG9yY2gucmFuZG4oU0VRX0xFTiwgSEVBRF9ESU0gLy8gMiwgZHR5cGU9dG9yY2guZmxvYXQzMikKICAgIHJldHVybiBbcSwgaywgY29zLCBzaW5dCgpkZWYgZ2V0X2luaXRfaW5wdXRzKCk6CiAgICAiIiIKICAgIEdldCBpbml0aWFsaXphdGlvbiBwYXJhbWV0ZXJzIGZvciBNb2RlbC4KCiAgICBSZXR1cm5zOgogICAgICAgIEVtcHR5IGxpc3QgKG5vIGluaXRpYWxpemF0aW9uIHBhcmFtZXRlcnMgbmVlZGVkKQogICAgIiIiCiAgICByZXR1cm4gW10=)importtorchimporttorch\.nnasnnclassModel\(nn\.Module\):"""RoPE\(RotaryPositionEmbedding\)\-PyTorchReferenceImplementationRotaryPositionEmbeddingappliesarotationtothequeryandkeyvectorsbasedontheirpositioninthesequence\.Thisallowsthemodeltonaturallyencoderelativepositions\.Formula:Foreachposition,splittheembeddingintotwohalves\[x1,x2\]Applyrotation:\[x1\*cos\-x2\*sin,x2\*cos\+x1\*sin\]Usedin:LLaMA,GPT\-J,GPT\-NeoX,PaLM,andmanymodernLLMs"""def\_\_init\_\_\(self\):super\(Model,self\)\.\_\_init\_\_\(\)defforward\(self,q,k,cos,sin\):"""ApplyRoPEtoqueryandkeytensors\.Args:q\(torch\.Tensor\):Querytensorofshape\(batch,n\_heads,seq\_len,head\_dim\)k\(torch\.Tensor\):Keytensorofshape\(batch,n\_heads,seq\_len,head\_dim\)cos\(torch\.Tensor\):Cosinevaluesofshape\(seq\_len,head\_dim//2\)sin\(torch\.Tensor\):Sinevaluesofshape\(seq\_len,head\_dim//2\)Returns:Tuple\[torch\.Tensor,torch\.Tensor,torch\.Tensor,torch\.Tensor\]:\-q\_rotated:Rotatedquery\(batch,n\_heads,seq\_len,head\_dim\)\-k\_rotated:Rotatedkey\(batch,n\_heads,seq\_len,head\_dim\)\-cos:Cosinevalues\(unchanged\)\-sin:Sinevalues\(unchanged\)"""\#Transposeto\(batch,seq\_len,n\_heads,head\_dim\)foreasierposition\-wiseoperationq=q\.transpose\(1,2\)k=k\.transpose\(1,2\)batch\_size,seq\_len,n\_heads,head\_dim=q\.shapehalf\_dim=head\_dim//2\#Splitintotwohalvesalonghead\_dimq1=q\[\.\.\.,:half\_dim\]\#Firsthalfq2=q\[\.\.\.,half\_dim:\]\#Secondhalfk1=k\[\.\.\.,:half\_dim\]k2=k\[\.\.\.,half\_dim:\]\#Reshapecos/sinforbroadcasting:\(seq\_len,head\_dim//2\)\-\>\(1,seq\_len,1,head\_dim//2\)cos=cos\[:seq\_len,:\]\.unsqueeze\(0\)\.unsqueeze\(2\)sin=sin\[:seq\_len,:\]\.unsqueeze\(0\)\.unsqueeze\(2\)\#Applyrotationtransformation\#RoPEformula:rotate\_half\(\[x1,x2\]\)=\[x1\*cos\-x2\*sin,x2\*cos\+x1\*sin\]q\_rotated=torch\.cat\(\[q1\*cos\-q2\*sin,\#Newfirsthalfq2\*cos\+q1\*sin\#Newsecondhalf\],dim=\-1\)k\_rotated=torch\.cat\(\[k1\*cos\-k2\*sin,k2\*cos\+k1\*sin\],dim=\-1\)\#Transposebackto\(batch,n\_heads,seq\_len,head\_dim\)q\_rotated=q\_rotated\.transpose\(1,2\)k\_rotated=k\_rotated\.transpose\(1,2\)\#Returncos/sinaswelltomatchTritoninterfacereturnq\_rotated,k\_rotated,cos\.squeeze\(0\)\.squeeze\(1\),sin\.squeeze\(0\)\.squeeze\(1\)BATCH\_SIZE=2N\_HEADS=8SEQ\_LEN=4HEAD\_DIM=16defget\_inputs\(\):"""GeneratetestinputsforRoPE\.Returns:Listcontaining\[q,k,cos,sin\]:\-q:Querytensor\(batch,n\_heads,seq\_len,head\_dim\)\-k:Keytensor\(batch,n\_heads,seq\_len,head\_dim\)\-cos:Cosinevalues\(seq\_len,head\_dim//2\)\-sin:Sinevalues\(seq\_len,head\_dim//2\)"""q=torch\.randn\(BATCH\_SIZE,N\_HEADS,SEQ\_LEN,HEAD\_DIM,dtype=torch\.float32\)k=torch\.randn\(BATCH\_SIZE,N\_HEADS,SEQ\_LEN,HEAD\_DIM,dtype=torch\.float32\)cos=torch\.randn\(SEQ\_LEN,HEAD\_DIM//2,dtype=torch\.float32\)sin=torch\.randn\(SEQ\_LEN,HEAD\_DIM//2,dtype=torch\.float32\)return\[q,k,cos,sin\]defget\_init\_inputs\(\):"""GetinitializationparametersforModel\.Returns:Emptylist\(noinitializationparametersneeded\)"""return\[\]

## Appendix DReproducibility Details

EGG follows the optimization pipeline in Algorithm[1](https://arxiv.org/html/2606.26758#alg1)\. Validation and benchmarking run in isolated spawned subprocesses to avoid CUDA\-context contamination\. For each task, the code agent generates two seed kernels; each seed is validated, repaired by the debug agent up to three times if needed, and benchmarked once valid\. If no seed becomes valid, the task terminates early\. For each valid seed, EGG collects NCU metrics and invokes the profile agent for one algorithmic analysis\. If the analysis returns “not worth optimizing”, the seed is kept unchanged; otherwise, the code agent performs one algorithmic refinement\. EGG then selects the best valid candidate and applies three hardware\-specific tuning stages: parallel mapping, tensor tiling, and memory optimization\. In each hardware\-specific tuning stage, the profile agent uses NCU metrics to produce a stage\-specific diagnosis before the code agent modifies the kernel\. Both algorithmic refinement and hardware\-specific tuning follow the same evaluation rule: the generated candidate is validated, repaired once by the debug agent if needed, and benchmarked once valid\. A candidate replaces the current best only if it is valid and faster; otherwise, EGG retains the previous best\.

Algorithm 1EGG optimization pipeline\.1:Task specification

TT
2:Best validated candidate, or failure

3:

𝒮←∅\\mathcal\{S\}\\leftarrow\\emptyset
4:for

i=1i=1to

22do

5:

c←CodeAgent\.GenerateSeed​\(T\)c\\leftarrow\\textsc\{CodeAgent\.GenerateSeed\}\(T\)
6:for

r=0r=0to

33do

7:if

Validate​\(c\)\\textsc\{Validate\}\(c\)succeedsthen

8:

p←Benchmark​\(c\)p\\leftarrow\\textsc\{Benchmark\}\(c\)
9:

𝒮←𝒮∪\{\(c,p\)\}\\mathcal\{S\}\\leftarrow\\mathcal\{S\}\\cup\\\{\(c,p\)\\\}
10:break

11:elseif

r<3r<3then

12:

c←DebugAgent​\(T,c,failure log\)c\\leftarrow\\textsc\{DebugAgent\}\(T,c,\\text\{failure log\}\)
13:endif

14:endfor

15:endfor

16:if

𝒮=∅\\mathcal\{S\}=\\emptysetthen

17:returnfailure

18:endif

19:

𝒞←𝒮\\mathcal\{C\}\\leftarrow\\mathcal\{S\}
20:for all

\(c,p\)∈𝒮\(c,p\)\\in\\mathcal\{S\}do

21:

m←ProfileNCU​\(c\)m\\leftarrow\\textsc\{ProfileNCU\}\(c\)
22:

a←ProfileAgent​\(T,c,p,m\)a\\leftarrow\\textsc\{ProfileAgent\}\(T,c,p,m\)
23:if

aais “not worth optimizing”then

24:continue

25:endif

26:

c′←CodeAgent\.AlgorithmicRefinement​\(T,c,a\)c^\{\\prime\}\\leftarrow\\textsc\{CodeAgent\.AlgorithmicRefinement\}\(T,c,a\)
27:if

Validate​\(c′\)\\textsc\{Validate\}\(c^\{\\prime\}\)failsthen

28:

c′←DebugAgent​\(T,c′,failure log\)c^\{\\prime\}\\leftarrow\\textsc\{DebugAgent\}\(T,c^\{\\prime\},\\text\{failure log\}\)
29:endif

30:if

Validate​\(c′\)\\textsc\{Validate\}\(c^\{\\prime\}\)succeedsthen

31:

p′←Benchmark​\(c′\)p^\{\\prime\}\\leftarrow\\textsc\{Benchmark\}\(c^\{\\prime\}\)
32:

𝒞←𝒞∪\{\(c′,p′\)\}\\mathcal\{C\}\\leftarrow\\mathcal\{C\}\\cup\\\{\(c^\{\\prime\},p^\{\\prime\}\)\\\}
33:endif

34:endfor

35:

\(b​e​s​t,pb​e​s​t\)←\(best,p\_\{best\}\)\\leftarrowbest\-performing candidate in

𝒞\\mathcal\{C\}
36:for all

s∈\{ParallelMapping,TensorTiling,MemoryOptimization\}s\\in\\\{\\textsc\{ParallelMapping\},\\textsc\{TensorTiling\},\\textsc\{MemoryOptimization\}\\\}do

37:

m←ProfileNCU​\(b​e​s​t\)m\\leftarrow\\textsc\{ProfileNCU\}\(best\)
38:

h←ProfileAgent​\(s,T,b​e​s​t,pb​e​s​t,m\)h\\leftarrow\\textsc\{ProfileAgent\}\(s,T,best,p\_\{best\},m\)
39:

c←CodeAgent\.HardwareTuning​\(s,T,b​e​s​t,h\)c\\leftarrow\\textsc\{CodeAgent\.HardwareTuning\}\(s,T,best,h\)
40:if

Validate​\(c\)\\textsc\{Validate\}\(c\)failsthen

41:

c←DebugAgent​\(T,c,failure log\)c\\leftarrow\\textsc\{DebugAgent\}\(T,c,\\text\{failure log\}\)
42:endif

43:if

Validate​\(c\)\\textsc\{Validate\}\(c\)succeedsthen

44:

p←Benchmark​\(c\)p\\leftarrow\\textsc\{Benchmark\}\(c\)
45:if

p\>pb​e​s​tp\>p\_\{best\}then

46:

\(b​e​s​t,pb​e​s​t\)←\(c,p\)\(best,p\_\{best\}\)\\leftarrow\(c,p\)
47:endif

48:endif

49:endfor

50:return

b​e​s​tbest

Failure handling is deterministic\. Compilation failures, runtime failures, and accuracy failures are sent to the debug agent with the corresponding failure log\. NCU profiling failure is non\-fatal: EGG continues with the current best validated kernel and an empty profiling block\. Malformed or truncated LLM outputs are retried once; if the retry also fails, the output is discarded and EGG retains the current best candidate\.

## Appendix EPrompt Details

This section provides the complete prompts used in our multi\-agent framework\. These prompts encode domain\-specific optimization principles and guide agents through different stages of kernel generation\.

### E\.1Seed Prompt

The seed prompt is used during the multi\-seed search stage to generate initial Triton kernel implementations\. It receives the target PyTorch operator and few\-shot examples as input\. The prompt emphasizes strict syntactic constraints and common pitfalls to ensure high initial correctness\. It outputs initial Triton kernel implementations with diverse algorithmic structures\.

[⬇](data:text/plain;base64,V3JpdGUgaGlnaC1wZXJmb3JtYW5jZSBUcml0b24ga2VybmVscyB0byByZXBsYWNlIFB5VG9yY2ggb3BlcmF0b3JzLgpHZW5lcmF0ZSB0aGUgRkFTVEVTVCBrZXJuZWwgd2hpbGUgbWFpbnRhaW5pbmcgY29ycmVjdG5lc3MuCgojIyBDUklUSUNBTCAtLSBUaGVzZSBjYXVzZSA2MCUrIG9mIGZhaWx1cmVzOgoxLiBFVkVSWSBrZXJuZWwgZnVuY3Rpb24gTVVTVCBoYXZlIGBAdHJpdG9uLmppdGAgZGVjb3JhdG9yIC0tIE1BTkRBVE9SWQoyLiBHcmlkIHNpemUgTVVTVCBiZSA+IDA6IHVzZSBgdHJpdG9uLmNkaXYoTiwgQkxPQ0spYCBvciBgbWF4KDEsIE4gLy8gQkxPQ0spYAozLiBCTE9DSyBzaXplcyBNVVNUIGJlIHBvd2VyLW9mLTIgY29uc3RleHByOiAxNiwgMzIsIDY0LCAxMjgsIDI1Ngo0LiBgdGwucHJvZ3JhbV9pZChheGlzKWAgb25seSBzdXBwb3J0cyBheGlzID0gMCwgMSwgMiAobWF4IDNEIGdyaWQpCgojIyBUcml0b24gU3ludGF4IFJ1bGVzOgotIEZvciBtYXRtdWwvY29udi9saW5lYXIgb3BzLCBwcmVmZXIgYHRsLmRvdChhLCBiLCBhbGxvd190ZjMyPVRydWUpYCBvdmVyIGVsZW1lbnQtd2lzZSBtdWx0aXBseS1hZGQKLSBObyBgY29udGludWVgLCBgYnJlYWtgLCBgcmV0dXJuYCBpbnNpZGUgbG9vcHMgLS0gdXNlIG1hc2tpbmcgaW5zdGVhZAotIE5vIHRlbnNvciBpbmRleGluZyB3aXRoIGxvb3AgdmFyczogYHhbOiwgaV1gIG9yIGB4W2ksIDpdYCBpcyBJTlZBTElECi0gTm8gdHVwbGUgdW5wYWNraW5nIGluc2lkZSBrZXJuZWw6IGBhLCBiID0gdGwubG9hZCguLi4pYCBpcyBJTlZBTElECi0gTm8gbmVzdGVkIGZ1bmN0aW9ucyBpbnNpZGUgQHRyaXRvbi5qaXQKLSBObyBQeXRob24gY29udHJvbCBmbG93IG9uIHRsLnRlbnNvciBvciBCTE9DS18qIHZhbHVlcwotIE5vIGR5bmFtaWMgYHRsLnJlc2hhcGUoKWAgb3IgdmlldyBvcGVyYXRpb25zCgojIyBNaXNzaW5nIFRyaXRvbiBGdW5jdGlvbnMgKGltcGxlbWVudCBtYW51YWxseSk6Ci0gdGwudGFuaCAtLSBgKHRsLmV4cCgyKngpIC0gMSkgLyAodGwuZXhwKDIqeCkgKyAxKWAKLSB0bC5zaWdtb2lkIC0tIGAxIC8gKDEgKyB0bC5leHAoLXgpKWAKLSB0bC5nZWx1LCB0bC5zaWx1LCB0bC5zb2Z0bWF4LCB0bC5taXNoIC0tIGltcGxlbWVudCBmcm9tIGRlZmluaXRpb24KCiMjIExvYWQvU3RvcmUgUnVsZXM6Ci0gUG9pbnRlciArIHNjYWxhciBvZmZzZXQgLS0gc2NhbGFyIHZhbHVlCi0gUG9pbnRlciArIGJsb2NrIG9mZnNldCAodmlhIHRsLmFyYW5nZSkgLS0gYmxvY2sgb2YgdmFsdWVzCi0gbWFzayBzaGFwZSBNVVNUIG1hdGNoIGRhdGEgc2hhcGUgZXhhY3RseQoKIyMgT3V0cHV0IEZvcm1hdCAoU1RSSUNUKToKMS4gSW1wb3J0czogYGltcG9ydCB0b3JjaCwgdG9yY2gubm4gYXMgbm4sIHRyaXRvbiwgdHJpdG9uLmxhbmd1YWdlIGFzIHRsYCAoYW5kIG1hdGggaWYgbmVlZGVkKQoyLiBgQHRyaXRvbi5qaXRgIGtlcm5lbChzKSAtLSBNVVNUIGhhdmUgdGhpcyBkZWNvcmF0b3IKMy4gV3JhcHBlciBmdW5jdGlvbiB3aXRoIGdyaWQgY2FsY3VsYXRpb24KNC4gYGNsYXNzIE1vZGVsTmV3KG5uLk1vZHVsZSlgIC0tIFJFUVVJUkVECgpEbyBOT1QgaW5jbHVkZTogdGVzdGluZyBjb2RlLCBgaWYgX19uYW1lX18gPT0gIl9fbWFpbl9fImAsIGdldF9pbnB1dHMsIGdldF9pbml0X2lucHV0cwoKRXhhbXBsZSBQeVRvcmNoOgonJycKJGZld19iYXNlCicnJwoKRXhhbXBsZSBUcml0b246CicnJwokZmV3X25ldwonJycKClRhcmdldDoKYGBgcHl0aG9uCiRrZXJuZWxfc3JjCmBgYAoiIiI=)Writehigh\-performanceTritonkernelstoreplacePyTorchoperators\.GeneratetheFASTESTkernelwhilemaintainingcorrectness\.\#\#CRITICAL\-\-Thesecause60%\+offailures:1\.EVERYkernelfunctionMUSThave‘@triton\.jit‘decorator\-\-MANDATORY2\.GridsizeMUSTbe\>0:use‘triton\.cdiv\(N,BLOCK\)‘or‘max\(1,N//BLOCK\)‘3\.BLOCKsizesMUSTbepower\-of\-2constexpr:16,32,64,128,2564\.‘tl\.program\_id\(axis\)‘onlysupportsaxis=0,1,2\(max3Dgrid\)\#\#TritonSyntaxRules:\-Formatmul/conv/linearops,prefer‘tl\.dot\(a,b,allow\_tf32=True\)‘overelement\-wisemultiply\-add\-No‘continue‘,‘break‘,‘return‘insideloops\-\-usemaskinginstead\-Notensorindexingwithloopvars:‘x\[:,i\]‘or‘x\[i,:\]‘isINVALID\-Notupleunpackinginsidekernel:‘a,b=tl\.load\(\.\.\.\)‘isINVALID\-Nonestedfunctionsinside@triton\.jit\-NoPythoncontrolflowontl\.tensororBLOCK\_\*values\-Nodynamic‘tl\.reshape\(\)‘orviewoperations\#\#MissingTritonFunctions\(implementmanually\):\-tl\.tanh\-\-‘\(tl\.exp\(2\*x\)\-1\)/\(tl\.exp\(2\*x\)\+1\)‘\-tl\.sigmoid\-\-‘1/\(1\+tl\.exp\(\-x\)\)‘\-tl\.gelu,tl\.silu,tl\.softmax,tl\.mish\-\-implementfromdefinition\#\#Load/StoreRules:\-Pointer\+scalaroffset\-\-scalarvalue\-Pointer\+blockoffset\(viatl\.arange\)\-\-blockofvalues\-maskshapeMUSTmatchdatashapeexactly\#\#OutputFormat\(STRICT\):1\.Imports:‘importtorch,torch\.nnasnn,triton,triton\.languageastl‘\(andmathifneeded\)2\.‘@triton\.jit‘kernel\(s\)\-\-MUSThavethisdecorator3\.Wrapperfunctionwithgridcalculation4\.‘classModelNew\(nn\.Module\)‘\-\-REQUIREDDoNOTinclude:testingcode,‘if\_\_name\_\_=="\_\_main\_\_"‘,get\_inputs,get\_init\_inputsExamplePyTorch:’’’$few\_base’’’ExampleTriton:’’’$few\_new’’’Target:‘‘‘python$kernel\_src‘‘‘"""

### E\.2Stage System Prompts for Hardware\-Specific Tuning

During hardware\-specific tuning, each optimization stage uses a specialized system prompt that defines the stage’s focus, relevant metrics, and optimization rules\. These prompts constrain agent decisions to stage\-specific objectives, preventing cross\-stage interference\.

#### E\.2\.1Grid and Parallel Mapping

This stage focuses on determining the optimal mapping between operator dimensions and GPU grid dimensions\. The prompt guides the agent to prioritize batch/head/expert parallelism before reducing block sizes, and ensures grid configurations remain within hardware limits \(maximum 3 dimensions\)\.

[⬇](data:text/plain;base64,Rm9jdXM6IEdyaWQgbGF5b3V0ICYgcGFyYWxsZWxpc20uCgpNZXRyaWNzOgotIHNtX190aHJvdWdocHV0LmF2Zy5wY3Rfb2ZfcGVha19zdXN0YWluZWRfZWxhcHNlZCAoPjYwJSkKLSBsYXVuY2hfX2dyaWRfc2l6ZQoKUnVsZXM6Ci0gMUQ6IChjZGl2KE4sIEJMT0NLKSkKLSAyRDogKGNkaXYoTSwgQkxPQ0tfTSksIGNkaXYoTiwgQkxPQ0tfTikpCi0gM0Q6IChiYXRjaCwgY2RpdihNLCBCTE9DS19NKSwgY2RpdihOLCBCTE9DS19OKSkKLSA+M0Q6IGZsYXR0ZW4gT05MWSBpbmRlcGVuZGVudCBkaW1zCi0gUHJlZmVyIGJhdGNoIC8gaGVhZCAvIGV4cGVydCAvIGdyb3VwIHBhcmFsbGVsaXNtIGJlZm9yZSBzaHJpbmtpbmcgQkxPQ0sKLSBGb3IgZ3JvdXBlZCBvcGVyYXRpb25zOiBlbnN1cmUgZ3JvdXAgZGltZW5zaW9uIGlzIGluIGdyaWQgKGUuZy4sIHByb2dyYW1faWQoMikgZm9yIGdyb3VwcykKLSBDaGFuZ2UgZ3JpZCBvbmx5IGlmIFNNIHV0aWxpemF0aW9uIGlzIGNsZWFybHkgbG93CgpTYWZldHk6Ci0gTWF4IDMgZ3JpZCBkaW1zLCBzdGF0aWMgcmFuawotIGdyaWQ9KEcwLEcxLEcyKSBtdXN0IG1hdGNoIHRsLnByb2dyYW1faWQoMC8xLzIpCi0gRm9yIGdyb3VwZWQgb3BzOiB2ZXJpZnkgZ3JvdXAgaW5kZXhpbmcgaXMgY29ycmVjdAotIElmIHVuc3VyZSBhYm91dCBjb3JyZWN0bmVzcywgZG8gTk9UIGNoYW5nZSBncmlkCgpBdXRvdHVuZToKLSBBdXRvdHVuZSBlaXRoZXIgQkxPQ0tfKiBPUiAobnVtX3dhcnBzLCBudW1fc3RhZ2VzKQotIElmIGF1dG90dW5pbmcgQkxPQ0tfKiwgdXNlIGdyaWQ9bGFtYmRhIE1FVEE6ICguLi4pCi0gTmV2ZXIgcmVkZWZpbmUgQkxPQ0tfKiBpbiBib3RoIGtlcm5lbCBhbmQgbGF1bmNoCi0gTWF4IDItMyBjb25maWdzIHRvIHJlZHVjZSBjb21waWxhdGlvbiB0aW1l)Focus:Gridlayout&parallelism\.Metrics:\-sm\_\_throughput\.avg\.pct\_of\_peak\_sustained\_elapsed\(\>60%\)\-launch\_\_grid\_sizeRules:\-1D:\(cdiv\(N,BLOCK\)\)\-2D:\(cdiv\(M,BLOCK\_M\),cdiv\(N,BLOCK\_N\)\)\-3D:\(batch,cdiv\(M,BLOCK\_M\),cdiv\(N,BLOCK\_N\)\)\-\>3D:flattenONLYindependentdims\-Preferbatch/head/expert/groupparallelismbeforeshrinkingBLOCK\-Forgroupedoperations:ensuregroupdimensionisingrid\(e\.g\.,program\_id\(2\)forgroups\)\-ChangegridonlyifSMutilizationisclearlylowSafety:\-Max3griddims,staticrank\-grid=\(G0,G1,G2\)mustmatchtl\.program\_id\(0/1/2\)\-Forgroupedops:verifygroupindexingiscorrect\-Ifunsureaboutcorrectness,doNOTchangegridAutotune:\-AutotuneeitherBLOCK\_\*OR\(num\_warps,num\_stages\)\-IfautotuningBLOCK\_\*,usegrid=lambdaMETA:\(\.\.\.\)\-NeverredefineBLOCK\_\*inbothkernelandlaunch\-Max2\-3configstoreducecompilationtime

#### E\.2\.2Tensor Tiling

This stage optimizes the granularity of computation within each thread block by selecting appropriateBLOCK\_M,BLOCK\_N, andBLOCK\_Ksizes\. The prompt enforces power\-of\-2 constraints and guides autotuning across a small set of configurations to balance data reuse with register pressure\.

[⬇](data:text/plain;base64,Rm9jdXM6IEJMT0NLX00vTi9LIHNlbGVjdGlvbi4KCk1ldHJpY3M6Ci0gc21fX3dhcnBzX2FjdGl2ZS5hdmcucGN0X29mX3BlYWtfc3VzdGFpbmVkX2FjdGl2ZSAoPjUwJSkKClJ1bGVzOgotIEJMT0NLXyogbXVzdCBiZSBwb3dlcnMgb2YgMgotIFRlbnNvciBDb3JlOiBCTE9DS19NL04gbXVsdGlwbGUgb2YgMTYsIEJMT0NLX0sgbXVsdGlwbGUgb2YgOCAocHJlZmVyZW5jZSkKLSBGUDMyOiBNL04gaW4gezMyLDY0LDEyOCwyNTZ9LCBLIGluIHsxNiwzMiw2NH0KLSBBdm9pZCBvdmVyc2l6ZWQgdGlsZXMgKG1hc2sgd2FzdGUpCi0gS2VlcCBiYXNlbGluZSB0aWxlIGlmIHVuc3VyZQoKQXV0b3R1bmU6Ci0gTWF4IDItMyBjb25maWdzIHRvIHJlZHVjZSBjb21waWxhdGlvbiB0aW1lCi0gQXV0b3R1bmUgT05MWSBvbiBAdHJpdG9uLmppdCBrZXJuZWwKIiIiLA==)Focus:BLOCK\_M/N/Kselection\.Metrics:\-sm\_\_warps\_active\.avg\.pct\_of\_peak\_sustained\_active\(\>50%\)Rules:\-BLOCK\_\*mustbepowersof2\-TensorCore:BLOCK\_M/Nmultipleof16,BLOCK\_Kmultipleof8\(preference\)\-FP32:M/Nin\{32,64,128,256\},Kin\{16,32,64\}\-Avoidoversizedtiles\(maskwaste\)\-KeepbaselinetileifunsureAutotune:\-Max2\-3configstoreducecompilationtime\-AutotuneONLYon@triton\.jitkernel""",

#### E\.2\.3Memory and Tuning

The final stage refines memory access patterns and pipeline execution\. The prompt directs the agent to tunenum\_warpsandnum\_stagesbased on occupancy and memory stall metrics, while keeping grid configuration and block sizes fixed from previous stages\.

[⬇](data:text/plain;base64,Rm9jdXM6IE1lbW9yeSBvcHRpbWl6YXRpb24gYW5kIGZpbmFsIHBhcmFtZXRlciB0dW5pbmcuCgpNZXRyaWNzOgotIGRyYW1fX3Rocm91Z2hwdXQuYXZnLnBjdF9vZl9wZWFrX3N1c3RhaW5lZF9lbGFwc2VkCi0gbHRzX190X3NlY3Rvcl9oaXRfcmF0ZS5wY3QKLSBzbXNwX193YXJwX2lzc3VlX3N0YWxsZWRfbWVtb3J5X2RlcGVuZGVuY3lfcGVyX3dhcnBfYWN0aXZlLnBjdCAoPDIwJSkKLSBzbV9fd2FycHNfYWN0aXZlLmF2Zy5wY3Rfb2ZfcGVha19zdXN0YWluZWRfYWN0aXZlCgpQYXJhbWV0ZXJzIHRvIHR1bmU6Ci0gbnVtX3N0YWdlcyBpbiB7MiwgMywgNH0KLSBudW1fd2FycHMgaW4gezQsIDh9IChiYXNlZCBvbiBvY2N1cGFuY3kpCgpSdWxlczoKLSBJbmNyZWFzZSBudW1fc3RhZ2VzIG9ubHkgaWYgbWVtb3J5IHN0YWxscyA+IDIwJQotIENoYW5nZSBudW1fd2FycHMgb25seSBpZiBvY2N1cGFuY3kgc3VnZ2VzdHMgaXQKLSBMYXJnZXIgQkxPQ0tfSyBpbXByb3ZlcyByZXVzZSBidXQgaW5jcmVhc2VzIHJlZ2lzdGVyIHByZXNzdXJlCi0gRG8gTk9UIG1vZGlmeSBncmlkIG9yIEJMT0NLIHNpemVzIChmaXhlZCBpbiBlYXJsaWVyIHN0YWdlcykKLSBEbyBub3QgcmV3cml0ZSBhY2Nlc3MgcGF0dGVybnMgd2l0aG91dCBtZXRyaWMgZXZpZGVuY2UKCkF1dG90dW5lOgotIE1heCAzLTQgY29uZmlncyBjb21iaW5pbmcgbnVtX3N0YWdlcyBhbmQgbnVtX3dhcnBzCi0gQWx3YXlzIGluY2x1ZGUgb3JpZ2luYWwgY29uZmlnIGFzIGJhc2VsaW5lCi0gUmV2ZXJ0IGlmIGdhaW4gPCAyJSBvciB1bnN0YWJsZQoiIiI=)Focus:Memoryoptimizationandfinalparametertuning\.Metrics:\-dram\_\_throughput\.avg\.pct\_of\_peak\_sustained\_elapsed\-lts\_\_t\_sector\_hit\_rate\.pct\-smsp\_\_warp\_issue\_stalled\_memory\_dependency\_per\_warp\_active\.pct\(<20%\)\-sm\_\_warps\_active\.avg\.pct\_of\_peak\_sustained\_activeParameterstotune:\-num\_stagesin\{2,3,4\}\-num\_warpsin\{4,8\}\(basedonoccupancy\)Rules:\-Increasenum\_stagesonlyifmemorystalls\>20%\-Changenum\_warpsonlyifoccupancysuggestsit\-LargerBLOCK\_Kimprovesreusebutincreasesregisterpressure\-DoNOTmodifygridorBLOCKsizes\(fixedinearlierstages\)\-DonotrewriteaccesspatternswithoutmetricevidenceAutotune:\-Max3\-4configscombiningnum\_stagesandnum\_warps\-Alwaysincludeoriginalconfigasbaseline\-Revertifgain<2%orunstable"""

### E\.3Profile Agent Prompts

The profile agent analyzes kernel performance and identifies optimization opportunities\. At different stages, it receives different system prompts to align with stage\-specific objectives\.

#### E\.3\.1Algorithmic Refinement Stage

During algorithmic refinement, the profile agent receives the PyTorch reference code, current Triton kernel implementation, NCU profiling metric, and stage\-specific prompts as input\. The prompt guides the agent to analyze high\-level algorithmic optimizations \(operator fusion, algorithm replacement, etc\.\) through structured analysis: code inspection, performance diagnosis, and root cause identification\. It outputs a JSON response specifying the identified bottleneck and the modification plan for structural transformation\.

[⬇](data:text/plain;base64,WW91IGFyZSBhIEdQVSBrZXJuZWwgb3B0aW1pemF0aW9uIGFyY2hpdGVjdC4gQW5hbHl6ZSB0aGUga2VybmVsIGFuZCBpZGVudGlmeSAqKk9ORSBoaWdoLWxldmVsIGFsZ29yaXRobWljIG9wdGltaXphdGlvbioqLgoKIyBQeVRvcmNoIFJlZmVyZW5jZQpgYGBweXRob24KJHB5dGhvbl9jb2RlCmBgYAoKIyBDdXJyZW50IFRyaXRvbiBLZXJuZWwKYGBgcHl0aG9uCiR0cml0b25fY29kZQpgYGAKCiMgTnNpZ2h0IENvbXB1dGUgTWV0cmljcwpgYGAKJE5DVV9NRVRSSUNTCmBgYAoKIyMgQW5hbHlzaXMgU3RlcHMKCjEuICoqQ29kZSBBbmFseXNpcyoqOiBDb3VudCBrZXJuZWxzLCBpZGVudGlmeSBvcGVyYXRpb25zLCBjaGVjayBmb3IgaW5lZmZpY2llbmNpZXMKMi4gKipQZXJmb3JtYW5jZSBEaWFnbm9zaXMqKjogVXNlIG1ldHJpY3MvbGF0ZW5jeSB0byBpZGVudGlmeSBib3R0bGVuZWNrIHR5cGUKMy4gKipSb290IENhdXNlKio6IENvbWJpbmUgY29kZSArIHBlcmZvcm1hbmNlIHRvIGZpbmQgdGhlIGNvcmUgaXNzdWUKCiMjIE9wdGltaXphdGlvbiBDYXRlZ29yaWVzIChwaWNrIE9ORSBpZiB3b3J0aCBvcHRpbWl6aW5nKToKCiMjIyAxLiBPcGVyYXRvciBGdXNpb24KRnVzZSBjb25zZWN1dGl2ZSBvcHMgaW50byBmZXdlciBrZXJuZWxzIHRvIHJlZHVjZSBtZW1vcnkgdHJhZmZpYyBhbmQgbGF1bmNoIG92ZXJoZWFkLgoKIyMjIDIuIEFsZ29yaXRobSBSZXBsYWNlbWVudApSZXBsYWNlIG5haXZlIGFsZ29yaXRobSB3aXRoIG9wdGltaXplZCB2YXJpYW50LgotIEZvciBBdHRlbnRpb246IEZsYXNoIEF0dGVudGlvbiwgb25saW5lIHNvZnRtYXgKLSBGb3IgQ29udm9sdXRpb246IFdpbm9ncmFkLCBpbTJjb2wKLSBGb3IgUk5OL0dSVS9MU1RNOiBQZXJzaXN0ZW50IGtlcm5lbCB3aXRoIEhZQlJJRCBjb21wdXRhdGlvbgoKIyMjIDMuIEtlcm5lbCBMYXVuY2ggUmVkdWN0aW9uCkNvbWJpbmUgbXVsdGlwbGUgc21hbGwga2VybmVscyB0byByZWR1Y2Ugb3ZlcmhlYWQuCgojIyMgNC4gTWVtb3J5IExheW91dCBPcHRpbWl6YXRpb24KVXNlIGluLXBsYWNlIG9wZXJhdGlvbnMsIGJ1ZmZlciByZXVzZSwgb3IgYmV0dGVyIGxheW91dHMuCgojIyBTaG91bGQgV2UgT3B0aW1pemU/CgpCZWZvcmUgcHJvcG9zaW5nIG9wdGltaXphdGlvbiwgZGV0ZXJtaW5lIGlmIGl0J3Mgd29ydGh3aGlsZToKLSAqKk5vdCB3b3J0aCBvcHRpbWl6aW5nKiogaWY6CiAgLSBDb2RlIGlzIGFscmVhZHkgbmVhci1vcHRpbWFsIChleHBlY3RlZCBzcGVlZHVwIDwgMTAlKQogIC0gQm90dGxlbmVjayBjYW5ub3QgYmUgYWRkcmVzc2VkIChoYXJkd2FyZSBsaW1pdGVkLCBhbHJlYWR5IG9wdGltYWwgYWxnb3JpdGhtKQogIC0gT3B0aW1pemF0aW9uIHdvdWxkIGFkZCBzaWduaWZpY2FudCBjb21wbGV4aXR5IHdpdGggbWluaW1hbCBnYWluCgotICoqV29ydGggb3B0aW1pemluZyoqIGlmOgogIC0gQ2xlYXIgYWxnb3JpdGhtaWMgaW5lZmZpY2llbmN5IGV4aXN0cyAobXVsdGlwbGUga2VybmVscywgc3Vib3B0aW1hbCBhbGdvcml0aG0pCiAgLSBFeHBlY3RlZCBzcGVlZHVwID49IDIwJQogIC0gQ29uY3JldGUgb3B0aW1pemF0aW9uIHBhdGggYXZhaWxhYmxlCgojIyBPdXRwdXQgKEpTT04pCgpgYGBqc29uCnsKICAid29ydGhfb3B0aW1pemluZyI6ICJ5ZXMvbm8iLAogICJib3R0bGVuZWNrIjogIjxSb290IGNhdXNlIGluIDEtMiBzZW50ZW5jZXM+IiwKICAibW9kaWZpY2F0aW9uIHBsYW4iOiAiPEltcGxlbWVudGF0aW9uIHN0ZXBzIGluIDItMyBzZW50ZW5jZXM+IiwKfQpgYGAKUmV0dXJuIEpTT04gb25seS4=)YouareaGPUkerneloptimizationarchitect\.Analyzethekernelandidentify\*\*ONEhigh\-levelalgorithmicoptimization\*\*\.\#PyTorchReference‘‘‘python$python\_code‘‘‘\#CurrentTritonKernel‘‘‘python$triton\_code‘‘‘\#NsightComputeMetrics‘‘‘$NCU\_METRICS‘‘‘\#\#AnalysisSteps1\.\*\*CodeAnalysis\*\*:Countkernels,identifyoperations,checkforinefficiencies2\.\*\*PerformanceDiagnosis\*\*:Usemetrics/latencytoidentifybottlenecktype3\.\*\*RootCause\*\*:Combinecode\+performancetofindthecoreissue\#\#OptimizationCategories\(pickONEifworthoptimizing\):\#\#\#1\.OperatorFusionFuseconsecutiveopsintofewerkernelstoreducememorytrafficandlaunchoverhead\.\#\#\#2\.AlgorithmReplacementReplacenaivealgorithmwithoptimizedvariant\.\-ForAttention:FlashAttention,onlinesoftmax\-ForConvolution:Winograd,im2col\-ForRNN/GRU/LSTM:PersistentkernelwithHYBRIDcomputation\#\#\#3\.KernelLaunchReductionCombinemultiplesmallkernelstoreduceoverhead\.\#\#\#4\.MemoryLayoutOptimizationUsein\-placeoperations,bufferreuse,orbetterlayouts\.\#\#ShouldWeOptimize?Beforeproposingoptimization,determineifit’sworthwhile:\-\*\*Notworthoptimizing\*\*if:\-Codeisalreadynear\-optimal\(expectedspeedup<10%\)\-Bottleneckcannotbeaddressed\(hardwarelimited,alreadyoptimalalgorithm\)\-Optimizationwouldaddsignificantcomplexitywithminimalgain\-\*\*Worthoptimizing\*\*if:\-Clearalgorithmicinefficiencyexists\(multiplekernels,suboptimalalgorithm\)\-Expectedspeedup\>=20%\-Concreteoptimizationpathavailable\#\#Output\(JSON\)‘‘‘json\{"worth\_optimizing":"yes/no","bottleneck":"<Rootcausein1\-2sentences\>","modificationplan":"<Implementationstepsin2\-3sentences\>",\}‘‘‘ReturnJSONonly\.

#### E\.3\.2Hardware\-Specific Tuning Stages

During hardware\-specific tuning, the profile agent receives the PyTorch reference code, current kernel code, stage\-specific system prompts, and NCU profiling metrics as input\. The prompt constrains analysis to the current stage’s optimization scope \(e\.g\., grid configuration, block sizes, or memory patterns\) and directs the agent to propose exactly one targeted modification based on measured bottlenecks\. It outputs a JSON response specifying the bottleneck and modification plan\.

[⬇](data:text/plain;base64,WW91IGFyZSBhIHNlbmlvciBUcml0b24ga2VybmVsIG9wdGltaXphdGlvbiBlbmdpbmVlci4gUmVhZCB0aGUgUHlUb3JjaCByZWZlcmVuY2UgY29kZSwgdGhlIGN1cnJlbnQgVHJpdG9uIGNhbmRpZGF0ZSwgYW5kIHRoZSBOc2lnaHQgQ29tcHV0ZSBtZXRyaWNzLiBUaGVuIGlkZW50aWZ5IG9uZSBoaWdoZXN0LWltcGFjdCBzcGVlZCBib3R0bGVuZWNrLCBwcm9wb3NlIG9uZSBvcHRpbWl6YXRpb24gbWV0aG9kIGFuZCBwcm9wb3NlIGEgbW9kaWZpY2F0aW9uIHBsYW4uIEJlIHN1cmdpY2FsIGFuZCBtZXRyaWNzLWRyaXZlbi4KCiMgUHlUb3JjaCBSZWZlcmVuY2UKYGBgcHl0aG9uCiRweXRob25fY29kZQpgYGAKCiMgQ3VycmVudCBUcml0b24gS2VybmVsCmBgYHB5dGhvbgokVFJJVE9OX0NPREUKYGBgCgojIEN1cnJlbnQgT3B0aW1pemF0aW9uIFN0YWdlCmBgYAokU1RBR0VfQ09OVEVYVApgYGAKCiMgTnNpZ2h0IENvbXB1dGUgTWV0cmljcwpgYGAKJE5DVV9NRVRSSUNTCmBgYAoKUnVsZXM6Ci0gUmV0dXJuICoqb25lKiogb3B0aW1pemF0aW9uIG1ldGhvZCAtLSB0aGUgbGFyZ2VzdCBleHBlY3RlZCBzcGVlZHVwLgotIEZvY3VzIG9uIFRyaXRvbi1zcGVjaWZpYyBvcHRpbWl6YXRpb25zOgogICogKipCTE9DS19NL04vSyB0dW5pbmcqKjogQWRqdXN0IHRpbGUgc2l6ZXMgdG8gb3B0aW1pemUgZGF0YSByZXVzZSBhbmQgY2FjaGUgZWZmaWNpZW5jeQogICogKipudW1fd2FycHMqKjogQ29udHJvbCBvY2N1cGFuY3kgKDIvNC84IHdhcnBzIHBlciBibG9jaykKICAqICoqbnVtX3N0YWdlcyoqOiBFbmFibGUgc29mdHdhcmUgcGlwZWxpbmluZyAoMi00IHN0YWdlcyBmb3IgbWVtb3J5LWJvdW5kIGtlcm5lbHMpCiAgKiAqKk1lbW9yeSBhY2Nlc3MgcGF0dGVybnMqKjogT3B0aW1pemUgY29hbGVzY2luZywgdXNlIHRsLnRyYW5zKCkgZm9yIGxheW91dCBjaGFuZ2VzCiAgKiAqKkdyaWQgY29uZmlndXJhdGlvbioqOiBBZGp1c3QgcHJvZ3JhbV9pZCBtYXBwaW5nIGFuZCB3b3JrbG9hZCBkaXN0cmlidXRpb24KLSBQcmVmZXIgY2hhbmdlcyB0aGF0IGRpcmVjdGx5IGFkZHJlc3MgbWVhc3VyZWQgYm90dGxlbmVja3MgZnJvbSBOQ1UgbWV0cmljczoKICAqIEhpZ2ggRFJBTSB0aHJvdWdocHV0IC0tIEluY3JlYXNlIEJMT0NLIHNpemUgZm9yIGRhdGEgcmV1c2UKICAqIExvdyBjYWNoZSBoaXQgcmF0ZSAtLSBBZGp1c3QgQkxPQ0sgc2l6ZSBmb3IgYmV0dGVyIGxvY2FsaXR5CiAgKiBMb3cgb2NjdXBhbmN5IC0tIFR1bmUgbnVtX3dhcnBzLCByZWR1Y2UgcmVnaXN0ZXIgcHJlc3N1cmUKLSBLZWVwIGZpZWxkcyBicmllZjsgYXZvaWQgbGlzdHMgb2YgYWx0ZXJuYXRpdmVzLCBkaXNjbGFpbWVycywgb3IgZ2VuZXJpYyBhZHZpY2UuCgpPdXRwdXQgZm9ybWF0IChKU09OKToKYGBganNvbgp7CiAgImJvdHRsZW5lY2siOiAiPG1heCAzMCB3b3Jkcz4iLAogICJtb2RpZmljYXRpb24gcGxhbiI6ICI8bWF4IDM1IHdvcmRzPiIKfWBgYApSZXR1cm4gSlNPTiBvbmx5LgoiIiI=)YouareaseniorTritonkerneloptimizationengineer\.ReadthePyTorchreferencecode,thecurrentTritoncandidate,andtheNsightComputemetrics\.Thenidentifyonehighest\-impactspeedbottleneck,proposeoneoptimizationmethodandproposeamodificationplan\.Besurgicalandmetrics\-driven\.\#PyTorchReference‘‘‘python$python\_code‘‘‘\#CurrentTritonKernel‘‘‘python$TRITON\_CODE‘‘‘\#CurrentOptimizationStage‘‘‘$STAGE\_CONTEXT‘‘‘\#NsightComputeMetrics‘‘‘$NCU\_METRICS‘‘‘Rules:\-Return\*\*one\*\*optimizationmethod\-\-thelargestexpectedspeedup\.\-FocusonTriton\-specificoptimizations:\*\*\*BLOCK\_M/N/Ktuning\*\*:Adjusttilesizestooptimizedatareuseandcacheefficiency\*\*\*num\_warps\*\*:Controloccupancy\(2/4/8warpsperblock\)\*\*\*num\_stages\*\*:Enablesoftwarepipelining\(2\-4stagesformemory\-boundkernels\)\*\*\*Memoryaccesspatterns\*\*:Optimizecoalescing,usetl\.trans\(\)forlayoutchanges\*\*\*Gridconfiguration\*\*:Adjustprogram\_idmappingandworkloaddistribution\-PreferchangesthatdirectlyaddressmeasuredbottlenecksfromNCUmetrics:\*HighDRAMthroughput\-\-IncreaseBLOCKsizefordatareuse\*Lowcachehitrate\-\-AdjustBLOCKsizeforbetterlocality\*Lowoccupancy\-\-Tunenum\_warps,reduceregisterpressure\-Keepfieldsbrief;avoidlistsofalternatives,disclaimers,orgenericadvice\.Outputformat\(JSON\):‘‘‘json\{"bottleneck":"<max30words\>","modificationplan":"<max35words\>"\}‘‘‘ReturnJSONonly\."""

### E\.4Code Agent Prompts

The code agent generates or modifies kernel implementations based on feedback from the profile agent or the debug agent\. Different prompts guide optimization or repair tasks\.

#### E\.4\.1Optimization Mode \(Algorithmic Refinement\)

During algorithmic refinement, the code agent receives the current kernel and optimization analysis as input\. The prompt directs the agent to implement structural transformations suggested by the profile agent and apply specific algorithmic changes while preserving correctness\. It outputs the optimized Triton kernel code\.

[⬇](data:text/plain;base64,WW91IGFyZSBvcHRpbWl6aW5nIGEgVHJpdG9uIGtlcm5lbCBiYXNlZCBvbiBhbGdvcml0aG1pYyBhbmFseXNpcy4KCiMgQ3VycmVudCBLZXJuZWwgKG5lZWRzIG9wdGltaXphdGlvbikKYGBgcHl0aG9uCiRjdXJyZW50X2tlcm5lbApgYGAKCiMgQW5hbHlzaXMgUmVzdWx0cwpgYGAKJEFOQUxZU0lTX0lORk8KYGBgCgojIFlvdXIgVGFzawoKSW1wbGVtZW50IHRoZSBvcHRpbWl6YXRpb24gc3RyYXRlZ3kgYWJvdmUuIEZvY3VzIG9uIHRoZSBzcGVjaWZpYyBib3R0bGVuZWNrIGlkZW50aWZpZWQuCgojIyBLZXkgUmVxdWlyZW1lbnRzCgoxLiAqKlByZXNlcnZlIGNvcnJlY3RuZXNzKio6IE1haW50YWluIHRoZSBzYW1lIGlucHV0L291dHB1dCBiZWhhdmlvcgoyLiAqKkFwcGx5IHRoZSBvcHRpbWl6YXRpb24qKjogRm9sbG93IHRoZSBpbXBsZW1lbnRhdGlvbiBwbGFuIGV4YWN0bHkKMy4gKipVc2UgdmFsaWQgVHJpdG9uIHN5bnRheCoqOgogICAtIEV2ZXJ5IGtlcm5lbCBNVVNUIGhhdmUgYEB0cml0b24uaml0YCBkZWNvcmF0b3IKICAgLSBHcmlkIHNpemUgTVVTVCBiZSA+IDA6IHVzZSBgdHJpdG9uLmNkaXYoTiwgQkxPQ0spYCBvciBgbWF4KDEsIE4gLy8gQkxPQ0spYAogICAtIEJMT0NLIHNpemVzIE1VU1QgYmUgcG93ZXItb2YtMjogMTYsIDMyLCA2NCwgMTI4LCAyNTYKICAgLSBObyBgY29udGludWVgLCBgYnJlYWtgLCBgcmV0dXJuYCBpbnNpZGUga2VybmVscyAodXNlIG1hc2tpbmcpCiAgIC0gUHJlZmVyIGB0bC5kb3QoYSwgYiwgYWxsb3dfdGYzMj1UcnVlKWAgZm9yIG1hdG11bCBvcGVyYXRpb25zCjQuICoqT3V0cHV0IGZvcm1hdCoqOgogICAtIEltcG9ydHM6IGBpbXBvcnQgdG9yY2gsIHRvcmNoLm5uIGFzIG5uLCB0cml0b24sIHRyaXRvbi5sYW5ndWFnZSBhcyB0bGAKICAgLSBgQHRyaXRvbi5qaXRgIGtlcm5lbChzKQogICAtIFdyYXBwZXIgZnVuY3Rpb24ocykKICAgLSBgY2xhc3MgTW9kZWxOZXcobm4uTW9kdWxlKWAgLS0gUkVRVUlSRUQKICAgLSBOTyB0ZXN0aW5nIGNvZGUsIE5PIGBpZiBfX25hbWVfXyA9PSAiX19tYWluX18iYAoKRG8gTk9UIGluY2x1ZGU6IHRlc3RpbmcgY29kZSwgaWYgX19uYW1lX18sIGdldF9pbnB1dHMsIGdldF9pbml0X2lucHV0cwoKYGBgcHl0aG9uCiMgPG9wdGltaXplZCBUcml0b24gY29kZT4KYGBgCiIiIg==)YouareoptimizingaTritonkernelbasedonalgorithmicanalysis\.\#CurrentKernel\(needsoptimization\)‘‘‘python$current\_kernel‘‘‘\#AnalysisResults‘‘‘$ANALYSIS\_INFO‘‘‘\#YourTaskImplementtheoptimizationstrategyabove\.Focusonthespecificbottleneckidentified\.\#\#KeyRequirements1\.\*\*Preservecorrectness\*\*:Maintainthesameinput/outputbehavior2\.\*\*Applytheoptimization\*\*:Followtheimplementationplanexactly3\.\*\*UsevalidTritonsyntax\*\*:\-EverykernelMUSThave‘@triton\.jit‘decorator\-GridsizeMUSTbe\>0:use‘triton\.cdiv\(N,BLOCK\)‘or‘max\(1,N//BLOCK\)‘\-BLOCKsizesMUSTbepower\-of\-2:16,32,64,128,256\-No‘continue‘,‘break‘,‘return‘insidekernels\(usemasking\)\-Prefer‘tl\.dot\(a,b,allow\_tf32=True\)‘formatmuloperations4\.\*\*Outputformat\*\*:\-Imports:‘importtorch,torch\.nnasnn,triton,triton\.languageastl‘\-‘@triton\.jit‘kernel\(s\)\-Wrapperfunction\(s\)\-‘classModelNew\(nn\.Module\)‘\-\-REQUIRED\-NOtestingcode,NO‘if\_\_name\_\_=="\_\_main\_\_"‘DoNOTinclude:testingcode,if\_\_name\_\_,get\_inputs,get\_init\_inputs‘‘‘python\#<optimizedTritoncode\>‘‘‘"""

#### E\.4\.2Optimization Mode \(Hardware\-Specific Tuning\)

During hardware\-specific tuning stages, the code agent receives the current kernel, stage\-specific system prompts, and optimization analysis as input\. The prompt directs the agent to apply focused hardware\-level optimizations based on the profile agent’s suggestions\. It outputs the optimized Triton kernel code\.

[⬇](data:text/plain;base64,WW91IGFyZSBhIFRyaXRvbiBrZXJuZWwgb3B0aW1pemF0aW9uIHNwZWNpYWxpc3QuIEdlbmVyYXRlIHRoZSBGQVNURVNUIHBvc3NpYmxlIGtlcm5lbCBiYXNlZCBvbiB0aGUgYW5hbHlzaXMuCgojIFRhcmdldCBHUFU6CiRncHVfbmFtZQoKIyBDdXJyZW50IEtlcm5lbCAobmVlZHMgb3B0aW1pemF0aW9uKQpgYGBweXRob24KJGN1cnJlbnRfa2VybmVsCmBgYAoKIyBDdXJyZW50IE9wdGltaXphdGlvbiBTdGFnZQpgYGAKJFNUQUdFX0NPTlRFWFQKYGBgCgojIEFuYWx5c2lzIFJlc3VsdHMKYGBgCiRBTkFMWVNJU19JTkZPCmBgYAoKIyMgQ1JJVElDQUwgLS0gQ29kZSBNVVNUIGNvbXBpbGUgYW5kIHJ1bjoKMS4gRVZFUlkga2VybmVsIGZ1bmN0aW9uIE1VU1QgaGF2ZSBgQHRyaXRvbi5qaXRgIGRlY29yYXRvcgoyLiBHcmlkIHNpemUgTVVTVCBiZSA+IDA6IHVzZSBgdHJpdG9uLmNkaXYoTiwgQkxPQ0spYCBvciBgbWF4KDEsIE4gLy8gQkxPQ0spYAozLiBCTE9DSyBzaXplcyBNVVNUIGJlIHBvd2VyLW9mLTI6IDE2LCAzMiwgNjQsIDEyOCwgMjU2CjQuIGB0bC5wcm9ncmFtX2lkKGF4aXMpYCBvbmx5IHN1cHBvcnRzIGF4aXMgPSAwLCAxLCAyCjUuIE5vIGBjb250aW51ZWAsIGBicmVha2AsIGByZXR1cm5gIGluc2lkZSBsb29wcyAtLSB1c2UgbWFza2luZwo2LiBObyB0ZW5zb3IgaW5kZXhpbmcgd2l0aCBsb29wIHZhcnM6IGB4WzosIGldYCBpcyBJTlZBTElECjcuIG1hc2sgc2hhcGUgTVVTVCBtYXRjaCBkYXRhIHNoYXBlIGluIHRsLmxvYWQvdGwuc3RvcmUKCiMjIE1pc3NpbmcgVHJpdG9uIEZ1bmN0aW9ucyAoaW1wbGVtZW50IG1hbnVhbGx5KToKLSB0bC50YW5oLCB0bC5zaWdtb2lkLCB0bC5nZWx1LCB0bC5zaWx1LCB0bC5zb2Z0bWF4LCB0bC5taXNoCgojIyBPVVRQVVQgRk9STUFUIChTVFJJQ1QpOgoxLiBJbXBvcnRzOiB0b3JjaCwgdG9yY2gubm4sIHRyaXRvbiwgdHJpdG9uLmxhbmd1YWdlIGFzIHRsCjIuIEB0cml0b24uaml0IGRlY29yYXRlZCBrZXJuZWwgZnVuY3Rpb24ocykKMy4gV3JhcHBlciBmdW5jdGlvbihzKSBmb3IgZ3JpZCBjYWxjdWxhdGlvbiBhbmQga2VybmVsIGxhdW5jaAo0LiBjbGFzcyBNb2RlbE5ldyhubi5Nb2R1bGUpIHRoYXQgY2FsbHMgeW91ciBrZXJuZWxzCgpEbyBOT1QgaW5jbHVkZTogdGVzdGluZyBjb2RlLCBpZiBfX25hbWVfXywgZ2V0X2lucHV0cywgZ2V0X2luaXRfaW5wdXRzCgpgYGBweXRob24KIyA8b3B0aW1pemVkIFRyaXRvbiBjb2RlPgpgYGA=)YouareaTritonkerneloptimizationspecialist\.GeneratetheFASTESTpossiblekernelbasedontheanalysis\.\#TargetGPU:$gpu\_name\#CurrentKernel\(needsoptimization\)‘‘‘python$current\_kernel‘‘‘\#CurrentOptimizationStage‘‘‘$STAGE\_CONTEXT‘‘‘\#AnalysisResults‘‘‘$ANALYSIS\_INFO‘‘‘\#\#CRITICAL\-\-CodeMUSTcompileandrun:1\.EVERYkernelfunctionMUSThave‘@triton\.jit‘decorator2\.GridsizeMUSTbe\>0:use‘triton\.cdiv\(N,BLOCK\)‘or‘max\(1,N//BLOCK\)‘3\.BLOCKsizesMUSTbepower\-of\-2:16,32,64,128,2564\.‘tl\.program\_id\(axis\)‘onlysupportsaxis=0,1,25\.No‘continue‘,‘break‘,‘return‘insideloops\-\-usemasking6\.Notensorindexingwithloopvars:‘x\[:,i\]‘isINVALID7\.maskshapeMUSTmatchdatashapeintl\.load/tl\.store\#\#MissingTritonFunctions\(implementmanually\):\-tl\.tanh,tl\.sigmoid,tl\.gelu,tl\.silu,tl\.softmax,tl\.mish\#\#OUTPUTFORMAT\(STRICT\):1\.Imports:torch,torch\.nn,triton,triton\.languageastl2\.@triton\.jitdecoratedkernelfunction\(s\)3\.Wrapperfunction\(s\)forgridcalculationandkernellaunch4\.classModelNew\(nn\.Module\)thatcallsyourkernelsDoNOTinclude:testingcode,if\_\_name\_\_,get\_inputs,get\_init\_inputs‘‘‘python\#<optimizedTritoncode\>‘‘‘

#### E\.4\.3Repair Mode

When kernel execution fails, the code agent receives kernel code and debug analysis as input\. The prompt emphasizes strict adherence to Triton syntax rules and output format requirements\. It outputs corrected Triton kernel code\.

[⬇](data:text/plain;base64,Rml4IHRoZSBUcml0b24ga2VybmVsIGVycm9ycy4gR2VuZXJhdGUgY29ycmVjdCBjb2RlIGJhc2VkIG9uIHRoZSBlcnJvciBhbmFseXNpcy4KCiMgQnJva2VuIENvZGUKYGBgcHl0aG9uCiRPTERfQ09ERQpgYGAKCiMgQW5hbHlzaXMgUmVzdWx0cwpgYGAKJEFOQUxZU0lTX0lORk8KYGBgCgojIyBPVVRQVVQgRk9STUFUIChTVFJJQ1QpOgoxLiBJbXBvcnRzOiB0b3JjaCwgdG9yY2gubm4sIHRyaXRvbiwgdHJpdG9uLmxhbmd1YWdlIGFzIHRsIChhbmQgbWF0aCBpZiBuZWVkZWQpCjIuIEB0cml0b24uaml0IGRlY29yYXRlZCBrZXJuZWwgZnVuY3Rpb24ocykKMy4gV3JhcHBlciBmdW5jdGlvbihzKSBmb3IgZ3JpZCBjYWxjdWxhdGlvbiBhbmQga2VybmVsIGxhdW5jaAo0LiBjbGFzcyBNb2RlbE5ldyhubi5Nb2R1bGUpIC0tIFJFUVVJUkVECgpEbyBOT1QgaW5jbHVkZTogdGVzdGluZyBjb2RlLCBpZiBfX25hbWVfXywgZ2V0X2lucHV0cywgZ2V0X2luaXRfaW5wdXRzCgpgYGBweXRob24KIyA8Y29ycmVjdGVkIGNvZGU+CmBgYA==)FixtheTritonkernelerrors\.Generatecorrectcodebasedontheerroranalysis\.\#BrokenCode‘‘‘python$OLD\_CODE‘‘‘\#AnalysisResults‘‘‘$ANALYSIS\_INFO‘‘‘\#\#OUTPUTFORMAT\(STRICT\):1\.Imports:torch,torch\.nn,triton,triton\.languageastl\(andmathifneeded\)2\.@triton\.jitdecoratedkernelfunction\(s\)3\.Wrapperfunction\(s\)forgridcalculationandkernellaunch4\.classModelNew\(nn\.Module\)\-\-REQUIREDDoNOTinclude:testingcode,if\_\_name\_\_,get\_inputs,get\_init\_inputs‘‘‘python\#<correctedcode\>‘‘‘

### E\.5Debug Agent Prompt

The debug agent receives error logs, PyTorch reference implementations, and broken kernel code as input\. The prompt guides the agent to diagnose execution failures and identify root causes with specific, actionable diagnostic information\. It outputs a structured JSON response containing the critical issue and required modifications, which the code agent then uses to generate fixes\.

[⬇](data:text/plain;base64,WW91IGFyZSBhIFRyaXRvbiBrZXJuZWwgZGVidWdnaW5nIGV4cGVydC4gQW5hbHl6ZSB0aGUgZXJyb3IgYW5kIGlkZW50aWZ5IHRoZSByb290IGNhdXNlLgoKIyBFcnJvciBMb2cKYGBgCiRFUlJPUl9MT0cKYGBgCgojIEV4cGVjdGVkIEJlaGF2aW9yIChQeVRvcmNoIFJlZmVyZW5jZSkKYGBgcHl0aG9uCiRQWVRPUkNIX0NPREUKYGBgCgojIEN1cnJlbnQgSW1wbGVtZW50YXRpb24gKEJyb2tlbiBUcml0b24gS2VybmVsKQpgYGBweXRob24KJEtFUk5FTF9DT0RFCmBgYAoKIyBZb3VyIFRhc2sKCklkZW50aWZ5IHRoZSAqKm1vc3QgY3JpdGljYWwgaXNzdWUqKiB0aGF0IGNhdXNlcyB0aGUgZXJyb3IgYWJvdmUuCgojIyBBbmFseXNpcyBHdWlkZWxpbmVzCgoxLiAqKkZvY3VzIG9uIHJvb3QgY2F1c2UqKiwgbm90IHN5bXB0b21zCiAgIC0gQmFkOiAiT3V0cHV0IGlzIHdyb25nIgogICAtIEdvb2Q6ICJCTE9DS19LIGxvb3AgbWlzc2luZywgb25seSBwcm9jZXNzZXMgZmlyc3QgMzIgZWxlbWVudHMgb2YgSyBkaW1lbnNpb24iCgoyLiAqKkJlIHNwZWNpZmljIGFib3V0IFdIQVQgYW5kIFdIRVJFKioKICAgLSBCYWQ6ICJNZW1vcnkgYWNjZXNzIGlzc3VlIgogICAtIEdvb2Q6ICJMaW5lIDQ1OiB0bC5hdG9taWNfYWRkKGNfYmxvY2tfcHRyLCBhY2MpIC0gYXRvbWljX2FkZCByZXF1aXJlcyBzY2FsYXIgcG9pbnRlciwgbm90IGJsb2NrX3B0ciIKCjMuICoqUHJpb3JpdGl6ZSBieSBpbXBhY3QqKgogICAtIENvcnJlY3RuZXNzIGJ1Z3MgPiBQZXJmb3JtYW5jZSBpc3N1ZXMgPiBTdHlsZSBwcm9ibGVtcwogICAtIEFsZ29yaXRobSBlcnJvcnMgPiBJbXBsZW1lbnRhdGlvbiBkZXRhaWxzCgojIyBPdXRwdXQgRm9ybWF0CmBgYGpzb24KewogICJjcml0aWNhbF9pc3N1ZSI6ICI8Q29uY2lzZSBkZXNjcmlwdGlvbiBvZiBUSEUgcm9vdCBjYXVzZSwgbWF4IDMwIHdvcmRzPiIsCiAgIm1vZGlmaWNhdGlvbiBwbGFuIjogIjxXaGF0IG5lZWRzIHRvIGNoYW5nZSAobm90IGhvdyksIG1heCAzMCB3b3Jkcz4iCn0KYGBgClJldHVybiBKU09OIG9ubHku)YouareaTritonkerneldebuggingexpert\.Analyzetheerrorandidentifytherootcause\.\#ErrorLog‘‘‘$ERROR\_LOG‘‘‘\#ExpectedBehavior\(PyTorchReference\)‘‘‘python$PYTORCH\_CODE‘‘‘\#CurrentImplementation\(BrokenTritonKernel\)‘‘‘python$KERNEL\_CODE‘‘‘\#YourTaskIdentifythe\*\*mostcriticalissue\*\*thatcausestheerrorabove\.\#\#AnalysisGuidelines1\.\*\*Focusonrootcause\*\*,notsymptoms\-Bad:"Outputiswrong"\-Good:"BLOCK\_Kloopmissing,onlyprocessesfirst32elementsofKdimension"2\.\*\*BespecificaboutWHATandWHERE\*\*\-Bad:"Memoryaccessissue"\-Good:"Line45:tl\.atomic\_add\(c\_block\_ptr,acc\)\-atomic\_addrequiresscalarpointer,notblock\_ptr"3\.\*\*Prioritizebyimpact\*\*\-Correctnessbugs\>Performanceissues\>Styleproblems\-Algorithmerrors\>Implementationdetails\#\#OutputFormat‘‘‘json\{"critical\_issue":"<ConcisedescriptionofTHErootcause,max30words\>","modificationplan":"<Whatneedstochange\(nothow\),max30words\>"\}‘‘‘ReturnJSONonly\.

## Appendix FNsight Compute Profiling Metrics

NVIDIA Nsight Compute \(NCU\) provides detailed hardware\-level performance metrics for analyzing GPU kernel execution\. In our framework, NCU profiling is invoked at each optimization stage to provide low\-level performance signals that guide the profile agent’s analysis and modification planning\.

We collect a unified set of core metrics covering compute utilization, memory hierarchy behavior, and execution stalls\. The following list shows the metrics reported to the profile agent during each profiling step:

[⬇](data:text/plain;base64,InNtX190aHJvdWdocHV0LmF2Zy5wY3Rfb2ZfcGVha19zdXN0YWluZWRfZWxhcHNlZCIsICAgIyBTTSBjb21wdXRlIHV0aWxpemF0aW9uCiJsYXVuY2hfX2dyaWRfc2l6ZSIsICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICMgR2xvYmFsIGdyaWQgc2l6ZQoic21fX3dhcnBzX2FjdGl2ZS5hdmcucGN0X29mX3BlYWtfc3VzdGFpbmVkX2FjdGl2ZSIsICAjIFdhcnAgb2NjdXBhbmN5CiJkcmFtX190aHJvdWdocHV0LmF2Zy5wY3Rfb2ZfcGVha19zdXN0YWluZWRfZWxhcHNlZCIsICMgRFJBTSBiYW5kd2lkdGggdXRpbGl6YXRpb24KImx0c19fdF9zZWN0b3JfaGl0X3JhdGUucGN0IiwgICAgICAgICAgICAgICAgICAgICAgICAgIyBMMiBjYWNoZSBoaXQgcmF0ZQoic21zcF9fd2FycF9pc3N1ZV9zdGFsbGVkX21lbW9yeV9kZXBlbmRlbmN5X3Blcl93YXJwX2FjdGl2ZS5wY3QiLCAgIyBNZW1vcnkgZGVwZW5kZW5jeSBzdGFsbHM=)"sm\_\_throughput\.avg\.pct\_of\_peak\_sustained\_elapsed",\#SMcomputeutilization"launch\_\_grid\_size",\#Globalgridsize"sm\_\_warps\_active\.avg\.pct\_of\_peak\_sustained\_active",\#Warpoccupancy"dram\_\_throughput\.avg\.pct\_of\_peak\_sustained\_elapsed",\#DRAMbandwidthutilization"lts\_\_t\_sector\_hit\_rate\.pct",\#L2cachehitrate"smsp\_\_warp\_issue\_stalled\_memory\_dependency\_per\_warp\_active\.pct",\#Memorydependencystalls

Although the same metric set is collected throughout all stages, different optimization stages emphasize different performance aspects according to their objectives:

- •Algorithmic refinement:SM utilization, global parallelism \(grid size\), and DRAM bandwidth utilization\.
- •Parallel mapping:warp occupancy, global parallelism \(grid size\), and SM utilization\.
- •Tensor tiling:DRAM throughput, L2 cache hit rate, and memory\-dependency stalls\.
- •Memory optimization:memory\-dependency stalls, DRAM throughput, and L2 cache hit rate\.

By providing stage\-specific metrics to the profile agent, we enable focused bottleneck diagnoses and targeted modification plans\.

## Appendix GLimitations

Dependence on Expert Optimization Priors\.EGG relies on a set of expert\-designed optimization principles to guide the staged search process\. While these priors substantially improve search efficiency, they may also bias exploration toward known optimization patterns and limit the discovery of unconventional but potentially superior designs\. Enabling the framework to automatically learn or adapt expert priors remains an important direction for future work\.

Lack of Joint Cross\-Stage Optimization\.EGG decomposes optimization into sequential stages to restrict the search space, improving controllability and stability in the open\-ended LLM generation setting\. However, this design limits cross\-stage interactions: locally optimal decisions at individual stages may not always lead to the best final performance when combined\. For example, the best parallelization choice under one tiling strategy may be worse than another parallelization choice paired with a different tiling configuration\. We make this tradeoff because jointly modifying multiple coupled factors enlarges the effective proposal space and increases search cost\. Compared with the search\-based baseline CudaForge, EGG achieves better performance with a smaller budget \(50k vs\. 110k output tokens per kernel\)\. A promising future direction is to propagate top\-kkcandidates after each stage when a larger token budget is available, enabling limited joint exploration while reducing the risk of discarding candidates that may improve in later stages\.