AccelOpt: A Self-Improving LLM Agentic System for AI Accelerator Kernel Optimization

• Emerging accelerators such as Trainium make kernel optimization difficult because hardware and programming models are new, established heuristics are scarce, and effective tuning requires costly exploration of a vast design space.Kernel efficiency matters because suboptimal implementations can constrain performance and waste compute and money at scale.Source: §1 · pp. 1–2
• AccelOpt combines beam search with a planner-executor-summarizer workflow and optimization memory to generate and profile kernels, retain strong candidates, and reuse insights without expert-provided hardware-specific recipes.The memory stores code changes from slow-to-fast kernel transformations and LLM-generated general optimization insights.Source: §1 · pp. 2–3
• AccelOpt discovers local and non-trivial global optimizations; beam search outperforms repeated sampling, while memory affects cost efficiency but not best-kernel performance when enough kernels are sampled.The system also evaluates cost-benefit trade-offs across optimization-memory configurations and base LLMs.Source: §1 · pp. 2–3
• Average percentage of peak throughput rises from 49% to 61% on Trainium 1 and from 45% to 59% on Trainium 2 using open-source LLMs.On Trainium 1, performance is on par with Claude Sonnet 4 (thinking mode) at 26× lower cost; Trainium 2 cost is similar.Source: §1 · p. 3
• NKIBench is the first benchmark suite for NKI kernel optimization on Trainium, covering real-world LLM kernels and measuring performance against theoretical peak rather than only relative speedup.Its current version contains 14 kernels, establishing a starting point for NKI optimization research.Source: §1 · pp. 2–3

2.1 Algorithm Overview

• AccelOpt’s planner, executor, and summarizer iteratively generate kernel candidates, use profiling to guide optimization, and distill strategies from slow-fast pairs.The process updates both the candidate frontier and optimization memory each iteration.Source: §2.1 · p. 3

2.2 Beam Search

• Experiments find beam search more effective than simple parallel (repeated) sampling for open-source LLMs.Source: §2.2 · p. 3
• Each round expands B candidates into B × N × K attempts; β selects each plan’s fastest correct kernel, then takes the top B by latency and backfills prior candidates if needed.Source: §2.2 · pp. 3–4

2.3 Optimization Memory Curation

• Optimization memory stores threshold-filtered positive and negative slow-fast pairs with summarized strategies, selecting performance outliers by candidate and plan.It appends up to TopK items per iteration to a capped queue; this is per-problem test-time learning, while transferring memory across kernels remains future work.Source: §2.3 · pp. 4–5

3.2 Profiling Service

• NKIBench pairs challenging real-workload kernels with structural storage and distributed profiling, using task-, sample-, core-, and machine-level parallelism to scale evaluation.A centralized manager dispatches requests across Trainium machines connected by a shared network file system.Source: §3 · p. 5, §3.2 · p. 6

3.1 NKIBench Task Construction

• NKIBench contains 14 representative NKI kernels from LLM workloads, combining compiler-generated baselines with four manually constructed or adapted kernels where compiler implementations were unavailable.The suite spans inference and training, from single operators to multi-operator chains and larger building blocks; one kernel comes from a non-transformer LLM.Source: §3.1 · p. 5

3.2 Profiling Service

• Profiling checks correctness across random-seed inputs and measures execution time excluding compilation, averaging warm-up runs and repeating rounds to mitigate fluctuations.The service selects a round with the smallest relative difference, or the first round within a predefined threshold.Source: §3.2 · p. 5

3.3 Peak Performance Calculation

• NKIBench estimates each task’s best achievable Trainium performance using roofline analysis, enabling assessment against hardware peak rather than latency-relative baselines alone.The estimate accounts for minimum tensor traffic and separates matmul from other FLOPs, with vector and scalar peak throughput considered together.Source: §3.3 · p. 6

4.1 Overall Performance

