STARK: Strategic Team of Agents for Refining Kernels
Authors: Juncheng Dong
Yang Yang
Tao Liu
Yang Wang
Feng Qi
Vahid Tarokh
Kaushik Rangadurai
Shuang Yang
Venue: ICLR 2026 Poster (2026)
Keywords: large language model
agent
kernel optimization
efficiency
Primary area: infrastructure, software libraries, hardware, systems, etc.
TL;DR: We propose an agent framework for optimizing GPU kernels.

Abstract: The efficiency of GPU kernels is central to the progress of modern AI, yet optimizing them remains a difficult and labor-intensive task due to complex interactions between memory hierarchies, thread scheduling, and hardware-specific characteristics. While recent advances in large language models (LLMs) provide new opportunities for automated code generation, existing approaches largely treat LLMs as single-shot generators or naive refinement tools, limiting their effectiveness in navigating the irregular kernel optimization landscape. We introduce an LLM agentic framework for GPU kernel optimization that systematically explores the design space through multi-agent collaboration, grounded instruction, dynamic context management, and strategic search. This framework mimics the workflow of expert engineers, enabling LLMs to reason about hardware trade-offs, incorporate profiling feedback, and refine kernels iteratively. We evaluate our approach on KernelBench, a benchmark for LLM-based kernel optimization, and demonstrate substantial improvements over baseline agents: our system produces correct solutions where baselines often fail, and achieves kernels with up to 16$\times$ faster runtime performance. These results highlight the potential of agentic LLM frameworks to advance fully automated, scalable GPU kernel optimization.

Published as a conference paper at ICLR 2026
STARK: STRATEGIC TEAM OF AGENTS FOR REFINING KERNELS
Juncheng Dong 1 Yang Yang', Tao Liu Yang Wang', Feng Qi!,
Vahid Tarokh Kaushik Rangadurail & Shuang Yang'
Meta Ranking Al Research
Duke University
*Work done during an internship at Meta
juncheng dong@duke edu, yzyangemeta...m.m com, shuangyang@mete com
ABSTRACT
The efficiency of GPU kernels is central to the progress of modern AI, yet optimiz- ing them remains a difficult and labor-intensive task due to complex interactions between memory hierarchies, thread scheduling, and hardware-specific character- istics. While recent advances in large language models (LLMs) provide new op- portunities for automated code generation, existing approaches largely treat L.LMs as single-shot generators or naive refinement tools, limiting their effectiveness navigating the irregular kernel optimization landscape. We introduce an LLM agentic framework for GPU kernel optimization that systematically explores the design space through multi-agent collaboration, grounded instruction, dynamic context management, and strategic search. This framework mimics the workflow of expert engineers, enabling LLMs to reason about hardware trade-offs, incorpo- rate profiling feedback, and refine kernels iteratively. We evaluate our approach on KernelBench, a benchmark for L.LM-based kernel optimization, and demon- strate substantial improvements over baseline agents: our system produces correct solutions where baselines often fail, and achieves kernels with up to 16x faster runtime performance. These results highlight the potential of agentic LLM frame- works to advance fully automated, scalable GPU kernel optimization.
INTRODUCTION
Artificial intelligence (AI) has advanced at an unprecedented pace, transforming both re- search and real-world applications. While in- novations in model architectures and training algorithms have been central to this progress, the efficiency of the computational infrastruc- ture that executes them is equally critical. At the of modern AI systems GPU ker- nels, which implement fundamental operations such matrix multiplication and convolution. Even modest improvements in GPU kernel effi- ciency can translate into significant reductions in training time, inference latency, and deploy- ment cost, making kernel optimization corner- stone for sustaining Al's rapid growth.
Sampling Agant Refiexion Agent
17.5
16%
15.0
13.7%
12.5
10.7x
10.0
5.0
2.5
2.35x
Figure 1: Speedup of STARK over bascline agents on KernelBench (L1-L3) with same number of at- tempts. Bars report GPU wall-clock speedups (x) rel- ative to the Sampling and Reflexion agents; higher better. STARK reaches up to 16x over Reflexion (L2) and 10.7x over Sampling (LI).
Despite their importance, designing and opti- mizing GPU kernels remains major challenge. The performance kernel depends on sub- de interactions between thread scheduling, memory hierarchy utilization, synchronization. and hardware-specite characteristics. Small changes in tiling strategies, loop unrolling, or memory alignment yield disproportionate effects runtime. As result, the kernel optimization land- scape is highly irregular, sccitecture-dependeer and difficult to navigate. Existing approaches largely fall into categories: manual optimization by expert engineers, which is effective but

---

Published as a conference paper at ICLR 2026
labor-intensive and difficult to scale; and automated compilers and domain-specift languages (DSLs) such as TVM and Triton (Chen et al. 2018 Tillet et al. 2019), which apply heuristics or search but often struggle with irregular operators and hardware variability (Zheng al. 2020a b).
The rapid progress of large language models (LLMs) opens a new opportunity for kernel optimiza- tion. Beyond their ability to generate correct code, LLMs can be guided to reason about hardware trade-offs, adapt to profiling feedback, and iteratively refine implementations. However, prior work has mostly treated LLMs single-shot code generators simple refinement tools (Ouyang al.) 2025), which underutilizes their potential for structured exploration of the kernel design space. To build more powerful agent, we identify and address three critical limitations in existing methods:
1.
Naive exploration strategy. Current agents typically refine code linearly, learning only from the immediately preceding attempt. This simplistic process neglects the rich history of prior attempts and fails effectively balance the exploration-exploitatio trade-off, often getting trapped in local optima.
2. Monolithic agent design. Kernel optimization is a multifaceted task requiring distinct capabilities for planning, implementation, and reflection. By assigning all these responsi- bilities to single, generalist LLM, current agents operate inefficiently.
3. Planning-implemmnnttioon gap. We observe a failure mode particularly acute in this do- main: LLMs frequently devise a correct high-level optimization plan (e.g., "apply memory tiling") but fail to translate it into valid low-level CUDA code, This gap stems from the relative scarcity of expert-level kernel code the models' training data.
To address these limitations, we introduce STARK (Strategic Team of Agents for Refining Kernels), a novel framework for automated GPU-kernel optimization. Our contributions are threefold:
Collaborative multi-agent workflow. We design a workflow with specialized agents for planning. coding, and reflection, mirroring expert development cycle and overcoming the inefficiencies of monolithic designs.
Bridging the instruction and dynamic context windows-the translate high-level strategies into precise, actionable code edits, ensuring robust coordination across agents.
gap. We propose two mechanisms grounded
Strategic search for refinement. We incorporate a search policy that balances exploration and exploitation over prior attempts, enabling systematic discovery of strong kemells.
We evaluate our framework on KernelBench Ouyang et al. (2025), benchmark designed to as- sess LLM-based GPU kernel optimization. Experiments show that combining these improvements leads to agent system significantly more competitive than the baseline agents in both runtime performance and success rate across diverse kernel problems, authoring competitive kernels for the challenging problems in KerneBBench where the baseline agents struggle to even find working solution. Notably, STARK achieves more than 10x speedup over kernels produced by the base- line agents (i.e., the optimized kernels run in under one-tenth the time of the baseline.). Overall, our work suggests that LLM--riven agents represent a promising step toward fully automated GPU kernel optimization.
RELATED WORK
Due to space constrain, we only review the most relevant prior work here and defer the complete discussion to Appendix F
The optimization of GPU kernels has progressed from empirical auto-tuning frameworks that perform black-box parameter searches (van Werkhoven 2019; Nugteren & Codreanu) 2015) and compiler-based approaches with static heuristics (Yang et | 2010). to the use of machine learning (ML). ML-based techniques have been used to replace hand-tuned heuristics in production compil- ers (Trofin et al. 2021), learn cost models to guide optimization (Chen et al., 2018). and even learn directly from raw source code without manual feature engineering (Cummins al. 2017). A sig- nificant leap was the use of deep reinforcement learning to discover fundamentally new algorithms, demonstrated by AlphaTTensor's success in finding faster matrix multiplication methods (Fawzi
2

---

