TongmingLAIC/AKO4ALL

What it solves

AKO4ALL automates the tedious process of GPU kernel optimization. Instead of a human expert manually profiling and rewriting code to squeeze out performance, it uses an AI agent to iteratively improve a kernel's speed while ensuring it remains correct.

How it works

It is implemented as a "skill" (a protocol document) for coding agents like Claude Code. When invoked in a directory containing a GPU kernel, the agent follows a closed-loop cycle: it creates a workspace, verifies a baseline, profiles the code using NVIDIA Nsight Compute (ncu), rewrites the kernel, benchmarks the result, and logs the progress. If performance plateaus, the agent may perform web searches for new optimization ideas or re-profile the code to find new bottlenecks. It supports various languages including Triton, CUDA, C++, and TileLang.

Who it’s for

It is designed for developers and researchers working with high-performance GPU kernels who want to automate the optimization process without building a full-scale infrastructure for every single kernel.

Highlights

• Expert-Level Performance: Demonstrated speedups over FlashInfer's expert kernels on operators like GQA and MLA decode/prefill on NVIDIA B200.
• Language Agnostic: Can optimize kernels written in Triton, CUDA, C++, Python, and more, and can even switch languages to find better performance.
• Integrated Evaluation: Includes a built-in evaluator (KernelBench) and anti-cheat mechanisms to prevent the agent from gaming metrics.
• Drop-in Integration: Functions as a single skill for Claude Code, requiring minimal setup to start an optimization run.