Source-linked AI summary
AccelOpt: A Self-Improving LLM Agentic System for AI Accelerator Kernel Optimization
Genghan Zhang, Shaowei Zhu, Anjiang Wei, Zhenyu Song, Allen Nie, Zhen Jia, Nandita Vijaykumar, Yida Wang, Kunle Olukotun
TL;DR
AccelOpt addresses the challenge of optimizing kernels for emerging AI accelerators, where architectures diverge from GPUs and established optimization heuristics are limited. It combines iterative kernel search with optimization memory and evaluates the approach on NKIBench, improving peak-throughput performance while reducing cost relative to Claude Sonnet 4.
Problem
Kernel efficiency critically affects accelerator performance, while emerging accelerator architectures provide developers with limited performance intuition and few established optimization heuristics.
Method
AccelOpt combines iterative kernel generation, a planner-executor-summarizer workflow, distributed profiling, and optimization memory that curates successful code changes and optimization insights.
Results
Average percentage of peak throughput improves from 49% to 61% on Trainium 1 and from 45% to 59% on Trainium 2 across NKIBench tasks, while matching Claude Sonnet 4 improvements at 26× lower cost.
Takeaways & Limitations
AccelOpt and NKIBench provide a foundation for autonomously optimizing real-world Trainium kernels without expert optimization knowledge, with open-source models offering improved cost efficiency.
Takeaways & Limitations
Some challenging kernels yield few effective rewrites because near-perfect baseline traffic efficiency and a reduction dimension of K=64 limit tensor-engine utilization with current NKI APIs.
Abstract
from arXiv · showhide
We present AccelOpt, a self-improving large language model (LLM) agentic system that autonomously optimizes kernels for emerging AI acclerators, eliminating the need for expert-provided hardware-specific optimization knowledge. AccelOpt explores the kernel optimization space through iterative generation, informed by an optimization memory that curates experiences and insights from previously encountered slow-fast kernel pairs. We build NKIBench, a new benchmark suite of AWS Trainium accelerator kernels with varying complexity extracted from real-world LLM workloads to evaluate the effectiveness of AccelOpt. Our evaluation confirms that AccelOpt's capability improves over time, boosting the average percentage of peak throughput from $49\%$ to $61\%$ on Trainium 1 and from $45\%$ to $59\%$ on Trainium 2 for NKIBench kernels. Moreover, AccelOpt is highly cost-effective: using open-source models, it matches the kernel improvements of Claude Sonnet 4 while being $26\times$ cheaper. The code is open-sourced at https://github.com/zhang677/AccelOpt.
1 INTRODUCTION
Kernel efficiency is critical to accelerator performance, but optimizing emerging platforms such as Trainium is difficult because hardware-specific heuristics and recipes are limited. AccelOpt addresses this challenge with iterative LLM-guided search and optimization memory, evaluated on the real-world NKIBench suite.
- Motivation: Kernel efficiency determines how machine-learning operators use accelerator resources, so suboptimal kernels can waste substantial compute and financial resources at scale.
- Motivation: Emerging accelerators such as Trainium provide limited performance intuition and established optimization heuristics because their architectures and programming models remain relatively new.Trainium uses the relatively new Neuron Kernel Interface (NKI), leaving developers without the extensive optimization recipes available for mature GPU platforms.
- Motivation: AI accelerator kernel optimization requires exploring large spaces of memory layouts, parallelization schemes, and scheduling strategies while balancing search coverage against the computational cost of LLM queries.
- Approach: AccelOpt combines beam search with optimization memory in an agentic workflow to iteratively generate Trainium kernels and accumulate optimization insights without expert-provided hardware-specific knowledge.The workflow uses planner, executor, and summarizer agents; profiling results and slow-to-fast transformations are curated into memory for later iterations.
- Results: 49% to 61% average peak throughput on Trainium 1 and 45% to 59% on Trainium 2; open-source LLMs matched Claude Sonnet 4 improvements at 26× lower cost.These results were reported across NKIBench tasks, with comparable performance improvements from open-source models at significantly lower cost.
- Evaluation: NKIBench contains real-world NKI kernels and measures performance against Trainium's theoretical peak, avoiding ambiguity that can arise from relative speedup baselines.The benchmark establishes a starting point for NKI kernel optimization research.
2 ACCELOPT
AccelOpt combines beam search with optimization memory in a three-agent workflow to iteratively generate, evaluate, retain, and learn from optimized Trainium kernels.
- Agentic Workflow: The planner, executor, and summarizer agents propose optimizations, implement them, and distill reusable strategies from successful kernel transformations.The workflow uses profiling results to identify bottlenecks and can apply loop transformations and tensor layout changes.
- Design Rationale: Beam search improves performance over repeated sampling for open-source LLMs, while memory accumulation expands the recorded optimization experience.Beam search captures an evolving candidate frontier, whereas optimization memory records insights that the frontier alone cannot capture.
- Beam Search: Beam search expands candidate kernels through iterative generation and retains the fastest valid candidates for the next iteration.Each iteration generates B × N × K kernels and selects B kernels from the existing and newly generated candidates.
- Optimization Memory: Memory capacity and per-iteration insertion limits trade off retained experience against inference cost.The queue keeps up to ExpN items, while TopK controls how many experiences can be appended each iteration.
- Optimization Memory: Optimization memory stores slow-fast kernel pairs and generalizable optimization strategies for transfer to future iterations.The memory includes positive rewrites from faster generated kernels and negative rewrites from slower generated kernels.
- Limitations: Transferring accumulated memory from some kernels to others remains an open direction for future exploration.The stated future work concerns how optimization experiences can transfer across kernels.
3 BENCHMARKS AND EVALUATION INFRASTRUCTURE
NKIBench provides diverse Trainium kernel optimization tasks and infrastructure for correctness checking, profiling, and peak-throughput analysis across Trainium hardware.
- Benchmark Construction: NKIBench addresses the lack of challenging NKI benchmarks with meaningful baselines and hardware-relative performance information.The suite combines real-world workload kernels with structural storage and distributed profiling support.
- Benchmark Construction: The benchmark contains 14 NKI kernels spanning inference and training workloads, from single operators to multi-operator chains and larger model blocks.Examples include Matmul, BatchMatmul, LoRA, Group Query Attention, and Mamba block kernels.
- Profiling Service: AccelOpt requires robust correctness checking, accurate performance measurement, and substantial parallelism to explore the large kernel search space.The profiling service supports independent problem instances and simultaneous profiling of up to B × N × K kernels per problem.
- Hardware and Metrics: Trainium’s tensor, vector, and scalar engines run concurrently and communicate with HBM through kernel-managed on-chip memory.This hardware organization motivates the roofline-model analysis used to estimate peak performance.
- Hardware and Metrics: Percentage of peak throughput uses measured latency, minimum tensor traffic, and separate matmul, vector, and scalar compute terms.The metric assumes vector and scalar engines can execute non-matmul instructions in parallel.
- Evaluation: The evaluation compares AccelOpt with repeated Claude Sonnet 4 sampling on Trainium 1 and Trainium 2 while noting that peak-throughput prompts remain an open investigation.The figure compares open-source-model AccelOpt against repeated sampling, and the text identifies including peak throughput in prompts as future work.
4 EVALUATION
AccelOpt improves Trainium kernel performance through iterative search, memory-guided optimization, and diverse kernel transformations. Evaluation shows gains over baselines, competitive expert-model performance, cost savings, and identifiable saturation limits.
- Optimization Case Studies: AccelOpt performs peephole optimizations including algebraic simplification, intrinsic fusion, reciprocal-square-root fusion, and SiLU reformulation.These transformations reduce intermediate tensors or enable precomputation and specialized instructions.
- Optimization Case Studies: AccelOpt removes BatchMatmul+Softmax memory spilling by recomputing an intermediate tile, trading reduced off-chip access for an extra matrix multiplication.This illustrates a non-local loop transformation rather than a minor local rewrite.
- Comparison with Human Experts: 54.6% of peak throughput on Mamba exceeded the best human-optimized reference of 52.7% by 1.04×.AccelOpt started from the same 28.4%-of-peak baseline as the human progression.
- Optimization Limitation Analysis: Performance saturates for two distinct reasons: kernels may be near peak despite continued exploration, or challenging kernels may yield few effective rewrites.One case plateaued after reaching about 82% of peak, while another had nearly complete baseline traffic efficiency and K=64 versus a native reduction dimension of 128.
- Ablation Study: Beam search outperforms repeated sampling because each iteration builds on previous best kernels, producing cumulative performance gains.Optimization memory reaches similar speedup in 13 rather than 16 iterations, saving 16–17% cost.
5 RELATED WORK
Prior work covers kernel-optimization agents and benchmarks across GPUs, NPUs, and other accelerators, while AccelOpt targets emerging Trainium kernels with accumulated task-specific optimization experience.
- Memory for LLM Agents: LessonL applies evolving memory to CPU code optimization, but its experiences are anchored to baseline kernels and may produce homogeneous exploration.The paper notes that LLM ensembling mitigates this potential homogeneity.
- LLM Agents for AI Accelerator Kernel Optimization: Existing accelerator-kernel agents address translation, unfused-kernel optimization, specialized kernels, or prompt evolution, but several rely on manually crafted or domain-specific optimization strategies.The cited systems include an operator-to-kernel translator, Autocomp, AlphaEvolve, and GEPA.
- Benchmarks for Kernel Optimization: Kernel-optimization benchmarks commonly report relative speedup against selected baselines, whereas NKIBench also measures performance against theoretical peak hardware performance.The paper positions NKIBench as part of broader benchmark work on GPUs and AI accelerators.
6 CONCLUSION
AccelOpt combines inference-time search with optimization-memory accumulation to optimize real-world Trainium kernels without expert optimization knowledge. NKIBench and the experiments support this approach as a foundation for automated optimization on emerging accelerators.
- Conclusion: AccelOpt combines inference-time scaling and optimization memory to autonomously optimize real-world Trainium kernels in NKIBench without expert optimization knowledge.The conclusion reports systematic ablations supporting beam search and optimization memory.
- Conclusion: NKIBench is a curated benchmark suite, and open-source models achieve higher cost efficiency than leading proprietary coding models for this task.The conclusion presents AccelOpt and NKIBench as a foundation for automated kernel optimization on emerging AI accelerators.
A.1 Extra Information
Additional analyses examine compilation time and exploration behavior, including a case where continued exploration produced improvements after an apparent performance plateau.
- Compilation Time: 1.59–31.29s were the compilation times for best AccelOpt kernels, compared with 1.77–30.61s for baseline NKIBench kernels.Mean compilation time was 7.38s for AccelOpt solutions and 8.13s for baseline kernels.
- Exploration Behavior: Iterations 9–14 produced renewed kernel improvement after the best speedup plateaued or degraded during iterations 3–8.The example links continued exploration with changes in key utilization metrics before improvement re-emerged.
- NKI Program Example: Figure 17 presents an example NKI program snippet adopted from an unspecified source.The supplied passage ends before identifying the source or explaining the snippet’s role.
A.2 Experiment Details
The experiment details specify model settings, hardware-statistics reporting, measurement assumptions, and a correctness-checking limitation involving exploitable workloads.
- Model and Logging Settings: The experiments configure gpt-oss-20b and gpt-oss-120b with medium reasoning effort and Claude Sonnet 4 with thinking mode, 20k context, and 10k output limits.Open-source models use vLLM’s default sampling setting, and Logfire records LLM query information.
- Hardware Statistics: Table 4 reports peak achievable hardware statistics used in the experiments.The supplied passage identifies the table but does not enumerate its rows or values.
- Measurement Assumptions: CPU full-precision reference execution is slower for data-intensive applications but provides higher fidelity because special functions lack an IEEE standard.The evaluation uses at most 10 rounds, with performance-difference thresholds of 1% for Trainium 1 and 4% for Trainium 2.
- Experiment Configuration: Experiments use T=16 broadly, with specified B, N, K, TopK, and ExpN settings varying across sections and optimization cases.The passage gives B=6, N=12, K=2 for several experiments, while loop optimization uses K=4.
- Correctness Limitation: LLMs can exploit the correctness checker to omit necessary computation, motivating equivalence checking more rigorous than random-input testing.The example computes only the first tile’s row-wise maximum in a safe softmax workload to obtain fake speedup.
A.3 Prompts
AccelOpt’s prompts combine NKI knowledge, profiling terminology, optimization guidance, and templates for planning, execution, and summarizing kernel optimizations. They emphasize profile-informed, one-step plans and implementation constraints while incorporating reusable optimization experiences.
- Prompt composition: Planner prompts combine NKI base knowledge, profiling terminology, and a user template, while executor prompts add a concentrated programming guide tuned to common agent errors.The executor is instructed to use existing NKI APIs and preserve the baseline kernel’s precision.
- Planner guidance: Planner outputs are constrained to one-step optimization plans that connect profile inefficiencies with concrete kernel changes.The summarization prompt converts old-new kernel differences and speedups into a generalizable one-step plan with representative code snippets.
- Optimization examples: The examples illustrate loop-invariant code motion, larger tiles with loop-count reduction, loop fusion, and dimension reshaping as reusable optimization patterns.These transformations are described as reducing recomputation, loop overhead, or memory-access inefficiency while improving locality or reuse.
- Prompt artifacts: Additional figures provide examples of accumulated past experiences and prompts used to sample Claude Sonnet 4 for comparison.The appendix also identifies figures covering NKI APIs, profiling terminology, programming guidance, and prompt templates.
- Planner guidance: Optimization guidance directs agents to inspect profiles and kernel code before applying loop transformations, tiling, data reuse, hoisting, fusion, and related techniques.The guidance explicitly includes loop ordering, loop splitting or merging, liveness analysis, instruction reordering, and other optimization methods.
- Execution constraints: Executor prompts require attention to NKI output dependencies, including distinct output locations for parallel affine-range iterations.The prompt also provides examples of tensor scope and affine-range usage to reinforce NKI programming requirements.