Skip to content

STARK: Strategic Team of Agents for Refining Kernels

🕒 Published (v1): 2026-01-01 · Source: ICLR · Venue: ICLR 2026 · link

Why this paper was selected

Multi-agent team discovers and applies GPU kernel optimizations automatically

Ask a follow-up

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

💬 Ask ChatGPT✦ Ask Claude

TL;DR

STARK is a multi-agent LLM framework for automated GPU kernel optimization that decomposes optimization into specialized plan, code, and debug agents coordinated via a tree-memory strategic search policy and a grounded instruction mechanism that anchors high-level edits to concrete code spans. Evaluated on KernelBench, it achieves up to 16× runtime speedup over the Reflexion baseline and 10.7× over the Sampling baseline while reaching 100% success rate across all difficulty levels where baselines frequently fail.

Problem

LLM-based kernel optimizers have three critical deficiencies: (1) naive exploration — linear iterative refinement ignores prior attempt history and traps optimization in local optima; (2) monolithic agent design — a single generalist agent cannot simultaneously excel at creative strategy generation (requiring high temperature/diversity) and precise CUDA code emission (requiring low temperature/fidelity); (3) planning-implementation gap — LLMs routinely propose correct high-level strategies (e.g., "apply memory tiling") but fail to translate them into valid low-level CUDA code, largely due to scarcity of expert-level kernel code in pretraining data.

Method

STARK operates in three integrated layers:

1. Multi-Agent Design (MAD). Three role-specialized agents share a base LLM (Claude Sonnet 4) with role-specific temperatures: - Plan agent (\(\tau = 0.8\)): proposes optimization strategies with grounded instruction anchors. - Code agent (\(\tau = 0.1\)): translates anchored plans into executable CUDA kernels. - Debug agent (\(\tau = 0.1\)): repairs compile/runtime failures using compiler diagnostics.

2. Strategic Search with Tree Memory. Kernel optimization is recast as search over a persistent tree \(\mathcal{T}\) whose nodes store candidate kernels and their evaluation outcomes (runtime, correctness, compiler logs). An adapted \(\epsilon\)-greedy policy selects nodes to expand, with four domain-specific modifications: (i) root throttling — cap direct children of the root at \(n_\text{root}\) to avoid redundant first-hop edits; (ii) dead-branch pruning — mark nodes with \(> n_\text{child}\) all-failing children as ineligible; (iii) high exploration rate — \(\epsilon \in [0.3, 0.4]\) to escape local traps; (iv) leaf-biased exploration — sample uniformly from expandable leaves with probability \(\epsilon\).

3. Grounded Instruction. The plan agent annotates the kernel source with span anchors (<<<IMPROVE BEGINS>>> … <<<IMPROVE ENDS>>>) marking exact sites for modification (e.g., a load/store, loop body, launch configuration). The code agent resolves each anchor into concrete CUDA. This narrows the coder's search space, reduces hallucinated guidance, and creates verifiable edit traces.

4. Dynamic Context Windows. Each agent receives a role-specific view of history at every step: - Plan agent: \(\mathcal{W}_\text{plan}(i) = \{i, n_\text{root}\} \cup D(i) \cup \text{Top-}r(\mathcal{C})\) — children of the selected node plus global leaderboard leaders for ambition calibration. - Code agent: \(\mathcal{W}_\text{code}(i) = \{i, n_\text{root}\} \cup D(i) \cup \{j : p(j) \in S(i)\}\) — siblings' children for cross-pollinating successful micro-optimizations. - Debug agent: \(\mathcal{W}_\text{debug}(i) = \{i, n_\text{root}\} \cup S(i)\) — sibling nodes sharing the same scaffold to transfer structural fixes.

Key Contributions

  • Multi-agent workflow (plan/code/debug) with role-specific temperature settings enabling simultaneous diversity and precision.
  • Adapted \(\epsilon\)-greedy search over a persistent tree memory with root throttling, dead-branch pruning, and leaf-biased exploration tailored to the kernel optimization landscape.
  • Grounded instruction: span-anchor mechanism bridging the planning-implementation gap with verifiable, localized code edits.
  • Dynamic, agent-specific context windows that surface role-relevant history (local lineage, global leaderboard, sibling scaffolds) without cross-contaminating agent objectives.
  • Evaluated on KernelBench L1–L3, establishing new state-of-the-art in both correctness and runtime efficiency.

Results

Evaluated on a representative KernelBench subset, budget \(B = 30\) attempts, using Claude Sonnet 4 for all agents and baselines:

  • Level 1 (single operators): STARK 100% success vs. 57.1% (Sampling), 92.6% (Reflexion); 3.03× speed over Torch Eager vs. 0.81× (Sampling), 1.24× (Reflexion); 71.4% Fast1 vs. 14.3% / 28.6%.
  • Level 2 (fused multi-operator kernels): STARK 100% success vs. 87.5% (Sampling), 100% (Reflexion); 2.69× speed over Torch Eager vs. 1.06× (Sampling), 0.88× (Reflexion); 100% Fast1 vs. 50% / 75%.
  • Level 3 (full ML architectures, e.g., ResNet, LSTM): STARK maintains 100% success and achieves up to 1.6× over Torch Eager; Sampling and Reflexion degrade significantly in both success rate and speed.
  • Wall-clock speedup over agents (Figure 1): STARK reaches up to 16× over Reflexion (L2), 13.7× over Reflexion (L3), 10.7× over Sampling (L1), and 6.54× over Sampling (L3).
  • \(\epsilon\)-greedy outperforms MCTS, evolutionary, and greedy policies under the same budget.

Limitations

  • Evaluated on a representative subset of KernelBench due to compute constraints; full-benchmark results not reported.
  • Agent-specific post-training (e.g., fine-tuning the code agent) is identified as a promising direction but not explored.
  • All experiments use Claude Sonnet 4; generalization to other base LLMs is not systematically studied.
  • The \(\epsilon\)-greedy policy comparison is empirical; theoretical analysis of the search strategy is absent.
  • Level 3 speedups (1.6×) are substantially lower than L1–L2, suggesting diminishing returns for highly irregular architectures.

Relevance to Agentic AI / LLM Agents

STARK is a high-quality instance of the specialized multi-agent architecture paradigm, demonstrating that decomposing a complex optimization task across agents with complementary temperature/capability profiles (creative planner + precise coder + structural debugger) yields large gains over both single-agent and naive refinement baselines. The grounded instruction mechanism — using span anchors to translate abstract plans into concrete, verifiable edits — is a transferable design pattern for any agentic workflow where LLMs must bridge high-level intent and low-level implementation. The tree-memory \(\epsilon\)-greedy search represents a principled middle ground between unguided best-of-K sampling and myopic iterative refinement, with domain-specific adaptations (root throttling, dead-branch pruning) that generalize to other irregular optimization landscapes. This work strengthens the case that structured agentic orchestration with shared memory and role specialization is a key lever for LLM performance on expert-level engineering tasks.