Published as a conference paper at ICLR 2026
LLM Agent GPU Kernel Optimization
Tree Memory & Strategie Scarch
Arch Source
Strotegic
Optimiznd GPu Kamnle
Memory Node
(1)
Search
LLM
opsindsed Model I...ff
alese Bodeal
(2)
Kernel
Kereel IP3
Agents
Compller Info Gornectess info Runtime
Plan
Comppier (hytte
Camplic bric ..)
Correctress Info
Debug
Runtime 20ts
Context Window
class (odei(en.modulle
Plan Agent
<< IICRREEVVE BEGINS>>> fuse the addition relu single operator relu(x) <<CIMPEVVE ENDS>>>
Code Agent
(c) Grounded Instruction Dynamic Context Window
Context Window
kernel optimizers rely on a monolithic agent with local iterative refinement. (b) STARK replaces this with a collaborative multi-agent workflow (plan/code/debug) coupled with strategic search over tree memory. (c) The plan agent issues grounded instructions that anchor edits to code spans; dynamic context windows surface role-speciiic history; and the debug agent repairs failures. See Section for details.
Figure Overview of STARK. (a) Prior
et al. 2022). While powerful, these prior works either optimize within fixed search space or oper- in purely formal domains. Our work addresses these limitations by operating directly source code to implement novel, structural changes.
The emergence of powerful Large Language Models (LLMs) has revolutionized programmatic in- teraction with source code, demonstrating a remarkable proficiency in generating code for diverse applications from competitive programming to compiler testing (Gul 2023 Zhong Wang 2024 Jain et al.) 2025). This capability has catalyzed a paradigm shift away from single-shor code genera- tion and toward the development of autonomous LLM agents. An agent enhances a base LLM with planning, memory, and tool-use capabilities to direct its own workflow (Weng) 2023). The success of frameworks like SWE-agent independently resolving complex GitHub issues has validated the power of this approach for software engineering (SWE) (Yang et al.) (222)). While the application LLM agents to SWE a burgeoning field of research (Yang et al. 2024 Antoniades et al. 2024 Yang al. 2025). their potential in the specialized domain of GPU kernel optimization remains largely unexplored. To fill this gap, we designed STARK, an agent framework with capabilities tailored the unique challenges of this domain.
PRELIMINARY
3.1 LLMS AND AUTOREGRESSIVE GENERATION
Given an input sequence x = (I1, T2, . )n) (e.g., ., the task instruction) as the context, an L.M with parameters generates an output sequence (y1. y2.... vmm) where yt V,t {l...., m) are tokens. Pretrained on a massive corpus of text, LLMs autoregressively generate the next token Yt conditioning on and all the previously generated token VCt (V1..... 1)-1)- Specifically, at each time t, the LLM first computes the logits 2g (vlp<e, x) for each token in the vocabulary y and generate y, following the conditional distribution
exp(zo(VilVer,z))
po(yaluee,
Ever Exprxep(2z01 wxpp=x(pppxx/eeeeoooooo Vet, x)77)
3

---

Published as a conference paper at ICLR 2026
The temperature parameter > modulates the randommess of an LLM's output. Higher values of т flatten the next token distribution in Equation|] [ encouraging creative and diverse responses. Conversely, lower values sharpen the distribution, promoting deterministic and high-fidelity outputs.
This trade-off is critical in complex tasks, as different sub-problems demand different behaviors. For instance, planning and exploration benefit from high temperature generate novel strategies, whereas tasks requiring precision and factual correctness, such as code implementation, necessitate low temperature ensure reliability. A single agent with fixed temperature is ill-equipped to handle this dichotomy. This observation is a core motivation for STARK's multi-agent design, which allows specialized agents to operate at distinct temperatures tailored their roles, i.e., high 7 for the creative plan agent and low r for the precise code agent.
3.2 KERNELBENCH
KernelBench Ouyang al. (2025) is recently proposed benchmark specifically designed for sessing LLMM-based GPU kernel optimization. Unlike prior evaluations that focus only on code cor- rectness small-scale operator tests, KernelBench provides principled and reproducible testbed that measures both correctness and runtime efficiency across a broad spectrum of GPU workloads. KernelBench comprises suite of optimization tasks, categorized into three difficulty levels. For each task, the objective is to create a custom GPU kernel that is functionally equivalent to a pro- vided PyTorch reference implementation while minimizing its wall-clock execution time. See an example of the KernelBench task in Appendix
Specifically, Level 1 tasks focus on single, common operators such matrix multiplication and convolution, serving a baseline for fundamental optimization capabilities; Level 2 tasks com- prise tasks with multiple operators fused into a single kerne,, testing the ability to manage more complex dataflows and scheduling; Level 3 tasks represent the highest difficulty, featuring popular full ML architectures such as the ResNet (He et al., 2016) and LSTM (Hochreiter & Schmidhuber, 1997). which involve highly irregular computations and intricate memory access patterns that challenging for both human experts and automated systems to optimize effectively.
4 STARK: STRATEGIC TEAM OF AGENTS FOR REFINING KERNELS
Framework Overview. We present STARK, agentic framework for GPU-kernel optimiza- tion. STARK organizes kernel refinement into three layers: (i) a multi-agent workflow that separates planning, coding, and debugging, (ii) coordination mechanisms with grounded instruction anchor planned edits to concrete code spans and dynamic context windows that surface role- cel-sppcifii history (e.g.. prior attempts, fallures, profiler feedback) each agent, and (iii) strategic search policy that balances exploration and exploitation across iterative attempts. Notably, multi-agent workflow and grounded instruction improve reliability even under single-attempt budget, whereas dynamic con- text windows and strategic search deliver most of their gains when multiple attempts are allowed. Figure provides an overview; the following subsections detail each component in turn.
4.1 MULTI-AGENT COLLABORATION
Optimizing GPU kernels is inherently multifaceted and mirrors expert team workflows. A single agent typically fails to balance correctness, performance, and exploration across vast, irregular design space. In particular, strategy discovery (e.g., fusion, vectorization, shared-memon tiling) benefits from higher-tenen temperature generation that encourages diversity whereas strategy realization, i.e., committing those ideas to code, requires low-temperature precision to avoid errors. We therefore adopt a multi-agent framework that enables role specialization through L.MMs.
Multi-Agent Design (MAD). Specifically, STARK decomposes kernel optimization into three roles plan, code, and debug. Using a role-specific context window (Section 4.4) with selected prior at- tempts and execution outcomes, the plan agent proposes targeted transformations to either the source kernel or a candidate chosen by the strategic search policy (Section 4.2), emitting grounded instruc- tions (Section 4.3) that anchor edits to explicit code spans. The code agent consumes grounded instructions and translates them into executable GPU-kernel code, conditioning on its own context window to improve adherence and code quality. The debug agent repairs promising but failing
4

---

Published as a conference paper at ICLR 2026
didates by consulting the plan agent's instructions and compiler/uuntime diagnostics, producing working kernel that realizes the intended transformation.
Benefits of MAD. Role specialization lets each agent use prompts and base LLMs matched to its ob- jective. In our instantiation, we choose Claude Sonnet with temperature r=0.8 for the plan agent encourage strategy diversity, and the same model with 7=0.1 1 for the code and debug agents to enforce precision. Despite this simple setup, MAD already performs strongly (see Section (5) We underscore that because the design modular, can swap in planners with richer kernel- optimization priors or code-specialist reasoning models to further improve results. In addition, modularity also exposes bottlenecks. We observe that the dominant bottleneck is code-synthesis fidelity: LLMs often need multiple attempts to faithfully implement given instruction. Finally, MAD makes targeted post-training straightforward: we can fine-tune the base L.LM for a specific agent (e.g., the code agent) without affecting the others, improving stability and predictability. How- ever, a systematic study of agent-specific post-training is orthogonal to our core contributions and is left to future work.
4.2 STRATEGIC SEARCH WITH TREE MEMORY
Prior LLM-driven kernel optimizers typically use either best-of-K sampling that generates multiple candidates independently and select the fastest correct one or iterative refinement which repeatedly edits the latest kernel (Ouyang et al. 2025). However, best-of-K is unguided and wasteful: all the new attempts ignore feedback from earlier attempts and repeatedly probe redundant regions of the design space. On the other hand, iterative refinement is feedback-aware but myopic: by building only on the most recent candidate, it is prone to getting trapped in narrow, suboptimal basins.
To address these limitations, STARK reframes kernel optimization as strategic search over persis- tent tree memory. We maintain a search tree T whose nodes store candidates and their observations (runtime, correctness, and compiler diagnosties). The root represents the source architecture; each edge corresponds to applying grounded instruction from the plan agent and realizing it via the code agent (or repairing via the debug agent). Each node n is assigned a score s(n) reflecting competi- tiveness; in implementation is better. For kernels that are incorrect or failing to compile, we give them scores of +oo. At each step, we (1) select node expand using strategic policy, (2) expand by invoking the plan/code (or debug) agents to produce a child candidate, (3) evaluate for correctness and runtime, and (4) record results in I' to inform subsequent selections. This converts ad-hoc trial-and-error into directed, feedback-driven process.
the straightforward kernel runtime s(n) and treat lower
Policy choice and an adapted e-greedy rule. We compared representative search policies includ- ing Monte-Cari Tree-Search (MCTS), evolutionary, greedy, and e-greedy policies and found that e-greedy consistently performs best under the same budget constraint. Importantly, we observe that kernel optimization poses Domain-speciiff challenges that are root dominance (it is very challeng- ing to even outperform the source architecture in the root node) and frequent compilation/runtime failures. To address these challenges, we adapt the canonical rule as follows: (1) Root throttling: cap the number direct children the root at Thoo to avoid redundant first-hop edits; once the cap is reached, the root is ineligible for selection; (2) branch pruning: if node has more than Metild children and all current children fail, mark the node ineligible to prevent wasting trials; (3) High exploration rate: use relatively large (empirically 0.3-0.4) to counteract local traps; (4) Leaf-biased exploration: with probability 6, sample uniformly from expandable leaves (not only failing nodes), encouraging discovery beyond the immediate failure set.
4.3 GROUNDED INSTRUCTION
We introduce grounded instruction for kernel enhancement. The plan agent must not only pro- pose an optimization, but also insert explicit span anchors in the kernel source that mark exactly where the change should occur. Each anchor is short, machine-checcable tag (i.e, <<<IMPROVE BEGINS>>> ... <<<IMPROVE ENDS>>>) wrapped around the target site, such as load/store, loop body, or the launch configuration. The code agent consumes this annotated scaf- fold and resolves each anchor by emitting concrete CUDA that realizes the instruction. Grounded in- struction tightens plan-code aligmment, curbs hallucinated guidance, and narrows the coder's s search space. It also improves traceability: every proposal leaves visible, verifiable footprint in the final
5

---

Published as a conference paper at ICLR 2026
code. In practice, we observe fewer misinterpretaations and markedly fewer faulty kernels. Despite its simplicity, the mechanism especially effective on Level KernelBench tasks with deeper ar- chitectures (e.g.. VGG).
4.4 DYNAMIC CONTEXT WINDOW
Past attempts provide rich, actionable signals for subsequent decisions, but different agents ben- efit from different views of this history. We therefore maintain dynamic, agent-specific context window that is rebuilt at each selection step for different agents. See Figure for visual demon- stration. Throughout this section,
node i be the node selected by the search policy defined in Section 4.2 We use W(i) to denote the context window containing subset of historical attempts and their evaluation outcomes (c.g., compiler information and runtime). As we always include the source architecture as part of the prompt for agents, W(i) always includes the root node rooo For naive search algorithm without dynamic context window, W(i) = {i, Rroot] only includes node in addition to the root.
Plan
Code
Debug Agent
Agent
Agent
Figure 3: Dynamic Context Window. Nodes with v represent selected nodes.
Tree relations. We use tree relations to build agent-specifer context windows. Let p(i) be the parent of node i. Define the siblings of S(i) = {j p(j) p(i) }. Moreover, define the set of child nodes of a node i as D()). We also maintain small global leaderboard C of top-performing nodes.
Plan agent (local & contrastive global context). For a selected node i, the plan agent conditions on context window Wptam()) (i) that aggregates node i's children and small set of global leaders from the leaderboard Formally,
Wplan (i) = {ú, (roott U D(1) U Top-r(C),
where D(1) contains evaluated children of i with their observations, and Top-r(C) returns the highest-scoring distinct kernels from the global leaderboard (excluding i's subtree) to discourage duplication.
This design serves three purposes. (i) reflection: the plan agent revise stack its prior instruc- tions rather than rediscovering them; (ii) ambition calibration: top competitors prevent redundant exploration and provide transferable motifs such as warp-shuffle reductions, vectorized LD/ST, and therrd- tiling; (ii) capability estimation: by inspecting how past instructions were real- ized or failed by the code agent, the next instruction is adapted to what the code agent reliably execute, improving first-pass success and avoiding instructions beyond current ability. To achieve this, we explicitly require the plan agent to adapt its instruction the code agent's 's demonstrated capabilities observed in D(i).
Code agent (extended context). For kernel code emission at node i, the code agent conditions on
Weode (i) = {i. Proot U D(i) U ( p()) E S(i)).
The nodes in (j p()) E S(i) are essentially the children node siblings. Our insight that these nodes typically share near-identical scaffolds with node from a common planning lineage, so successful patches and micro-optimizations transfer with high probability; conversely, seeing failures in closely related contexts helps the coder avoid repeating the same mistakes. Hence, this extended window serves two aims: (i) reduce implementation errors by letting the coder imitate successful patches from closely related scaffolds and avoid previously observed failure modes; (ii) surface stronger implementations by transferring micro-optimizations (c.g., warp-shuffle motifs, vectorized LD/ST, shared-memory tiling) that have already worked cousin nodes.

---

Published as a conference paper at ICLR 2026
Debug agent (local context), For fault repair, construct the context window for the debug code as
Waehug (z) {2, Nlyoe)) U S(i),
We choose this design mainly for two reasons. Most fixes are structural and local, e.g., off-by- guards, stride/indexing alignment, launch-parameter tweaks, mared-memory sizing often transfer directly among siblings that share the same scaffold. Moreover, restricting the window to S avoids distracting the debug agent with globally unrelated kernels, improving precision and reducing hallucinated edits.
4.5 FRAMEWORK OVERVIEW
Here we provide overview of framework STARK and describe its execution process. Algo- rithm presents its pseudocode.
At high level, STARK repeatedly (i) selects a promising node (a prior attempt) from search tree, (ii) builds agent-specific context windows from local history and global leaders, (iii) asks the planning agent to propose concrete optimization along with grounded instruction anchors inserted into scaffold, (iv) asks the code agent to realize those anchors into an executable kernel, (v) if the selected node has a problematic kernel, we build the debugger's dynamic context window and request minimal fix. The new attempt is evaluated, appended as child node, and the leaderboard is updated. We repeat this process until we reach a pre-specified max attempts B.
Algorithm STARK: Strategic Team of Agents for Refining Kernels
Require: Budget B (max attempts), selection policy Tselect (adapted e-greedy). leaderboard size r Initialize search tree T with root nox (PyTorch reference)
2: Initialize leaderboard C <- {
3: for 1,,,..... do
4: 6 TTeteea (T,C)
pick node to refine > compile fail or unit-test fail recorded at i
5: if HASBUG( (t) then
6: Waehug (t) BUILDCONTEXIDERUCC...
7: kernel' DEBUGAGENT((Wane i.kernel, i.logs)
8:
(ok, correct, runtime, logs) + EVALUATE((ernee')
>> compile, correctness check,
timing
9:
(plan, anchors) (i.plans, i.anchors)
10:
else
Wptan (i) f- BUI.DDCONTEXTPPLANWY T,C)
12:
(plan, anchors) + PLANAGENT(
13:
Wende (i) f BULLDCOTEEETCOOEII
kernel' <- CODEAGENT plan, anchors)
(ok, correct, runtime, logs) EVALUATE((ermnel)
16: end if
ADDCHILD(T, i, kernel', plan, anchors, ok, correct, runtime, logs) C+ UPDATELEADERSIC j, r)
18:
19:
end for
20:
return BEST(C)
fastest correct, grounded kernel
EXPERIMENTS
We use KernelBench (Ouyang et al.) (025), a recently proposed benchmark consisting of compre- hensive and challenging GPU kernel tasks, to validate the effectiveness our proposed approaches.
Baselines and Metrics. We compare our framework STARK with the following list of approaches:
Torch Eager: the out-of-box PyTorch modules without any compilation or optimization.
Torch Compile: We torch. compile to produce optimized versions of the given PyTorch modules. While torch. compile offers different compilation modes, we com- pare two of the most representative and competitive default and max-autotune.

---

Published as a conference paper at ICLR 2026
Sampling Agent: the single agent framework originally proposed and used by Kernel- Bench to evaluate the difficulty of the tasks in KernelBench and the ability of LLMs to write efficient kernels. This agent repeatedly samples responses when given the source model to optimize and chooses the best generated custom kerne as the solution.
Reflexion Agent: this agent follows the Reflexion paradigm (Shinn et al., 2023). where at each optimization step, it tries to update its last attempt using its corresponding observa- tions such as the compiler and runtime information.
We report the following metrics to comprehensively understand the agents' performances: (i) Fast rate is the percentage of the problems for which the agent can generate kernels that are least as fast as the torch baselines; (ii) Success rate represents the percentage of the problems for which the agent can generate compiled and correct kernels; (iii) Speed: To better understand how good the generated kernels are, we also report the average speed across all tasks.
Comparison with Torch Baselines. In Table 1 we present the results about success rate, Fast rate and speed over all 3 levels of KernelBench challenges. For each task, we let all agents to have maximum of B 30 attempts. Due to limited computation resource, representative subset of KernelBench (Ouyang et al. 2025). We use Claude Sonnet 4 as the base LLMs for all the LLMM-based baselines and agents. Due to space constraint, defer implementation and evaluation details to Appendix [B
evaluate the
The results in Table I demonstrate that our proposed framework, STARK, consistently outperforms both the Sampling and Reflexion baselines across all KernelBench difficulty levels. At Level 1, STARK not only achieves perfect 100% success rate but also delivers up 3.0x speedup over Torch Eager baselines, while Sampling and Reflexion agents frequently generate kernels that are slower than the baselines. This advantage becomes even more pronounced Level 2, where the complexity of the kernels increases. Here, STARK maintains a perfect success rate and achieves speedups of 2.7x, whereas the Reflexion agent, despite attaining 100% correctness, produces ker- nels that run slower than the baseline. At Level 3, which involves the most irregular and challenging tasks, both Sampling and Reflexion degrade significantly, with success rates falling and runtimes dropping below baseline. In contrast, STARK continues to maintain full success while producing kernels that outperform the Torch implementations by up to 1.6x. These results highlight that STARK not only generates correct kernels but also delivers substantial performance improvements, even as task difficulty increases.
Torch Eager
Default
Max-autotune
Level
Success
Fast; f
Speed
Fast
Speed
Fast;
Speed
Sampling Agent
57.1% 17 92.6%
14.3%
0.81x
7.1%
0.46x
7.1%
0.81x
Reflexion Agent
28.6%
1.24x
14.3% 78.6%
0.57x
35.7%
0.92x
STARK
100%
71.4%
3.03x
2.37x
78.6%
2.76x
Level
Success
Fast
Speed
Fast,
Speed
Fasti
Speed
Sampling Agent
87.5% 100%
50%
1.06x
37.5% 62.5%
0.91x
37.5%
0.91x
Reflexion Agent
75%
0.88x
0.78x
62.5%
0.78x
STARK
100%
100%
2.69x
87.5%
2.51x
87.5%
2.52x
Level
Success
Fast
Speed
Fast
Speed
Fast;
Speed
Sampling Agent
100%
50%
0.87x
12.5%
0.67x
12.5%
0.66x
Reflexion Agent
67.5%
25%
0.79x
12.5%
0.62x
12.5%
0.61x
STARK
100%
87.5%
1.58x
87.5%
1.27x
87.5%
1.26x
Table 1: Performance of LLM Agents on the KemelBench Tasks. Fasti represents the percentage of problems for which the agent can generate custom kernels that are correct and as fast the Torch baselines (higher is better). Speed is computed the ratio of the kernel runtime of the baseline that of the generated kernel.
Comparison between Agents. We investigate deeper into the behavior of our agent STARK with the two baseline agents to better understand their optimization behaviors. A deeper analysis of
8

---

Published as a conference paper at ICLR 2026
compile and correctness rates, shown in Table provides further insight into why STARK succeeds where baselines struggle. While all agents achieve relatively high compile rates (mostly above 80%), the fraction of kernels that are both compilable and correct varies widely. The Sampling agent, for example, compiles over 90% of its outputs on Level 1 but only 43% of these are functionally correct. Reflexion improves correctness slightly through iterative refinement, but its correctness rate remains below 55% at all levels. In contrast, STARK achieves the highest correctness rates across the board, reaching 61.2% Level tasks. This suggests that STARK 's structured planning and feedback-drive refiement not only increase the chance of generating efficient kernels but also reduce wasted attempts on invalid or incorrect code. Finally, Figure Sigure[]]at highlights the dramatic runtime improvements of STARK relative to baseline agents. On Level tasks, STARK achieves over 10x speedup compared to Sampling and a 13.7x speedup over Reflexion. On Level 2, these gains rise as high as 16x, and even at the most challenging Level 3 tasks STARK maintains 5-6x improvements. These relative gains indicate that while baselines occasionally achieve correctness, they rarely deliver true runtime efficiency. By contrast, STARK 's ability to jointly optimize for correctness and speed allows to close both gaps simultaneously. Taken together, these findings confirm that multi-agent collaboration and strategic search are key enablers for scaling LLMs the demands of GPU kernel optimization.
Compile Rate
Correct Rate
KernelBench Level
2
3
2
3
Sampling Agent
90.8%
97.0%
84.9%
43%
44.0%
15.1%
Reflexion Agent
86.0%
86.2%
78.9%
48.3%
53.4%
28.4%
STARK
84.5%
90.7%
83.4%
50.6%
61.2%
35.5%
Table 2: Percentages of Successfully Compiled and Correct Kernels.
Ablations. We ablate the agentic components of our system. We compare (i) Search Agent, which is single-agent model equipped with our strategic search, and (ii) MA-only, which employs the multi-agent workflow (plan/code/debug with grounded instruction) using best-of- K sampling in- stead of search. As shown in Table 3 both variants outperform the Sampling baseline, confirming that each component helps. When combined in STARK, the effects compound: strategic search ex- ploits the structured proposals produced by the multi-agent workflow, yielding the largest gains. In Appendix [ we verify the robustness
STARK to different foundation LLMs and prompts. We also conduct experiments in Appendix E to test the saturated performance of the baseline agents, showing the necessity of STARK's 's advanced agentic features.
Torch Eager
Default
Max-autotune
Fast, f
Speed
Fast,
Speed
Fast t
Speed
Sampling Agent
50%
0.87x
12.5%
0.67x
12.5%
0.66x
Search Agent
67.5%
0.89x
25%
0.71x
25%
0.70x
MA-Only
67.5%
1.11x
25%
0.92x
25%
0.91x
STARK
87.5%
1.58x
87.5%
1.27x
87.5%
1.26x
Table Ablation on the Proposed Agentic Features.
CONCLUSION
In this work, we introduced an agentic framework for GPU kernel optimization that combines multi- agent role play, dynamic context management, and strategic search. Our evaluation on KernelBench demonstrated that the proposed framework consistently outperforms baseline methods in both suc- cess rate and runtime efficiency, across tasks of varying complexity. These results highlight the value of moving beyond single-agent or unguided sampling approaches, and point to the promise of collaborative, feedback- driven optimization. Looking forward, we envision that agentic LLM frameworks will play an increasingly important role in automated system optimization. Extend- ing our approach to broader classes of operators, diverse hardware architectures, and cross-kernel

---

Published as a conference paper at ICLR 2026
scheduling decisions natural directions for future research. More broadly, work suggests that multi-agent LLMs can meaningfully accelerate the co-design of Al algorithms and infrastructure, pushing the boundaries of what is possible in efficient large-scale computation.
ETHICS STATEMENT
This study complies with the ICLR Code of Ethics, All datasets employed are publicly available and open-source under licenses that permit research use. No private personally identifiable in- formation was accessed, and no new data were collected from human subjects. The research does not pose privacy, security, fairness concerns. The authors declare no conflicts of interest and no external sponsorship.
REPRODUCIBILITY STATEMENT
We document all algorithmic and implementation details in the paper and appendix, including the exact prompts for every agent, full hyperparameters, and ablation settings.
REFERENCES
Antonis Antoniades, Albert Örwall, Kexun Zhang, Yuxi Xie, Anirudh Goyal, and William Wang. Swe-search: Enhancing software agents with monte carlo tree search and iterative refinement. arXiv preprint arXiv:2410.20285, 2024.
Ian Buck. Gpu program optimization. In Matt Pharr (ed.), GPU Gems 2. NVIDIA, 2008
Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Hangu Shen, Meghan Cowan, Leyuan Wang, Yuwei Hu, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. Tvm: An automated end-to-end optimizing compiler for deep learning. In 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18), 2018.
Chris Cummins, Pavlos Petoumenos, Zheng Wang, and Hugh Leather. End-to-end deep learning of optimization heuristics. In 2017 26th International Conference on Parallel Architectures and Compilation Techniques (PACT), pp. 219-232 IEEE, 2017.
with DeepMind.
Discovering
faster
algorithms https: www wwwpwdddedd..eemme cwm blog/
matrix
multiplication
alphatensor.
2022. Accessed: 2024.
Alhussein Fawzi, Matej Balog, Aja Huang, Thomas Hubert, Bernardino Romera-Paredes, Mo- hammadamin Barekatain, Alexander Novikov, Francisco JR Ruiz, Julian Schrittwieser, Grzegorz Swirszez, et al. Discovering faster matrix multiplication algorithms with reinforcement learning. Nature, 610(7930):47-53 2022.
Qiuhan Gu. Llm-based code generation method for golang compiler testing. In Proceedings of the s1 ACM Joint European Software Engineering Conference and Symposium on the Foundations of Software Engineering. pp. 2201-2203, 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and patlern recognition, pp- 770-778, 2016.
Stijn Heldens and Ben van Werkhoven. Kernel launcher: C++ library for portable cuda applications. In 2023 IEEE International Parallel and Distributed Processing Sym- posium Workshops (IPDPSW), pp. 556-565. IEEF, 2023.
Erik Hellsten, Artur Souza, Johannes Lenfers, Rubens Lacouture, Olivia Hsu, Adel Ejjeh, Fredrik Kjolstad, Michel Steuwer. Kunle Olukotun, and Luigi Nardi. Baco: A fast and portable bayesian compiler optimization framework, 2023. URL htt.hhtts:://aa
wwwp:/wwwwwwwoo///oe/oo/oooese// org. 11142
10

---

Published as a conference paper at ICLR 2026
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735-1780, 1997.
Naman Jain, King Han, Alex Gu, Wen-Ding Li, Fanjia Yan, Tianjun Zhang, Sida Wang, Ar- mando Solar-Lezama, Koushik Sen, and Ion Stoica. Livecodebench: Holistic and contamina- tion free evaluation of large language models for code. In The Thirteenth International Confer- ence on Learning Representations, 2025. URL https: https///pperevveeee ee. net. hww.weeeee.rrrrteeeeeeee forum? d=
chfJJYC311.
Carlos E. Jimenez, John Yang, Alexander Wettig, Shunyu Yao, Kexin Pei, Ofir Press, and Karthik Narasimhan. SWE-bench: Can language models resolve real-world github issues? In Interna- tional Conference on Learning Representations, 2024.
Jakub Kurzak, Stanimire Tomov, and Jack Dongarra. Autotuning gemm kemells for the fermi gpu. IEEE Transactions on Parallel and Distributed Systems, 23(1 ((11))2045-2054 2012.
LangChain Inc. Langgraph: Build resilient language agents graphs, 2025. URL https:// Version 0.6.7; accessed 2025-09-24.
github. wwt.ecocccccooo wttp.///aaaagcaaaaaaaaaaaeeoeeea
Teodor V. Marinov, Alekh Agarwal, and Mircea Trofin. Offline imitation learning from multi- ple baselines with applications to compiler optimization, 2024. URL httpp:://xxxii......t
abs /als/2403.20 9462
Cedric Nugteren and Valeriu Codreanu. Citune: A generic auto-tuner for openel kernels. In 2015 IEEE International Parallel and Distributed Processing Symposium Workshop, pp. 846- 851. IEEE, 2015.
NVIDIA Developer. Improving gemm kernel auto-tuning efficiency nvidia gpus with heuristics and cutlass 4.2. httpps://dveeloer..vv..di htw.a/acee...........ooooccaaaaa
2024. Accessed: 2024.
Anne Ouyang, Simon Guo, Simran Arora, Alex L Zhang. William Hu, Christopher Re, and Azalia Mirhoseini. Kernelbench: Can LLMs write efficient GPU kernels? In Forty-second International Conference on Machine Learning, 2025. URL https: /tttps://openreviee net forumorrmioo ?id= veoNliQTI,
Ari Rasch, Markus Haidl, and Sergei Gorlatch. Atf: A generic auto-tuning framework. In 2017 IEEE 19th International Conference on High Performance Computing and Communications (HPPC), pp. 64-71. IEEE, 2017.
Noah Shinn, Federico Cassano, Ashwin Gopinath, Karthik Narasimhan, and Shunyu Yao. Reflexion: Language agents with verbal reinforcement learning. Advances Neural Information Processing Systems, 36:8634-8652. 2023.
Philippe Tillet, Hsiang- Tsung Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, pp. 10-19, 2019.
Mircea Trofin, Yundi Qian, Eugene Brevdo, Zinan Lin, Krzysztof Choromanski, and David Li. Mlgo: machine learning guided compiler optimizations framework. arXiv preprint arXiv:2101 100:22011.4800 2021.
Ben Werkhoven. Kernel tuner: A search-optimizing gpu code auto-tuner. In Fature Generation Computer Systems, volume 90, pp. 347-358. Elsevier, 2019.
Lilian Weng. Llm-pewered autonomous agents.
.o/
post s/ wttss/22233 2202222200066 06-33-66-20 20:00066--2-1/////- 2023. Accessed: 2024.
John Yang, Carlos E Jimenez, Alexander Wettig, Kilian Lieret, Shunyu Yao, Karthik R Narasimhan, and Ofir Press. SWE-agent: Agent-computer interfaces enable automated software engineering. In Tiihth- Conference on Neural Information Processing Systems, 2024.

---

Published as a conference paper at ICLR 2026
John Yang, Kilian Lieret, Carlos E Jimenez, Alexander Wettig, Kabir Khandpur, Yanzhe Zhang, Binyuan Hui, Ofir Press, Ludwig Schmidt, and Diyi Yang. Swe-smith: Scaling data for software engineering agents. arXiv preprint arXiv:2504.21990 2025.
Yi Yang, Ping Xiang, Jingfei Kong, and Huiyang Zhou. A gpgpu compiler for memory optimiza- tion and parallelism management. In Proceedings of the 31st ACM SIGPLAN Conference on Programming Language Design and Implementation, pp. 86-97. ACM, 2010.
Lianmin Zheng. Chengfan Jia, Minmin Sun, Zhao Wu, Cody Hao Yu, Ameer Haj-Ali, Yida Wang, Jun Yang, Danyang Zhuo, Koushik Sen, et al. Ansor: Generating
tensor programs deep learning. In 14th USENIX symposium on operating systems design and imple- mentation (OSDI 20), pp. 813-879, 2020a.
Size Zheng. Yun Liang, Shuo Wang, Renze Chen, and Kaiwen Sheng. Flextensor: An automatic schedule exploration and optimization framework for tensor computation on heterogeneous sys- tem. In Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages and Operating Systems, pp. 859-873, 2020b.
Li Zhong and Zilong Wang. Can Ilm replace stack overflow? study on robustness and reliability of large language model code generation. Proceedings of the AAAI conference on artificial intelligence, volume 38, pp. 21841-21844 2024.
12

---

Published as a conference paper at ICLR 2026
A USE OF LARGE LANGUAGE MODELS (LLMS)
Large Language Models (LLMs) were used solely as general-purpose assistive tools to improve the clarity and readability of the manuseript. Specifically, we used LLM to help rephrase and polish text that we had already drafted.
B IMPLEMENTATION AND EVALUATION
Agent Implementation. We use C1 aude Sonnet 4 as the base LLMs for all agent baselines and agents of STARK. For both sampling and reflexion agents, we follow KernelBench to set temperature T= 0.7 during generating, with other generation parameters such as top-p set to the default value. For STARK, we use Claude Sonnet 4 with temperature T = 0.8 for the plan agent, and T = 0.1 for the code and debug agents. For all tasks, all agent baselines and STARK have a maximum of B = 30 attempts to optimize each task. Regarding the hyperparameters of STARK, we choose the root throtting number to be 5, dead-branch pruning number to be = 0.3 for search, r for the leaderboard C. To prevent exploding context to the LLMs, we set an upper limit for the number nodes in the dynamic context window: whenever the dynamic context window has more than 5 nodes, randomly sample 5 from all the nodes in the window. We implement STARK with LangGraph (LangChain Inc. 2025).
Runtime Evaluation. We evaluate all the Pytorch baselines and LLM -LM-generated kernels on the same NVIDIA A100 40GB GPU. We use the source code provided by KernelBench at its official tep benchmark the kernels' runtime. In particular, to ensure stable measurement, runtime is measured with CUDA events after warm-up runs using fixed input shapes; we choose a large number of 100 ensure accurate measurement.
warm-up
Representative Subset of KernelBench. Due to resource constraints, we use the pro- vided list of representative problems of KernelBench to evaluate runtime of kernels. Specifi- cally, we use problems 1, 1,3,6 6, 18, 23, 26, 33, 40, 42, 48, 54, 57, 65, 77, 82, 86 of Level 1, problems 1,2, 8, 18, 23, 28, 33, 43 of Level and problems 7, 15, 18, 24, 28, 39, 45, 50 of Level
C PROMPTS
Our prompts follow the templates of KernelBench (Ouyang et al. 2025), which has four compo- nents: system message, in-context example, architecture source code, and instruction.
As we have multiple agents STARK with different tasks, they require different prompts to fulfill their tasks. Specifically, we reuse the system message and in-context example from KernelBench for all agents and include the architecture source code regardless of which node is selected for opti- mization. To motivate the agents an additional instruction in the system message to consider using existing highly optimized kernels. We show the system prompt in Figure [5] the in-context example in Figures I1 and 12 In addition, we include the information within the dynamic context window and role-specific instructions for dif- ferent agents. See Figure 4 for the prompt template of STARK. We show the role-specific instruction
the already optimized modules such cuBLAS,
include
in Figures 8 9 and 10.
D EXAMPLE KERNELBENCH TASKS
Here show three examples of KernelBench tasks, for each level. See Figures 13 and 15 for example tasks in Level 1, 2, and 3. We refer interested readers to Ouyang et al. (2025) for the complete list.
E EXTRA ABLATIONS
We conduct extra ablations using the representative Level 2 problems of KernelBench.
httpes:///iituubbccmm//////ll hag. Hench

---

Published as a conference paper at ICLR 2026
(System Message)
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architesture is: (Example Architecture Source Code)
The example new arch with custom CUDA kernels looks like this: (Example New Architecture Source Code)
You given the following architecture: (Arcaitecture Source Code]
10
I
Here is your latest attempt: (Source Code of the Selected Node)
12
13
14 [Dynamic Context Window] You should use the following observations zegarding your historical attempts to provide better implementaiions:
Learn from the failed examples to avoid bugs and write successful kernels. Learn from the successful examples to design improved kernels.
16
17
18
**Kernel Source Code #1** (Source Code of Historical Actempt)
19
20
21
**Compiler Observation**
22
[Compiler Log]
24
**ernel Execuation Result**
35
[Runtime or Correctness Brror)
26
**Kernel Source Code #2**
skipped)
29
30
(Role-specific Instruction)
Figure 4: Prompt Template for Agents.
Robustness of STARK to the Foundational LLMs. While we focus on instantiating STARK using Claude Scnnet 4, we test the robustness of STARK by instantiating baseline agents and STARK with OpenAI GPT-4. In Table 4 we observe that STARK 's gain remains with the more economic OpenAI GPT-4.1 model.
Table 4: Robustness STARK when instantiated with OpenA models
Metric
Sampling Agent
Reflexion Agent
STARK
Fast,
75
50
87.5
STARK's Speedup over baselines
4.2x
5.7x
Robustness of STARK to the System Prompt. We ask OpenAI GPT-5. to generate two differ- ent yet semantically similar alternatives to the KernelBench system prompt. The resultant prompts shown in Figures and In Table we observe that STARK 's performance is robust to varying system prompts.
Table 5: Performance Comparisons with Varying System Prompts.
Metric
STARK
STARK (Prompt V1)
STARK (Prompt V2)
Fast, f
87.5
87.5
87.5
Speedup over the reflexion agent
16x
14.5x
17.1x
14

---

Published as a conference paper at ICLR 2026
System Message
You are an expert in writing efficient coce. You write custom CUDA kernels to replace the pytorch operators in the given architecture get speedups.
You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged, You may replace multiple opezators with custom implementations, consider opezator fusion opportunities (combining multiple operators into single kernel, for example, combining matmul+relun or al gorithmie changes (such online softmax), You are only limited by your imagination.
You should consider using CUDA' s existing highly optimized kernels and operations whenever appropriate. Try building on these optimized blocks and further improve it with your custom kernels.
Figure System Message for All Agents.
# System Message
LLM agent specializing in GPU optimization. Your goal is to speed up given PyIorch architecture by generating custom CUDA kernels to replace its existing operators.
You
You may freely choose which operators to target, whether to rewrite operator, many operators,
none. You can pursue kernel fusion opportunities (e.g., folding linear layers and elementwise ops together) or restructure the computation algorithmically to achieve higher throughput or lower memory traffic.
There are no constraints on what transformations you may propose- innovation is encouraged.
Use CUDA's s existing highly optimized kernel libraries whenever advantageous; treat them as strong baselines that your custom kernels can refine, fuse, or extend to achieve even greater efficiency.
Figure 6: Alternate System Message V1 for STARK.
Performance of the Baseline Agents with Extra Searching Budgets. To ensure the fairness of experiments and to test whether the advanced agentic features of STARK are truly necessary, conduct an ablation where the baseline agents use 4x more attempts than STARK, making them equal in API costs. In Table [ we observe that, even with 4x more attempts, the baseline agents cannot effectively enhance their performance due to severely diminishing return. In particular, the reflexion agent keeps stuck in similar local optimums and struggles to improve while the sampling agent's is efficiency in generating strong kernels is very low so that the Fast, rates remain unchanged, and the fastest generated kernels have minimal speed improvement. These results exactly demon- strate the challenge of GPU kernel optimization, and advanced agentic features such as the ones proposed by STARK necessary for effectively finding better GPU kernels with more budgets.
F RELATED WORK
Optimizing GPU kernels extract maximum performance from underlying hardware is long- standing and formidable challenge. The vast, non-convex, and hardware-specifie search space of possible kernel implementations necessitates sophisticated optimization strategies. The evolution of
15

---

Published as a conference paper at ICLR 2026
System Message
You are GPU kernel engineer tasked with accelerating a FyTorch model by selectively replacing its operators with custom CUDA kernels.
You have full discretion in deciding which operators should be rewritten, fused, or left as-is. You may implement multiple custom kernels, explore operator Eusion (e.g.r matmul + activation), or introduce more efficient algorithmic variants (such as online softmax or tiled reductions).
Your design space is unrestri unrestricted stdedetttte-eeittty and performance intuition should guide your choices.
Waenever beneficial, Leverage CUDAs high-perforrance primitives tensor cores, outlass, CuBLAS, cuDNN, etc.). Build upon these optimized components and push performance further with your own implementations.
Figure Alternate System Message V2 for STARK.
Table 6: Performance of Baseline Agents with Extra Budgets.
Sampling Agent
Reflexion Agent
STARK
Metric
Base
4x
Base
4x
Fasti f
37.5
37.5
62.5
62.5
87.5
STARK's Speedup over baselines f
2.35x
2.31 x
16x
16x
these strategies can be broadly categorized into three paradigms: empirical auto-tuning, compiler- and model-driven optimization, and most recently, generative approaches using Large Language Models (LLMs). Our work builds upon this trajectory by introducing fully autonomous agent that manages the entire optimization lifecycle.
F.1 EMPIRICAL AND COMPILER BAPSLEER-BBEEDD OPTIMIZATION
The foundational approach to GPU performance tuning is empirical auto-tuning, which treats the problem as a black-box search over a set tunable parameters, such as thread block dimensions, memory tiling factors, and loop unrolling factors (van Werkhoven 2019). Traditional methods often rely on an exhaustive or brute-force search, where thousands of potential kernel configurations are generated, compiled, and benchmarked identify the top performer (Kurzak et al.] 2012). While While effective, this process is prohibitively time-consuming: for instance, an exhaustive search for single GEMM kernel can take over 700 minutes to complete (NVIDIA Developer 2024).
To mitigate this cost, heuristic-driven methods prune the search space. NVIDIA's nvMatmul Heuristics, for example, uses a predictive model recommend a small subset of configurations, achieving near-optimal performance in fraction of the time (NVIDIA Developer 2024). Frameworks like Kernel Tuner (van Werkhoven, (019). ATF (Rasch al. 2017). and CLTune (Nugteren & Codreanu (015) provide robust environments for orchestrating these searches and support more advanced strategies like Bayesian Optimization, which builds prob-
abilistic performance model to guide the search more intelligently (Hellsten et al. 2023 Heldens & van Werkhoven 2023).
Concurrently, compiler-based approaches aim to automate optimization through series of program transformations applied to an intermediate representation (Yang et al. 2010). GPU compilers em- ploy passes for memory coalescing, data prefetching, vectorization, and loop optimizations adapt naive code to the hardware architecture (Buck) (008).. While these approaches excel at finding op- timal configurations within predefined search space, they cannot fundamentally alter the kemnl's
16

---

Published as a conference paper at ICLR 2026
Instruction
Optimize the architecture named Model with custom CUDA operaters
Give explicit and actonable advice to improve the efficiency, in terma of the GPU wall-clock time, the architecture named Model
Give ONE advice of the top priority! Don't over-request. Include necessary details such as how to change pointers and indices or how to achieve shared memory tiling so that your acvice can be correctly implemented.
After your advice, modi fy and return the given source code the following way: Identify the code block whose efficiency can be improved that is where your advice should be implemented)
10
and mark it with comments <<<IMPROVE BEGINES> the beginning and <<<IMMPOVE END>>> The markers <<<<IMPROVE BEGIN>>> and '<<<IMPROVE END>>>' should be valid comments for the marked coding language. For example, when marking source code of custom kernels, you need to comments for the C++ language /// <<<IMPROVE END>>>'; when marking source code of Python, you should use '# <<<ImpRove BEGIN>>>' and '# <<<IMPROVE END>>>' Add your advice as comments at the identilied code block to help the following agent's implementation
the end
12
'// <<<IMPROVE BEGIN>>>' and
13
14
There will be another agent focusing on improving the efficiency of the identified code block.
15
Return the complete code block with the identified code block as its subpart.
17
You should consider using CUDA' s existing highly optimized kernels and operations whenever appropriate. Try building on these optimized blocks and further improve it with your custom kernels. When presented with multiple prior attempts, you should consider exploration of more diverse optimization strategies.
18
19
Pay careful attention to the impl implementation agent's capability demonstrated from the historical implementations, Adjust your advice accordingly to ensure that it can successfully implement.
20
Figure 8: Instruction for the Plan Agent.
algorithm. Our work introduces an agent that reasons about performance bottlenecks to implement novel, structural code changes, moving beyond simple parameter tuning.
F.2 MACHINE LEARNING FOR CODE OPTIMIZATION
Machine learning (ML) has emerged as a powerful tool to transcend the limitations of hand-crafted heuristics. Early work focused on using ML to make better decisions within existing compiler and tuning frameworks. Systems like TVM employ learned cost model to predict the performance of kernel variants, guiding the search process and avoiding exhaustive empirical evaluation (Chen et al.) 2018 More recent efforts have integrated ML directly into production compilers. Google's MLGO framework uses reinforcement learning (RL) to train policies for classic compiler optimizations like function inlining and register allocation, demonstrating significant improvements code size and performance over decades-old, manually-tuned heuristics in LLVM (Trofin al.) 2021 Marinov ett al., 2024). These models can leam from massive code corpora and discover complex feature interactions that opaque to human experts (Cummins et al.] 2017).
17

---

Published as a conference paper at ICLR 2026
Instruction
Optimize the architecture named Model with custom CUDA operators!
Think about the given advice from human experts and implement the ones that you believe are correct and you are confident implementing.
Only focus the code block marked with <E<IMPROVE BEGIN>>> and "<<<IMPROVE END>>>'
Write custom cuda kernel to replace the pytorch operators within the marked code block to improve its efficiency
Name your optimized output architecture ModelNew. Output the new code in codeblacks.
10
explain your Implemmntation and how you follow the advice.
11
Using the given tool to return your final structured answer.
12
13
Please generate real code, NOT pseudocode, make sure the code compiles and is fully functional.
14
15
NO testing code!
Figure 9: Instruction for the Code Agent.
Instruction
Fix the issues of your impl implemmnttaioon on named ModelNew, whigh should improve efficiency the source model named Model.
ModelNew and Model should have the same functionality, that is, the same input-output mapping, The given architectuse ModelNew ieither does not campile, or has zun-time error, or has different functionality to the source Model.
Use the giver observations to infer bugs and then fix them. Explain how the bugs happen and how you fix it. Return the fixed bug-free code and your optimized output architecture ModelNew.
Figure 10: Instruction for the Debug Agent.
more profound application ML has been in algorithmic discovery. DeepMind's AlphaTen-- framed the search for faster matrix multiplication algorithms single-player game, using deep RL agent based on AlphaZero to navigate the enormous search space of tensor decomposi- tions (Fawzi that outperform human-derived state-of-the-arr methods, including improving upon Strassen's 's algo- rithm for 4 matrices for the first time in over 50 years (Fawzi al. 2022 DeepMind (022). This work marked a critical shift from using ML to configure cxisting optimization strategies to using it invent new ones from first principles. However, AlphaTensor operated in a clean, formal mathe- matical domain. Translating this power to the messy, syntactic, and
2022). This approach successfully discovered novel, provably correct algorithms
domain of GPU kernel programming presents distinct challenge. Our work addresses this by employing an agent that operates directly on source code, navigating the complexities of syntax, compilation, and hardware-speccific performance characteristics.
F.3 LLM-POWERED AUTONOMOUS AGENTS
The capabilities of LLMs have given rise to a new paradigm of autonomous agents. An LLM agent uses core LLM as its "brain" controller, augmented with capabilities for planning, memory, and tool use to perform complex tasks autonomously (Weng 2023). The key distinction from simple L.M prompting is the agent's ability to decompose high-level goal into sequence of manageable
18

---

Published as a conference paper at ICLR 2026
## Example Architecture Source Code import torch
import torch.or as nn import torch.nn. functional F
class Model (
def _init__ (self) -> None:
super () init ()
10
11
def forward(self, a, b): return a b
12
15
get iot_inputs(( ():
16
# randomly generate input tensors based on the model architecture
corch.randn (1, 128) cuda () torch. randn (1, 128) cuda () return [a, b]
18
19
20
def get_init_irputs
23
# randomly generate tensors required for initialization based
on the model architecture
return 24
Figure 11: In-context Example Architecture for All Agents.
subtasks, execute them iteratively, and use reflection to gauge progress and self-correct. This agentic workflow involves the L.M interacting with an external environment through set of tools, such as a code interpreter or web search API, to gather information and perform actions. While this paradigm powerful for general problem-solving its application to specialized domains like soft- ware engineering requires tailored tools and reasoning processes. Our work specializes this agentic concept for the domain of performance optimization, which presents unique challenges not found in general-purpose purpose agent tasks.
F.4
LLMs FOR CODE OPTIMIZATION AND GENERATION
The advent of powerful LLMs has opened a new frontier in performance engineering. To grant LLMs greater autonomy, the agentic paradigm has been adapted specifically for software engineer- ing. The success of systems like SWE-agent, which autonomously resolves complex bugs in large GitHub repositories,
demonstrated the viability of this approach (Yang et al. 2024). SWE-agent equips an L.LM with specialized Agent-Compute Interface (ACI) containing tools for file naviga- tion, editing, and test exccution, enabling it to perform simple code generation (Yang et al. 2024 Jimenez et al. 2024p. While these agents icant step towards autonomous software engineering, their focus has primarily been on functional correctness, such bug fixing. Our work extends this agentic software engineering paradigm to the non-functional, performance-orieen agent that not only interacts with codebase but also with hardware profiling tools, allowing to autonomously diagnose performance bottlenecks, form hypotheses, and conduct experiments to iteratively improve kernel efficiency, thus acting
tasks far beyond the scope of
signif-
domain of GPU kernel optimization. We introduce an
autonomous performance engineer.
19

---

Published as a conference paper at ICLR 2026
## Example New Architecture Source Code
import torch
import torch. as an import torch.nn, from worch.utils.com werch..tils.cppeexteaeee cpp_extension import Load inline
wrr.cttioon functional as F
nefine the custom CUDA kernel for element-wise addition seleenttiis__ddd_sorrre "*"
10
Vinclude #include <cuda runtime.r.
global void elementwise anddise__dd_____eeeee kernel (const float* a, const float 13
* b, float* out, int size) {
14
int 1dx = blockidx.x blockDim.x + threadIdx.x7
if (idx size)
out[idx] a[idx] b[idx];
17
18
19
torch: Tensor elementwise add cuda (ttorch::eenae a, torch: Tensor
b)) {
21
auto size a. nunumel (); auto out torch: [
22
const int block size 256;
25
const int num_biocks (size block_size 1) block_size/
26
27
biock (a. ptrr float>(), b. d. out (tt size);
28
29
return out;
30 }
33 elementwise add cpp app_source
34 "torch;: Tensor elementwise _uda Tenaor b);"
cuda (torch: : Tennor a, torch::
36
37
# Compile the inline CUDA code for element wise addition
38
el elementwise_add load_inling inline(
name==aalemettwiss__ddee
source, oudaaeooureeeeemmnttttteeteeeeee
41
42
43
verbose-Irue,
extra_eflags=
extra_ldflags-("1)
46
47
48
class ModelNew (Mrdule)
50
def (self
None:
51
super in__ttit______
52
self. halffeelomentwise_dde = elementwise_add
def (self, a, return sel felem.tiii wtw.eeeewwisss...deeeeeeteeeeeee
55
(a, b)
Figure 12: In-context Optimized Example Architecture for All Agents.
20

---

Published as a conference paper at ICLR 2026
import torch
import torch. nn
class Model (nn. (onuled
Simple model that performs LogSoftmax activation. "*"
def _nnit_ dim: int 1): super (Model, .init___ii____
self.din dim
forward(self, x: torch. Tensor) -> torch. Tensor: ""*
13
Applics LogSoftmax activation to the input tensor.
Args:
x (torch. Tensor): Input tensor shape (batch_size, dim)
Returns:
torch. Tensor: Output tenser with LogSoftmax applied, same shape as input.
return torch torch.log_sofcmaxxt (x, dim-self.dim)
batch_ size 4096
25
dim = 393216
def get_inputs ();
torch. rand (batch_size, dim)
29
return [x]
30
def set_init_io [i_iintt_inptss ():
return $ special initialization inputs needed
Figure 13: Example KernelBench Level Task.
21

---

Published as a conference paper at ICLR 2026
import torch
import torch.nn nn as nn
class Model (nn. Monule);
*""
Model that performs matrix multiplication, division, summation, and scaling.
def
init_ input_size, hidden_size, scaling_factor super (Model, self) self)__init__) self.weigt Pn. (torch. (hidden_siae input_size)) self.scaling factor factor scaling factor
10
def forward(self, x):
Args:
(torch. Input tensor of shape (batch_size, input_size).
Returns:
18
torch. Tensor: Output tensor of shape (batch_size, hidden_sizee.
torch.matmin matmul (x, self. weightt))h Genm
21
x = x / # Divide
- torch.su sum (x, dim-1, keepdim-True) # Sum
self.scaline self.scaling_ffctoo # Scaling
return
26
batch_size 1024
input_size 8192
hidden_size 8192
30
scaling_factor 1.5
def get_inputs ():
return [torch. rand (batch_size, input_size)
35 def get init inputs ():
return linput size, hidden size, cealing.facton
Figure 14: Example KernelBench Level Task.
22

---

Published as a conference paper at ICLR 2026
import torch
import torch.nn as nn import torch.nn. functional forchion..unctionnll as F import math
class NewGELU (nn. :onule)
Implementation of the GELU activation function currently in Google BERT repo (identical to OpenaI GPT)
Reference: Gaussian Error Linear Units (GELU) paper: https://
azxiv..rg//bss/66666 0844 08415
10
def (init_(se1f)
super (NewGELU, self) f) init_
14
def (sew,, x)
return 0.5 (1.0 + torch. (math. sqrt (2.0 / math.pi (x 0.044715 torch.pow (x, 3.0))))
17
class CausalselfAttcntioon (an. (on.
18
A vanilla muiti-hes masked self-attention layer with projection the end.
is possible to use torch.nn. hwt.:.mm.attihhhhttteeeeeee here but
am including an explicit implementation here to show that there is nothing too scary here.
"""
23
24
def _ (self, n_embd, n head, attn petn_pproo, resid pdrop, max :ex_seelen):
super _ init
assert r_embd $ n_hoad
27
$ key, query, value projections for all heads, but in batch
self.c_attn Ln. (n_embd, 3 n_embd)
output projection
30
self.c_proj linerr (n_embd, n_embd)
# reguiarization
self dropop = nn. Dropout (attn pdrop) (attn_pdrop)
self drepout nn.Dropout (resid, (resid_pdrop) pdrop)
causal mask to ensure that attention only applied to the left in the input sequence self. max max_seclen, max_seqlen))
("bias", torch.trih tril (torch. ones
.view max_seqlen,
max_ max_seqlen)
37
self.n_head head self.n_emte n_embd
skipped)
class Model (nn. Monuuedule): an unassuming Transformer block
42
skipped
46
def forward x):
= + self.attn (s (x) self.mlpt 1 (x)
return
50
[... sk ipped]
Figure 15: Example KernelBench Level 3 Task.
23


Reviews:
[{"review_id": "iJh7vpc0WS", "reviewer": "Reviewer_sbZq", "summary": "This paper introduces STARK, an LLM-based agentic framework for automated GPU kernel optimization, addressing the inefficiency of existing methods (manual optimization, automated compilers, single-agent LLMs) in navigating irregular kernel design spaces. STARK’s core innovations include: a collaborative multi-agent workflow (plan/code/debug agents with role-specific temperatures), grounded instruction (anchoring optimization plans to code spans) and dynamic context windows (agent-specific historical feedback), and strategic search (tree memory with adapted ε-greedy policy to balance exploration/exploitation). Evaluated on KernelBench (3 difficulty levels), STARK outperforms baselines (Sampling/Reflexion agents, Torch Eager/Compile) by achieving 100% success rate across levels, up to 16× speedup over Reflexion, and resolving challenging Level 3 tasks where baselines fail. It demonstrates that multi-agent collaboration and feedback-driven search unlock LLMs’ potential for scalable, fully automated GPU kernel optimization.", "questions": "As shown in the weakness.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Innovative multi-agent framework decomposes tasks into planning, coding, debugging roles with tailored parameters. Combined with grounded instructions and dynamic context, it bridges the planning-implementation gap, achieving 100% success rate across all KernelBench levels.\n2. Strategic search with tree memory uses an adapted ε-greedy policy to avoid redundancy and local optima, gaining up to 16× speedup over baselines. Ablation studies confirm gains from its combination with multi-agent workflow.\n3. Evaluated on multi-difficulty KernelBench, compared with multiple baselines, and assessed via metrics like success rate. It fully verifies performance and provides a clear benchmark for similar research.", "weaknesses": "1. Evaluated exclusively on NVIDIA A100 GPUs, with no validation on other architectures (e.g., AMD GPUs, edge AI hardware).\n2. Lacks head-to-head testing against state-of-the-art ML/RL-based optimizers (e.g., AlphaTensor, TVM’s Ansor) and modern LLMs (e.g., GPT-4, DeepSeek-R1).\n3. Provides no analysis of why STARK’s strategies (e.g., grounded instruction) work, no qualitative breakdown of optimization decisions or failure modes.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lvkzkrwhWY", "reviewer": "Reviewer_tSQ9", "summary": "STARK proposes multi agent framework for iterative refinement GPU kernels for performance optimization. STARK consists of 3 different agents namely for planning, coding, and debugging implemented using the same LLM claude-4-sonnet but distinct sampling parameter e.g., temperature. Planning agents prepares a plan to explore optimization strategies. Coder agents works on top of planning agent output to produce code that best aligns with strategies produced by the planning agent. Debugging agent repairs the code based on execution and hardware feedback. STARK leverages grounded instructions to guide planning/coding agent to suggest/generate edits in specific parts of the code as desired. Furthermore, dynamic context window is leveraged to feed relevant but contrastive in-context learning samples to all planning, coding, and debug agents. STARK populates dynamic context window by maintaining a persistent tree of prior attempts (generated code) to guide future planning/coding attempts to produce better kernels. STARK compares against torch eager & compile mode and reflexion agent-based implementation. STARK achieves 100% success rate across first 3 levels of kernelbench benchmark and outperforms baseline run time by up to 3x.", "questions": "- How does this scale to low resource GPU kernel languages such as Triton?\n- How does strategic tree search differ from ABMCTS? What is the comparison in performance w.r.t. ABMCTS?\n- Why are siblings (S(i)) are expected to have similarities? Some evidence indicating its importance in Wdebug and Wcode would be useful.\n- How does hardware feedback improve the overall framework output further?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- Multi agent realization of one of the toughest coding tasks of GPU kernel generation.\n- Clear division of work across 3 different agents and their seamless integration to produce better kernel.\n- Tree of prior attempts that helps to systematically identify good candidates for iterative refinement instead of blind refinement. \n- grounded instructions and dynamic context windows that help carefully plan and target code edits.", "weaknesses": "- No hardware profiler feedback integration. Since hardware feedback can give lot more insight to planning agent.\n- Though appendix lists some newly generated codes, the paper lacks detailed case study on how these agents end up discovering strategies that beat even torch.compile.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jS7jth6MSj", "reviewer": "Reviewer_QJHD", "summary": "The paper proposes STARK, a multi-agent LLM framework for GPU-kernel optimization. It splits work into plan / code / debug agents, coordinates them with grounded instructions. On KernelBench, STARK markedly outperforms baseline agents in both success rate and runtime speed.", "questions": "1. How would STARK perform if adding another optimization/tuning agent?\n2. How sensitive are results to ε, root throttling cap, dead-branch threshold, and context-window size?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The plan/code/debug split is well-motivated and practical.\n2. The grounded instruction concretize edits and tighten plan–implementation alignment.\n3. The empirical results demonstrates kernels generated by STARK having consistent improvement in Fast1, success rate, and speed, as well as compile and correct rate.", "weaknesses": "1. All agents use Claude Sonnet 4; it’d help to show robustness to other base LLMs and to ablate prompt dependencies.\n2. The multi-agent workflow + search adds overhead. Wall-clock optimization time (not just kernel runtime) is not reported.\n3. STARK outperforms baseline agent kernels, but it's unknown its kernels' performance against expert written kernels.\n4. (misc) Typos such as \"Kernel Execuation Result\" and \"Due to space constrain\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cmCWpWdMA2", "reviewer": "Reviewer_5yP2", "summary": "This paper introduces STARK, a multi-agent LLM framework for GPU kernel optimization. STARK adopts a collaborative multi-agent workflow with specialized planner, coder, and debugger agents operating at different temperatures, and grounded instruction and dynamic context windows to leverage historical context. The better search algorithm with tree memory uses an adapted ε-greedy policy to balance exploration and exploitation. The proposed approach outperforms sampling agents and reflexion agents.", "questions": "1. How novel is the strategic search with tree memory? Can you compare your approach with the literature? How big is the delta here? I understand that the kernel optimization only uses best-of-K sampling / iterative refinement before (as you wrote in Line 229).\n2. How does this work compare against \"Astra: A Multi-Agent System for GPU Kernel Performance Optimization\" in terms of the proposed method? Can you summarize and discuss the differences/novelty compared to this work?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This paper targets the important problem of GPU kernel optimization.\n2. The empirical gains over other baselines are significant.", "weaknesses": "1.  The ablation study does not seem to be enough. E.g., is \"Grounded instruction\"'s ablation study separately conducted? Is each agent's separate role ablated?\n2. I believe directly using Claude Code (as an agentic baseline, not just a foundational model) can be an important baseline to consider adding.\n3. Any manual study on the code changes that the agents make?\n4. Is the comparison and evaluation fair? Are all baselines using the same computing resources (e.g., same API costs)?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)