autokernel

AutoKernel

Autoresearch for GPU kernels. Give it any PyTorch model, go to sleep, wake up to optimized Triton or CUDA C++ kernels.

Inspired by @karpathy/autoresearch -- which demonstrated autonomous AI agents for LLM training research. AutoKernel applies the same philosophy to GPU kernel optimization: agent modifies one file, runs a fixed evaluation, keeps or reverts, repeats forever.

How It Works

Give AutoKernel any PyTorch model. It will:

1. Profile the model to find which GPU kernels are bottlenecks
2. Extract each bottleneck as a standalone Triton or CUDA C++ kernel
3. Optimize each kernel autonomously (edit, benchmark, keep/revert -- forever)
4. Verify end-to-end correctness and report the total speedup

The agent reads program.md -- the "research org code" -- which contains comprehensive instructions for autonomous operation. It edits kernel.py one kernel at a time, runs bench.py (fixed benchmark with 5-stage correctness checks + roofline analysis), and either keeps or reverts the change. The orchestrator decides when to move to the next kernel using Amdahl's law.

Each experiment takes ~90 seconds. That's ~40 experiments/hour, ~320 overnight, across all kernels.

Quick Start

Requirements: NVIDIA GPU (tested on H100/A100/RTX 4090), Python 3.10+, uv.

# Install uv (if you don't have it)
curl -LsSf https://astral.sh/uv/install.sh | sh

# Clone and setup
git clone https://github.com/RightNow-AI/autokernel.git
cd autokernel
uv sync

# One-time setup: test data + baselines
uv run prepare.py

# Profile a model (ships with GPT-2, LLaMA, BERT -- no transformers needed)
uv run profile.py --model models/llama_7b.py --class-name LlamaModel \
 --input-shape 1,512 --dtype float16

# Extract top bottleneck kernels
uv run extract.py --top 5

# Verify benchmark works
uv run bench.py

Running the Agent

Spin up Claude, Codex, or any coding agent in this directory:

Read program.md and let's kick off a new experiment. Start with setup.

The agent will:

1. Profile your model and present the optimization plan
2. Create a branch (e.g., autokernel/mar10-llama7b)
3. Optimize each bottleneck kernel in priority order
4. Verify end-to-end correctness and report total speedup

program.md is intentionally comprehensive so the agent can run 10+ hours without getting stuck. It includes a 6-tier optimization playbook, decision framework, crash handling, and Amdahl's law reasoning.

The Pipeline

                 profile.py              extract.py           bench.py (loop)         verify.py
Any PyTorch  ──>  Rank kernels  ──>  Generate baseline  ──>  Optimize each  ──>  End-to-end
   model          by GPU time       Triton/CUDA kernels     kernel (agent)       verification

Tool What it does profile.py Profiles any PyTorch model with torch.profiler, ranks kernels by GPU time, classifies as compute/memory-bound extract.py Extracts top-N bottleneck kernels into standalone Triton or CUDA C++ kernel files (--backend triton|cuda) orchestrate.py Multi-kernel scheduler: decides which kernel to optimize next using Amdahl's law, tracks aggregate progress bench.py Fixed benchmark: 5-stage correctness (smoke, shape sweep, numerical stability, determinism, edge cases) + performance + roofline verify.py Plugs optimized kernels back into the model, checks end-to-end correctness, reports total speedup

Supported Kernels

9 kernel types covering the core operations of modern deep learning:

Kernel Description Key Metric matmul Dense matrix multiplication (M x K) @ (K x N) TFLOPS softmax Row-parallel numerically stable softmax GB/s layernorm Layer normalization with affine transform GB/s rmsnorm RMS normalization (LLaMA-style) GB/s flash_attention Scaled dot-product attention with causal masking TFLOPS fused_mlp SwiGLU-style fused MLP (gate + up + down) TFLOPS cross_entropy Fused cross entropy loss GB/s rotary_embedding Rotary position embeddings (RoPE) GB/s reduce Parallel reduction (sum) GB/s

Each has a PyTorch reference in reference.py, a starter Triton kernel in kernels/, and a starter CUDA C++ kernel in kernels/cuda/.

Example Models

Self-contained model definitions ship with AutoKernel (no transformers library needed):

Model File Params Usage GPT-2 Small models/gpt2.py 124M --class-name GPT2 --input-shape 1,1024 LLaMA (compact) models/llama_7b.py 160M --class-name LlamaModel --input-shape 1,512 LLaMA 7B models/llama_7b.py 7B --class-name LlamaModel7B --input-shape 1,2048 BERT-base models/bert_base.py 110M --class-name BertModel --input-shape 8,512 Custom models/custom.py -- Template for your own model

For HuggingFace models (uv sync --extra models):

uv run profile.py --module transformers --class-name AutoModelForCausalLM \
 --pretrained meta-llama/Llama-2-7b-hf --input-shape 1,2048 --dtype float16

KernelBench Integration

AutoKernel integrates with KernelBench, the standard benchmark for evaluating AI-generated GPU kernels (250+ problems across 4 difficulty levels). While most KernelBench evaluations use one-shot LLM generation, AutoKernel runs 50-300+ iterative refinement experiments per problem -- systematically exploring the optimization space instead of guessing.

# Install KernelBench dependencies
uv sync --extra kernelbench

