Skip to content
Search

AutoKernel

lastmod 2026-03-12created 2026-03-121 edits

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

Autonomous research system for GPU kernels

About

  • Applies Karpathy's Autoresearch to GPU kernel optimization
  • A GPU kernel autoresearch tool that takes a PyTorch model and automatically optimizes Triton or CUDA C++ kernels
  • Provides a fully automated pipeline that profiles, extracts, optimizes, and verifies a model's bottleneck kernels
  • Adjusts priorities with orchestration based on Amdahl's law; each experiment takes about 90 seconds, allowing hundreds of iterations overnight
  • Integrated with KernelBench, running 50–300 iterative experiments on 250+ problems, enabling systematic search rather than one-shot generation
  • Supports efficient, reproducible kernel optimization with dual Triton and CUDA C++ backends, correctness-first verification, and a single-file modification structure

How it works

  • AutoKernel takes a PyTorch model and performs the following steps
    • Identifies GPU bottleneck kernels through profiling
    • Extracts them into standalone Triton or CUDA C++ kernels
    • Repeats modify, benchmark, keep/revert through an automated optimization loop
    • Verifies correctness, then reports the overall speedup
  • Modifies kernel.py based on the instructions in program.md, and runs a 5-stage correctness check and roofline analysis through bench.py
  • Each experiment takes about 90 seconds: about 40 per hour, about 320 overnight

Running and components

  • Requirements: NVIDIA GPU (H100/A100/RTX 4090), Python 3.10+, uv package
  • Main scripts
    • profile.py: ranks kernels by GPU time
    • extract.py: extracts the top bottleneck kernels
    • bench.py: 5-stage correctness and performance benchmark
    • orchestrate.py: multi-kernel scheduling based on Amdahl's law
    • verify.py: full-model verification and speedup report
  • program.md includes a 6-tier optimization strategy, crash handling, and a decision-making framework, enabling long autonomous runs

Supported kernels and example models

  • Supports 9 kernel types: matmul, softmax, layernorm, rmsnorm, flash_attention, fused_mlp, cross_entropy, rotary_embedding, reduce
  • Each kernel includes a PyTorch reference implementation (reference.py) and initial Triton/CUDA versions (kernels/, kernels/cuda/)
  • Example models: GPT-2, LLaMA (including 7B), BERT-base, custom models
  • HuggingFace models can also be integrated with uv sync --extra models

KernelBench integration

  • Integrated with KernelBench (Stanford Scaling Intelligence Lab) to run the standard benchmark for AI-generated GPU kernels
  • AutoKernel systematically searches the optimization space with 50–300 iterative experiments per problem
  • Main tools
    • bridge.py: loads problems and generates initial kernels
    • bench_kb.py: correctness and performance evaluation
    • scorer.py: computes scores for entire levels
    • program_kb.md: KernelBench-specific agent instructions

HuggingFace Hub export

  • Optimized kernels can be exported to the HuggingFace Hub and loaded easily with get_kernel()
  • export_hf.py supports exporting and uploading CUDA kernels

Design principles

  • Dual Triton + CUDA C++ backends: Triton for fast iteration, CUDA for peak performance
  • Correctness first: reverts immediately if results don't match PyTorch
  • Orchestration based on Amdahl's law sets priorities by contribution to overall performance
  • Single-file modification structure (kernel.py) makes change tracking and recovery easy
  • Records experiment results simply and readably in a TSV log (results.tsv)

Result log format

  • results.tsv records each experiment's number, kernel type, throughput (TFLOPS), latency, percentage of GPU peak, speedup over PyTorch, correctness, VRAM usage, and description

Project background

  • Inspired by Andrej Karpathy's autoresearch concept, it applies an autonomous AI agent structure for LLM research to GPU kernel optimization
  • The KernelBench integration builds on research from the Stanford Scaling Intelligence Lab, and AutoKernel performs iterative optimization instead of one-shot generation
  • The project was developed by the Forge team at RightNow AI and released under the MIT license

See also

Favorite site

author Zer0 sitemap