• 61% versus 49% of peak throughput on Trainium 1 and 59% versus 45% on Trainium 2; AccelOpt matches Claude Sonnet 4 thinking mode while costing 26× less.AccelOpt performs comparably with Claude Sonnet 4 across most kernels.Source: §4.1 · p. 7, Figure 7 · p. 7

4.2 Optimization Case Study

• AccelOpt discovers peephole rewrites and multi-step loop transformations, including removing BatchMatmul+Softmax spilling while avoiding extra recomputation and improving vector-engine utilization.In a Conv2D kernel outside NKIBench, it raised throughput from 9.54% to 48.8% of peak; 33.6% of 131 student teams solved an extra-credit challenge based on its optimization.Source: §4.2 · pp. 7–8
• 54.6% versus 52.7% of peak on Mamba and 29.6% versus 21.1% on RoPE show AccelOpt exceeding the cited human-optimized references by 1.04× and 1.4×, respectively.The Mamba kernel used a different loop order than the best human version; the authors attribute AccelOpt's scalability to exploring many kernels in parallel.Source: §4.2 · p. 8

4.3 Optimization Limitation Analysis

• About 82% of peak throughput was enough to explain one speedup plateau despite continued exploration, while another kernel was difficult to optimize because its small, on-chip workload and K=64 limited utilization.In a separate case, agents continued making meaningful rewrites despite stalled latency; performance metrics varied, indicating exploration remained active.Source: §4.3 · p. 9

4.4 Ablation Study of AccelOpt Components

• Beam search yields cumulative performance gains over repeated sampling, and optimization memory achieves similar speedup in 13 rather than 16 iterations, saving 16–17% cost while improving candidate quality.Memory raises cumulative Fast@p and helps produce stronger candidate pools and higher best speedups with fewer iterations.Source: §4.4 · pp. 9–10, Figure 13 · p. 10
• 1.137× versus 1.235× speedup shows the Reflexion-style baseline underperforms AccelOpt, while consuming more tokens by reflecting on every generated kernel.The paper notes that a fair comparison with non-LLM search-based optimizers is difficult because they typically require manual insight, kernel rewrites, and explicit search-space and pruning choices.Source: §4.4 · pp. 10–11

4.5 Cost Analysis

• Increasing memory capacity is more cost-efficient than increasing update eagerness, and its benefit depends on the model: Qwen3-Coder-30B gains 4.6% speedup for $12.33, versus 0.6% for $13.81 with gpt-oss-120b.The model ensemble achieves the best results at both ExpN=8 and ExpN=16; changing planners produced no substantial speedup differences, suggesting executor capability is a more promising target.Source: §4.5 · p. 11, Table 1 · p. 11

4.6 Generalization to Other Platforms

• 1.27× average speedup over best Triton baselines, with 3.19× peak speedup on a GQA decoding kernel, extends AccelOpt's evaluation to 24 H100 Triton kernels.The system is described as platform-agnostic when a profiling service and platform-specific base prompts are available; the paper focuses on single-core kernels without cross-chip communication.Source: §4.6 · p. 12

Memory for LLM Agents

• LessonL (Liu et al., 2025) uses baseline-anchored experience items, potentially causing homogeneity mitigated by ensembling; AccelOpt evolves memory with candidate kernels, potentially increasing diversity.Source: §5 · p. 12

LLM Agents for AI Accelerator Kernel Optimization

• Zhang et al. (2025b) translates operators without optimizing, Autocomp (Hong et al., 2025) uses hand-crafted planners for unfused kernels, and AlphaEvolve (Novikov et al., 2025) targets TPU kernels but is not public.GEPA (Agrawal et al., 2025) automatically discovers and injects NPU best practices into prompts, while AccelOpt’s optimization memories may offer more detailed task-specific insights.Source: §5 · p. 12

Benchmarks for Kernel Optimization

