AccelOpt

KO | EN

gitlite — search

#javascript #php #ruby #hacktoberfest #nodejs #python #blog #gulp #java #css #r #angularjs

★ 57 Open GitHub ↗

[MLSys 2026] AccelOpt: Self-improving Agents for AI Accelerator Kernel Optimization

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# AccelOpt: Self-improving Agents for AI Accelerator Kernel Optimization

[![Paper](https://img.shields.io/badge/arXiv-2511.15915-b31b1b.svg)](https://arxiv.org/pdf/2511.15915)
[![Blog Post](https://img.shields.io/badge/Blog-accelopt.html-blue)](https://ppl.stanford.edu/accelopt.html)
[![HuggingFace](https://img.shields.io/badge/🤗-NKIBench-yellow)](https://huggingface.co/datasets/Genghan/NKIBench)

AccelOpt is 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.

🚧 This repository is still under construction.

## Setup
### AWS Trainium 
#### Machines
EC2 Instance: trn1.32xlarge

AMI: Deep Learning AMI Neuron (Ubuntu 22.04)
#### Install
```
source /opt/aws_neuronx_venv_pytorch_2_7/bin/activate # Check the PyTorch version of your AMI
pip install logfire
pip install openai-agents
git clone git@github.com:zhang677/AccelOpt.git
cd AccelOpt
python setup.py install
```
#### Run
`experiments/full_complete_local` shows how to run AccelOpt on NKIBench with a local served gpt-oss-120b.

### NVIDIA GPU 
#### Machines
EC2 Instance: p5.48xlarge

AMI: Deep Learning OSS Nvidia Driver AMI GPU PyTorch 2.8 (Ubuntu 24.04)
#### Install
```
# Create a PyTorch environment first
pip install logfire
pip install openai-agents
git clone https://github.com/zhang677/flashinfer-bench.git # DPS is not consistent with flashinfer-trace now
cd flashinfer-bench
pip install -v -e .
cd ..
git clone git@github.com:zhang677/AccelOpt.git
cd AccelOpt
python setup.py install
```
#### Run
`experiments/flb_full_complete_local` shows how to run AccelOpt on [FlashInfer-Bench](https://github.com/flashinfer-ai/flashinfer-bench) with a local served gpt-oss-120b.

## Kernel Database
Developed concurrently, NKIBench and FlashInfer-Bench share similar kernel schema, demonstrating the benefits of structual storage. `./templates/schema_v2.txt` shows the SQL schema we use to manage generated kernels for Flashinfer-Bench. `./accelopt/db_utils.py` captures the common usage patterns of this kernel database with [supabase](https://supabase.com/).

# NKIBench: Agent Environment for AWS Trainium accelerators

NKIBench is the first agent environemt for AWS Trainium accelerators that contain a profiling service and a set of NKI kernels with varying complexity extracted from real-world LLM workloads. Huggingface dataset can be found here [Genghan/NKIBench](https://huggingface.co/datasets/Genghan/NKIBench).

## Dataset
All the kernels are under `/NKIBench`. Kernels are
grouped by operator name and configuration in a structured
storage format as shown in the `/NKIBench/summary.json`

## Profiling
AccelOpt provides a `NKIKernel` class that is pluggable to any AI optimizers. `/tests` shows how to use the profiling API for single and a group of NKI kernels. 
```python
from accelopt.kernel_wrapper import NKIKernel
nki_kernel = NKIKernel(program_path, base_numpy_path)
result = nki_kernel.profile(save_fields)
```

NKIBench estimates the best achievable performance offered by the Trainium hardware, which offers additional insights on how effective AccelOpt has been in exploring the entire optimization landscape. The best achievable performance is calculated in `experiments/full_complete_local/calculate_percentage_of_peak.py`

# Results 
## NKIBench
<div align="center">
    <img src="img/ratio_bars_trn1.png" alt="TRN1-result" style="width:85%;">
</div>
Per-task kernel improvement achieved using Claude Sonnet 4 and AccelOpt (gpt-oss-120b + Qwen3-Coder-480B) on Trainium 1 of NKIBench.

## FlashInfer-Bench
<div align="center">
    <img src="img/compare_bars_with_flashinfer.png" alt="TRN1-result" style="width:75%;">
</div>
Triton kernel improvement achieved using AccelOpt with Fixer on H100 SXM5 over FlashInfer across 8 attention workload of FlashInfer-Bench. Fixers and executors use the same model and all the other agents use gpt-oss-120b. These experiements used `./templates/fib/complete_local_prompts/run_single_loop.sh` with `PROMPT_BASE_DIR=./prompts/fib`.

# Citation

This work was accepted by MLSys 2026. This is a follow-up to *Adaptive Self-improvement LLM Agentic System for ML Library Development* (ICML 2025) [[paper](https://proceedings.mlr.press/v267/zhang25at.html)] [[blog](https://zhang677.github.io/blog_md/aspl.html)] [[code](https://github.com/zhang677/PCL-lite)]. If you find this project useful, please cite:
```
@inproceedings{zhang2025adaptive,
  title={Adaptive Self-improvement LLM Agentic System for ML Library Development},
  author={Zhang, Genghan and Liang, Weixin and Hsu, Olivia and Olukotun, Kunle},
  booktitle={Proceedings of the 42nd International Conference on Machine Learning},
  year={2025}
}

@inproceedings{zhang2026accelopt,
 author = {Zhang, Genghan and Zhu, Shaowei and Wei, Anjiang and Song, Zhenyu and Nie, Allen and Jia, Zhen and Vijaykumar, Nandita and Wang, Yida and Olukotun, Kunle},
 booktitle = {Proceedings of Machine Learning and Systems},
 editor = {A. Chowdhery and Z. Jia},
 pages = {541--568},
 publisher = {MLSys},
 title = {AccelOpt: A Self-Improving LLM Agentic System for AI Accelerator Kernel Optimization},
 url = {https://proceedings.mlsys.org/paper_files/paper/2026/file/0f8426558905746fc38da5e335700aec-Paper-Conference.pdf},
 volume = {8},
 year = {2026}
}

```

AccelOpt

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README