# Fetch Level 1 problems from HuggingFace
uv run kernelbench/bridge.py fetch --source hf --level 1

# Set up a specific problem for optimization
uv run kernelbench/bridge.py setup --level 1 --problem 1 --source hf

# Evaluate (correctness + speedup vs PyTorch reference)
uv run kernelbench/bench_kb.py

# Batch score an entire level (computes fast_p metric)
uv run kernelbench/scorer.py --level 1

The agent reads kernelbench/program_kb.md for KernelBench-specific optimization instructions: how to write ModelNew classes, when to use CUDA C++ vs Triton, fusion strategies per problem level, and the edit-bench-keep/revert loop adapted for the KernelBench fast_p metric.

Tool What it does kernelbench/bridge.py Loads problems from HuggingFace or local repo, caches them, generates starter kernel.py kernelbench/bench_kb.py Evaluates ModelNew vs Model: 5-trial correctness + CUDA event timing + stability + determinism kernelbench/scorer.py Batch evaluation across a level, computes fast_p at thresholds (1.0x, 1.5x, 2.0x, 3.0x, 5.0x) kernelbench/program_kb.md Agent instructions for KernelBench mode

HuggingFace Kernels Export

Export optimized kernels to the HuggingFace Hub for easy distribution. Users can then load your kernels with a single line:

from kernels import get_kernel
module = get_kernel("your-username/kernel-name")

# Export an optimized CUDA kernel
uv run export_hf.py --name my_matmul

# Upload to Hub (requires `pip install kernels` and `huggingface-cli login`)
cd workspace/hf_export/my_matmul
kernels upload . --repo_id your-username/my_matmul

Project Structure

…

Design Choices

Dual backend: Triton + CUDA C++. Triton for fast iteration (Python-like syntax, compiles in seconds). CUDA C++ for maximum performance (direct access to tensor cores via wmma, PTX intrinsics, shared memory bank-conflict-free layouts). Triton regularly reaches 80-95% of cuBLAS; CUDA C++ can match or exceed it. Both backends share the same kernel_fn() interface -- bench.py runs identically on either.

Correctness first. The benchmark checks kernel output against PyTorch before measuring performance. A fast but wrong kernel is immediately reverted. This prevents the agent from "optimizing" by producing garbage.

Amdahl's law orchestration. The orchestrator prioritizes by impact. A 1.5x speedup on a 60% kernel (1.25x end-to-end) beats a 3x speedup on a 5% kernel (1.03x end-to-end). It moves on when diminishing returns set in.

Single file to modify. The agent only touches kernel.py. Scope stays manageable, diffs reviewable, reverts clean.

TSV logging. Results go to a plain results.tsv file. Human-readable, git-friendly, trivially parseable, no infrastructure.

Results Format

Every experiment is logged to results.tsv (tab-separated):

Column Description experiment Sequential experiment number (0 = baseline) tag Short identifier kernel_type Which kernel (e.g., matmul) throughput_tflops Measured throughput (higher is better) latency_us Execution time in microseconds pct_peak Percentage of GPU theoretical peak speedup_vs_pytorch Speedup vs PyTorch/cuBLAS correctness PASS, FAIL, TIMEOUT, or CRASH peak_vram_mb Peak GPU memory usage description What was tried

Credits

This project is autoresearch for GPU kernels -- directly inspired by Andrej Karpathy's autoresearch, the original experiment in autonomous AI research agents for LLM training. Karpathy showed that an AI agent can run hundreds of experiments overnight, methodically exploring a search space and logging every result. AutoKernel applies that same loop -- agent edits one file, runs a fixed evaluation, keeps or reverts -- to the domain of GPU kernel optimization with Triton and native CUDA C++.

KernelBench integration is based on the work of Simon Guo, Sean Resta, et al. at Stanford's Scaling Intelligence Lab. Their paper "KernelBench: Can LLMs Write GPU Kernels?" (2025) established the standard benchmark for evaluating AI-generated GPU kernels. AutoKernel extends this by applying iterative optimization (300+ experiments per problem) instead of one-shot generation. KernelBench dataset and evaluation protocol: ScalingIntelligence/KernelBench.

Built by RightNow AI. For enterprise GPU optimization, check out RightNow Enterprise.

Changelog

v1.3.0

• AMD ROCm GPU support: MI300X, MI325X, MI350X, MI355X detection and specs (thanks @andyluo7)
• Fixed verify.py SyntaxError on Python 3.13+
• Fixed CUDA flash_attention ignoring sm_scale parameter
• Fixed CUDA cross_entropy returning wrong dtype
• Fixed Triton rotary_embedding broadcasting truncation
• Fixed Triton reduce output shape for non-last-dim reductions

v1.2.0

• Enhanced profiler: --export-trace, --memory-snapshot, --torch-compile-log flags
• HuggingFace Kernels export via export_hf.py

v1.1.0

• Native CUDA C++ backend with 9 starter kernels (tensor cores, warp intrinsics, shared memory tiling)
• KernelBench integration (250+ standardized GPU kernel problems)
• --backend triton|cuda flag for extract.py

v1.0.0

• Initial release: Triton kernel optimization pipeline with 5-stage correctness harness

See CHANGELOG.md for full details.

License

MIT