• Prior GPU and accelerator benchmarks (Ouyang et al., 2025a; Wen et al., 2025) measure relative speedup; NKIBench also reports peak-throughput ratio as an absolute hardware-specific standard.Recent work also targets interface usability (Xing et al., 2026) and evaluation robustness (Lange et al., 2025).Source: §5 · p. 12

6 CONCLUSION

• AccelOpt is presented as the first self-improving LLM agentic system for emerging-accelerator kernel optimization that combines search with memory accumulation.The paper identifies AccelOpt and NKIBench as a foundation for automated kernel optimization on emerging AI accelerators.Source: §6 · p. 12
• Combining inference-time scaling with optimization memory enables autonomous optimization of real-world Trainium kernels in NKIBench without expert optimization knowledge.Source: §6 · p. 12
• Ablation studies confirm that beam search and optimization memory help obtain high-performing kernels with improved cost efficiency.Source: §6 · p. 12
• Open-source models achieve higher cost efficiency than leading proprietary coding models on this task.Source: §6 · p. 12

A.1 Extra Information

• AccelOpt solutions took 1.59–31.29s (mean 7.38s), versus 1.77–30.61s (mean 8.13s) for baseline NKIBench kernels.The paper notes that compilation time does not affect kernel quality because kernels are usually reused after one-time compilation.Source: §A.1 · p. 16
• In one run, the best speedup plateaued or degraded from iterations 3–8, then improved from iteration 9 through iteration 14 despite continued exploration.Key utilization metrics changed during the plateau.Source: Figure 16 · p. 16
• The appendix includes an example NKI program snippet adopted from an official NKI example.Source: Figure 17 · p. 16

A.2 Experiment Details

• gpt-oss-20b and gpt-oss-120b use medium reasoning, while Claude Sonnet 4 uses thinking mode with max sequence length 20k, max output length 10k, and temperature 1.0.Open-source models use the default sampling setting in vLLM, and logfire records LLM query information.Source: §A.2 · p. 16
• Table 4 presents peak achievable hardware statistics.Source: Table 4 · p. 16
• CPU full-precision execution serves as ground truth despite being slower; experiments allow at most 10 rounds, use 1% and 4% performance-difference thresholds on Trainium 1 and 2, respectively, and run 2 warmups plus 10 repetitions.The CPU reference is considered higher fidelity because there is no IEEE standard for special functions such as exponential.Source: §A.2 · p. 16
• Section 4 uses T=16; before Section 4.5, Qwen3-Coder-480B is the executor and gpt-oss-120b handles other agents and Trainium 1 unless noted.Sections 4.1 and 4.3 use B=6, N=12, K=2, T=16, Topk=8, ExpN=16 on Trainium 1 and 2; selected Section 4.2 optimizations include peephole settings B=6, N=12, K=2, Topk=8, ExpN=8 on Trainium 2 and loop settings B=6, N=12, K=4.Source: §A.2 · p. 16
• LLMs, especially gpt-oss, exploited the correctness checker by omitting necessary safe softmax computation, motivating equivalence checks beyond random-input testing.One example computed only the row-wise maximum of the first tile in each row chunk to produce a fake speedup.Source: §A.2 · p. 16

A.3 Prompts

• The planner prompt combines NKI basics, profiling terminology, and a user template.These components are shown in Figures 19–21.Source: §A.3 · p. 16, Figure 19 · p. 18, Figure 20 · p. 18
• The executor prompt combines the planner’s NKI base knowledge with a concentrated programming guide tuned to agents’ common errors and a user template.The guide is adopted from the public NKI programming document and shown in Figures 22–24.Source: §A.3 · p. 16, Figure 22 · p. 20, Figure 24 · p. 22
• Because planner output follows a stable pattern, AccelOpt passes it directly into the executor prompt without extra formatting and randomizes profiling-item order across samples.The item-order changes are intended to increase planner randomness.Source: §A.3 · p. 16
• NKI requires parallelizable affine_range loop iterations to write to different memory locations.The prompt excerpt gives an ndarray allocation and affine_range loop as code context.Source: §A.3 · p. 27