💬 NLP

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

AccelOpt is a self-improving LLM agentic system that autonomously optimizes AI accelerator kernels by leveraging an experience-based memory, achieving significant throughput gains on AWS Trainium hardware while matching the performance of expensive proprietary models at a fraction of the cost.

Original authors: Genghan Zhang, Shaowei Zhu, Anjiang Wei, Zhenyu Song, Allen Nie, Zhen Jia, Nandita Vijaykumar, Yida Wang, Kunle Olukotun

Published 2026-04-17

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Genghan Zhang, Shaowei Zhu, Anjiang Wei, Zhenyu Song, Allen Nie, Zhen Jia, Nandita Vijaykumar, Yida Wang, Kunle Olukotun

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a brand new, incredibly powerful supercomputer (like Amazon's Trainium). It's built to solve massive AI problems, like training the next generation of chatbots. But there's a catch: this computer speaks a very specific, difficult language called "Neuron Kernel Interface" (NKI).

To get this computer to run fast, you need to write tiny, low-level instructions called kernels. Think of these kernels as the "recipes" for how the computer cooks data. If the recipe is bad, the computer burns the food (wastes time and money). If the recipe is perfect, it cooks a feast in seconds.

The Problem:
Writing these perfect recipes is incredibly hard. It usually requires a team of elite experts who have spent years studying the computer's blueprints. But for these new supercomputers, the blueprints are new, the experts are scarce, and the "recipes" we have right now are often slow and inefficient. It's like trying to drive a Formula 1 car with a manual written for a bicycle.

The Solution: AccelOpt
The authors of this paper built a system called AccelOpt. Instead of hiring a human expert to write every recipe, they built a self-improving AI agent that learns how to write these recipes on its own.

Here is how AccelOpt works, using a simple analogy:

1. The Team of AI Agents

AccelOpt isn't just one robot; it's a small team of three AI specialists working together:

The Planner: This agent looks at a slow recipe and says, "Hey, I see we are walking back and forth to the fridge too many times. Let's move the fridge closer." It comes up with a strategy to fix the code.
The Executor: This agent actually tries to rewrite the code based on the Planner's strategy. It's the hands-on worker.
The Summarizer: This agent is the teacher. After the Executor tries a new version, the Summarizer looks at the "Before" and "After." If the new version is faster, the Summarizer writes down a note: "Moving the fridge closer saved 5 minutes. Remember this trick!"

2. The "Optimization Memory" (The Brain)

This is the secret sauce. Most AI systems try to solve a problem from scratch every time. AccelOpt is different. It keeps a notebook (Optimization Memory).

Every time the team finds a way to make a kernel faster, they write it in the notebook.
Every time they try something that makes it slower, they write that down too (so they don't repeat the mistake).
In the next round of work, the team reads the notebook first. They don't start from zero; they start with all the lessons they learned yesterday.

3. The "Beam Search" (Exploring Many Paths)

Imagine you are trying to find the fastest route through a maze.

Old way: You pick one path, walk it, and if it's a dead end, you go back and try another.
AccelOpt's way (Beam Search): You send out 6 different explorers at the same time. Each one tries a slightly different turn. You keep the 6 best paths and send them out again for the next round, discarding the dead ends. This way, you explore the maze much faster and find the exit (the fastest code) without getting stuck.

The Results: Why This Matters

The researchers tested AccelOpt on real-world AI tasks using Amazon's Trainium chips.

Before AccelOpt: The AI was only using about 45-49% of the computer's potential speed. It was like driving a Ferrari in a school zone.
After AccelOpt: The AI boosted performance to 59-61% of the computer's peak speed.
The Cost: Usually, to get this kind of improvement, you'd need to hire a human expert or use a very expensive, proprietary AI model (like the "Claude Sonnet 4" mentioned in the paper). AccelOpt achieved the same results using open-source models (free or cheap AI) and cost 26 times less than the expensive alternatives.

The Big Picture

Think of AccelOpt as a self-teaching mechanic.
In the past, if you bought a new, weirdly shaped engine, you had to wait for a human master mechanic to figure out how to tune it. With AccelOpt, you can just hand the engine to the robot. The robot tries a million different tweaks, remembers what worked, forgets what didn't, and eventually tunes the engine to run at peak performance—all without a human ever touching a wrench.

This is a huge step forward because as AI gets bigger and we build more custom chips, we won't have enough human experts to tune every single one. We need systems like AccelOpt that can learn and improve on their own.

1. Problem Statement

The rapid proliferation of AI accelerators (e.g., AWS Trainium, Google TPU, custom NPUs) has created a critical bottleneck: kernel optimization.

The Challenge: High-performance kernels require navigating a complex design space involving memory hierarchies, parallelism, and architecture-specific constraints. For emerging accelerators like AWS Trainium (programmed via Neuron Kernel Interface, or NKI), there is a severe lack of expert knowledge, established heuristics, and optimization recipes compared to mature platforms like GPUs.
The Cost: Suboptimal kernels lead to significant waste of compute resources and financial costs. Manual optimization is slow, expensive, and does not scale to the hundreds or thousands of kernels required for diverse LLM workloads.
The Gap: Existing LLM-based approaches often rely on human-engineered prompts or specific optimization examples. There is a need for a system that can autonomously explore the optimization space and self-improve without relying on pre-existing expert knowledge.

2. Methodology: AccelOpt

AccelOpt is a self-improving LLM agentic system designed to autonomously optimize kernels for emerging AI accelerators. It combines inference-time scaling (Beam Search) with a dynamic optimization memory.

Core Architecture

The system operates in iterative cycles, utilizing a three-agent workflow:

Planner: Analyzes the baseline kernel and profiling data (e.g., memory bandwidth, engine utilization) to generate optimization plans.
Executor: Implements the plans, generating new candidate kernels.
Summarizer: Distills successful "slow-to-fast" kernel pairs into generalizable optimization insights (pseudocode and strategy descriptions).

Key Mechanisms

Beam Search: Instead of simple repeated sampling, AccelOpt uses beam search to maintain a frontier of the top $B$ $B$ candidate kernels at each iteration.
- For each candidate, the Planner generates $N$ distinct optimization plans.
- The Executor attempts $K$ implementations per plan.
- The system selects the top $B$ kernels from the expanded pool ( $B + B \times N \times K$ ) to proceed to the next iteration. This ensures cumulative progress rather than random guessing.
Optimization Memory: A curated queue of experiences that stores:
- Positive Rewrites: Pairs of slow baseline kernels and faster optimized versions.
- Negative Rewrites: Pairs of baseline kernels and slower generated versions (to learn what not to do).
- Insights: Summarized general strategies (e.g., "Loop Invariant Code Motion," "Tile Size Adjustment") extracted by the Summarizer.
- This memory is injected into the Planner's context for future iterations, allowing the system to learn from its own history and avoid past mistakes.

Benchmark: NKIBench

To evaluate AccelOpt, the authors constructed NKIBench, the first benchmark suite for NKI kernel optimization.

Source: 14 kernels derived from real-world LLM workloads (e.g., Mamba, RoPE, BatchMatmul, LoRA) from models like DeepSeek, Qwen, and Falcon.
Metric: Unlike previous benchmarks that measure relative speedup against a baseline, NKIBench measures performance against the theoretical peak hardware throughput (calculated via roofline model analysis). This provides an absolute standard for optimization quality.

3. Key Contributions

First Self-Improving System for Emerging Accelerators: AccelOpt is the first system to combine search (beam search) with memory accumulation for kernel optimization on new hardware (Trainium) without requiring expert-provided optimization recipes.
NKIBench Benchmark: A new, rigorous benchmark suite featuring real-world LLM kernels and theoretical peak performance baselines, enabling absolute performance evaluation.
Cost-Effective Performance: Demonstrated that open-source LLMs (specifically Qwen3-Coder-480B and gpt-oss-120b) can achieve performance comparable to leading proprietary models (Claude Sonnet 4) at a fraction of the cost.
Ablation Insights: Validated that beam search is superior to repeated sampling and that optimization memory significantly improves cost-efficiency by reducing the number of iterations needed to find high-quality kernels.

4. Experimental Results

The evaluation was conducted on AWS Trainium 1 and Trainium 2 hardware.

Performance Gains:
- Trainium 1: Average peak throughput increased from 49% to 61%.
- Trainium 2: Average peak throughput increased from 45% to 59%.
- AccelOpt matched the performance of Claude Sonnet 4 (in thinking mode) but was 26× cheaper when using open-source models.
Optimization Capabilities:
- Local Optimizations: Successfully performed peephole optimizations (algebraic simplification, intrinsic fusion like rsqrt).
- Global Optimizations: Discovered non-trivial global changes, such as loop transformations, memory layout changes, and eliminating redundant computations (e.g., moving matrix transposes out of reduction loops).
- Human Comparison: On specific kernels (Mamba, RoPE), AccelOpt autonomously achieved results exceeding or matching human-optimized reference versions (e.g., 54.6% vs. 52.7% peak for Mamba).
Cost Analysis:
- Using open-source models with optimization memory reduced costs by 3.3× compared to using Claude Sonnet 4 with repeated sampling.
- Increasing memory capacity ($ExpN$) was found to be more cost-effective than increasing the eagerness of memory updates ($TopK$).

5. Significance and Impact

Democratization of Optimization: AccelOpt lowers the barrier to entry for optimizing kernels on new, specialized AI hardware. It removes the dependency on scarce human experts who possess deep, platform-specific knowledge.
Scalability: The system can explore thousands of optimization paths in parallel, a scale impossible for human engineers, making it suitable for the diverse and rapidly evolving landscape of AI accelerators.
Educational Value: The system was successfully used in a graduate-level course (Stanford CS149) to teach parallel computing principles, demonstrating its utility as a tool for human learning.
Future Directions: The paper suggests that while the current focus is on single-core kernels, the methodology is platform-agnostic and can be extended to communication primitives and other hardware architectures (e.g., GPUs, TPUs) where LLMs have even more training data.

In conclusion, AccelOpt represents a paradigm shift in compiler optimization, moving from static, rule-based or manually guided tuning to a dynamic, self-improving agent system capable of autonomously mastering the complexities of emerging AI hardware.