Auto GPU Kernel 🏆

Autonomous GPU-kernel discovery & optimizer. Supports flashinfer format and a custom format that builds its own harness around any natural language task against a given git repository.

Technical Report

Ranked #1 on MLSys 2026 - FlashInfer AI Kernel Generation Contest for the DeepSeek Sparse Attention (DSA) track with an average speedup of 34.93x. Submissions + optimized kernels can be found at archive.

• kopt tool runs the optimization agent.
• kbench tool owns validation, benchmarks, and agent history.

Why use auto-gpu-kernel over other harnesses or /goal ?

auto-gpu-kernel is designed by discovering kernel-engineering best practices for agents over months of effort and research during FlashInfer kernel generation contest. It addresses common pitfalls of agents such as believing noisy results to be improvements or missing various optimizations due to measurement noise. Surely the research and optimization loop can be more fancy i.e. by adding many sub-agents, but my experience has shown those things matter less, as long as there is a robust validation pipeline for agent's solution and the speed-up. Other harnesses usually fail on setting a good verification pipeline, the agents either hack it over-time, or they get stuck at local-minimums.

My design motto is: "Don't include any long instructions generated by AI. Those can be generated by AI itself over-time with enough attempts. Only give the key principles the agent needs to follow, so the agent knows what really matters."

Install

You need Python 3.11+, uv, and OMP.

uv venv --python 3.12
source .venv/bin/activate

# Arbitrary local repositories
uv pip install -e ".[agent]"

# FlashInfer: choose local on Linux with a GPU, or Modal from any machine
# uv pip install -e ".[agent,local]"
# uv pip install -e ".[agent,modal]"

Make sure omp + your model provider is ready.

omp

Run a FlashInfer kernel

Download a trace set:

git lfs install
git clone https://huggingface.co/datasets/flashinfer-ai/mlsys26-contest

If using a cloud provider, upload it:

modal setup
modal volume create flashinfer-trace
modal volume put flashinfer-trace ./mlsys26-contest/

Create a project and run three iterations:

kopt init ./mlsys26-contest/definitions/dsa_paged/dsa_topk_indexer_fp8_h64_d128_topk2048_ps64.json \
  --backend modal --gpu B200

cd work/dsa_topk_indexer_fp8_h64_d128_topk2048_ps64
kbench bench --quick
kopt run . -n 3 --model anthropic/claude-opus-5 --thinking low

Useful benchmark commands:

kbench bench --quick       # smallest and largest workload
kbench bench --stride 2    # half of the workloads
kbench bench               # full run
kbench ab --a experiments/exp_3/solution_fused.py

Development

examples/mlx-attention/ is a small naive MLX attention with tests and a benchmark. configs/mlx_attention.toml targets it, so you can exercise the whole task-mode loop on Apple Silicon with no remote repo or GPU cloud:

uv pip install mlx pytest
kopt init-task configs/mlx_attention.toml
kopt run work/mlx-attention -n 5 --model anthropic/claude-opus-5 --thinking low

The first run spends one setup turn generating harness/, measures the pristine baseline (about 2.7 ms geomean on an M4 Pro; the fused mx.fast op does it in ~0.9 ms), then starts optimizing attention.py.

Run against any Git repository

Write a short task.toml. Describe the job in plain language; the setup agent will inspect the repository and build the benchmark harness.

[task]
name = "my-project"
workdir = "repo"
objective = """
Speed up inference without changing the public API or model outputs.
The optimizer may change code below src/runtime/.
"""
measure = """
Measure end-to-end latency for the representative example in examples/serve.py.
Lower latency is better. Include warmup and synchronize the GPU before timing.
"""
validate = """
Run the existing correctness tests and compare the example's output with the untouched
repository. Outputs must match exactly.
"""
hints = "Start with allocations and repeated kernel launches in the decode loop."

[task.repo]
url = "git@github.com:my-org/my-project.git"
base = "main"
branch = "me/auto-optimize"
# or, for a plain local directory (relative to this file), snapshotted as a fresh repo:
# path = "../my-project"

[task.hardware]
gpus = 1
gpu = "H100"

The setup agent implements the generated-task adapter by writing harness/validate.py and harness/benchmark.py before the first run. Generated-task adapters currently run on the local machine, so the target repository's dependencies and any required GPU must be available there.

Then scaffold and run:

kopt init-task task.toml            # creates work/<task.name>
kopt run work/my-project -n 20 --model anthropic/claude-opus-5 --thinking low

On the first run, kopt uses one setup turn to inspect the untouched clone and generate the validation and benchmark adapters. Kbench then runs pristine quick and full baselines. The 20 requested optimization iterations begin after setup.

To inspect or rerun what the setup agent made:

cd work/my-project
cat harness/README.md
kbench bench --quick            # quick validation + quick measurement
kbench bench                    # full validation + metric of record
kbench ab --a <git-ref>         # same current harness for A and B

Watch a run

kopt watch .

Open http://127.0.0.1:8765.

config.toml          human task brief
harness/validate.py  generated quick/full correctness adapter
harness/benchmark.py generated quick/full measurement adapter
harness/prepared.json pristine baseline metadata
.omp/                project-local agent instructions
.kopt/runs/          agent logs
.kopt/bench.jsonl    benchmark history
experiments/         experiment notes and snapshots

Architecture

FlashInfer and arbitrary repositories use the same BenchmarkAdapter:

• FlashInferAdapter packages kernel sources and sends them through a local or Modal execution backend.
• GeneratedTaskAdapter runs the repository-specific validation and benchmark scripts created by the setup agent.
• Kbench supplies quick/full execution, normalized measurements, history, and paired A/B for both adapters.

See kbench/README.md and kopt/README.md for the small class diagrams.