PaperlayerView exact revision on arXiv

Source-linked AI summary

AutoKernel: Autonomous GPU Kernel Optimization via Iterative Agent-Driven Search

Jaber Jaber, Osama Jaber

arXiv:2603.21331v1cs.LGcs.PF

TL;DR

GPU kernel optimization is labor-intensive because performance depends on scarce microarchitectural expertise and library coverage can lag new model operations. AutoKernel automates model profiling, Amdahl-prioritized kernel search, and correctness-gated refinement across Triton and CUDA C++, with community deployments reaching leading performance results. It remains limited to individual kernels on a single GPU and inherits the capabilities of its underlying LLM.

  • Problem

    GPU kernel optimization requires scarce expertise, while evolving model operations can outpace vendor-library support.

  • Method

    AutoKernel profiles complete PyTorch models, ranks bottleneck kernels by Amdahl’s law, and iteratively edits, validates, benchmarks, and keeps or reverts Triton or CUDA C++ candidates.

  • Results

    AutoKernel-optimized kernels achieved leading community results, including first place on the vectorsum_v2 NVIDIA B200 leaderboard and 1.63× to 2.15× speedup over CUTLASS fused kernels for FP4 matmul on H100.

  • Takeaways & Limitations

    The system turns kernel optimization into an overnight autonomous process that combines model-level prioritization, correctness validation, and iterative search.

  • Takeaways & Limitations

    AutoKernel currently optimizes individual kernels on a single GPU, while distributed kernels and multi-device memory management remain out of scope.

Abstract

from arXiv · show

Writing high-performance GPU kernels is among the most labor-intensive tasks in machine learning systems engineering. We present AutoKernel, an open-source framework that applies an autonomous agent loop to GPU kernel optimization for arbitrary PyTorch models. Given a model, AutoKernel profiles it to identify computational bottlenecks, ranks them by Amdahl's law impact, and iteratively refines Triton or CUDA C++ kernel implementations through hundreds of experiments without human intervention. A five-stage correctness harness covering smoke tests, shape sweeps, numerical stability, determinism verification, and edge-case coverage ensures every candidate kernel is validated before any speedup is recorded. The system comprises over 9,000 lines of Python, 18 starter kernel implementations across two backends, a six-tier optimization playbook, and integration with the KernelBench benchmark suite. AutoKernel covers nine kernel types spanning the dominant operations in modern transformer architectures. On an NVIDIA H100, our Triton kernels outperform both PyTorch eager and torch.compile (max-autotune) on the majority of tested configurations: 5.29x over eager on RMSNorm, 2.82x on softmax, and 2.21x on cross-entropy, while beating torch.compile by 2.83x, 3.44x, and 2.94x respectively. In community deployment, an AutoKernel-optimized kernel achieved first place on the vectorsum_v2 B200 leaderboard. The full system is available at https://github.com/RightNow-AI/autokernel.

1 Introduction

AutoKernel targets the labor-intensive gap between GPU hardware capability and default library performance by automating expert kernel optimization. It profiles complete PyTorch models, prioritizes bottlenecks by Amdahl’s law, and iteratively searches kernel implementations with correctness checks.

  • 60 to 80% of total GPU time is typically consumed by matrix multiplications, while normalization, softmax, and positional embeddings account for much of the remainder.
  • GPU kernel optimization requires expertise in arithmetic intensity, memory behavior, register pressure, occupancy, tiling, synchronization, and tensor-core instruction selection.
  • Fewer than 20% of KernelBench cases matched PyTorch baseline performance under one-shot frontier-LLM generation.
  • AutoKernel uses an edit, benchmark, keep-or-revert loop in which each iteration takes roughly 90 seconds and overnight runs produce 300 to 400 experiments.
  • The framework adapts the autoresearch keep-or-revert paradigm to kernel implementations, using correctness-gated benchmarking instead of validation loss.
  • Amdahl’s law ranks kernels by end-to-end impact, so optimization effort targets kernels whose runtime contribution offers the greatest potential speedup.
  • The system combines an open-source pipeline, five-stage correctness harness, dual Triton/CUDA C++ backends, Amdahl-based orchestration, and a six-tier optimization playbook.

2 Related Work

Related systems span Triton and CUDA programming, compiler-generated kernels, benchmark infrastructure, hardware-aware agents, and evolutionary search. AutoKernel distinguishes itself by combining model-level profiling, dual backends, and a transparent correctness-gated keep/revert loop.

  • Triton abstracts GPU programming around block-level tensor operations, while CUDA C++ provides explicit control over warp primitives, tensor cores, and shared-memory layouts.
  • AutoKernel combines model-level profiling with dual Triton/CUDA C++ backend support, a combination identified as unique in the comparison table.
  • AutoKernel supports both Triton for rapid iteration and CUDA C++ for maximum control within one framework.
  • KernelBench provides standard evaluation infrastructure, while GEAK, CudaForge, KernelFoundry, and GPU Kernel Scientist explore agentic, hardware-aware, or evolutionary optimization.
  • Unlike multi-agent or learned-policy systems, AutoKernel uses edit, benchmark, keep/revert decisions with a five-stage correctness harness, trading architectural complexity for transparency and reliability.
  • AutoKernel applies the autoresearch paradigm to a different search space and evaluation function: kernel implementations and correctness-gated performance.

3 System Design

AutoKernel profiles models, extracts bottleneck kernels, and runs an agent-driven optimization loop across Triton and CUDA C++ implementations. A correctness-gated benchmark, Amdahl-based orchestration, and explicit expert playbooks govern which candidates persist.

  • The system comprises over 9,200 lines of Python, 18 kernel implementations, four model definitions, and a 909-line agent instruction document organized across three phases.
  • The profiler accepts local files, HuggingFace identifiers, or custom model classes and records per-kernel GPU time with shape information.
  • Kernel names are pattern-matched into nine operation types across cuBLAS, CUTLASS, Triton, and ATen implementations.
  • The extractor generates standalone kernel files containing starter code, model-specific shape variants, roofline formulas, and dtype-specific tolerances.
  • Each candidate is edited in a single kernel file, benchmarked for correctness and performance, then kept only when it passes and exceeds the current best by more than 1%.
  • Amdahl’s law ranks optimization plans using each kernel’s fraction of total GPU time and achieved speedup, with what-if projections at 1.5×, 2×, 3×, and 5×.
  • The six-tier playbook covers block-size tuning, memory access, computation, advanced techniques, architecture-specific methods, and kernel-specific strategies.
  • The orchestrator moves to another kernel after five consecutive reverts, 90% of GPU peak, two hours, or 2× speedup.

4 Five-Stage Correctness Verification

AutoKernel measures performance only after candidates pass five correctness stages designed to catch compilation, shape, numerical, determinism, and edge-case failures. The harness spans shape and dtype sweeps, adversarial inputs, repeated executions, and non-power-of-two dimensions.

  • All five correctness stages must pass before performance is measured.
  • Stage 1 runs a small-input smoke test to catch compilation errors, shape mismatches, and gross numerical bugs.
  • Stage 2 tests 8 to 10 configurations and three data types to expose boundary, tile-remainder, and dtype-specific bugs.
  • Stage 3 probes numerical stability with adversarial inputs such as large identical softmax rows, extreme matmul ranges, and near-zero normalization variance.
  • Stage 4 executes identical inputs three times and requires bitwise-identical outputs to catch reduction races and nondeterministic atomics.
  • Stage 5 tests non-power-of-two dimensions including 1023, 4097, and 1537 to expose masking and tile-remainder errors.
  • Tolerance thresholds are dtype-specific: 10^-2 for FP16, 2 × 10^-2 for BF16, and 10^-4 for FP32.

5 Dual Backend: Triton and CUDA C++

AutoKernel supports both Triton and CUDA C++ backends for GPU kernel optimization, exposing a shared interface so benchmarking is backend-independent.

  • Nine Triton starter kernels use a Python-like DSL and allow optimization of block sizes, warps, stages, accumulator precision, and loop structure.They are JIT-compiled in 1 to 5 seconds and routinely reach 80 to 95% of cuBLAS throughput for matmul.
  • Nine CUDA C++ starter kernels provide direct access to tensor cores, warp shuffles, vectorized loads, shared-memory layouts, double buffering, and register controls.Compilation includes architecture auto-detection, hash-based caching, and threadsafe builds.
  • Both backends expose the same kernel_fn() interface, allowing the benchmark to run identically across implementations.

6 Kernel Coverage

AutoKernel covers nine supported kernel types and evaluates each against a PyTorch reference using throughput and roofline-utilization measurements.

  • The system supports nine kernel types spanning the dominant operations targeted by the framework.The section refers to the complete list in Table 3.
  • Each kernel has a PyTorch reference implementation that serves as the correctness oracle.Benchmarking also measures throughput in TFLOPS or GB/s and roofline utilization against detected GPU peak.
  • Four self-contained model definitions ship with the system: GPT-2 124M, LLaMA 160M/7B, BERT-base 110M, and a custom template.They require no external dependencies such as transformers.

7 Experimental Evaluation

On an H100, AutoKernel’s Triton starters outperform the baselines on many memory-bound configurations, while matmul remains a harder optimization target. Community results further show competitive performance against established GPU implementations.

  • 7.1 Kernel Performance: Evaluation uses FP16 CUDA-event timing on an NVIDIA H100, comparing PyTorch eager with torch.compile max-autotune and AutoKernel Triton starters.Measurements use 200 iterations per configuration with trimmed means, and Table 4 shows 16 representative configurations from 34 tested.
  • 7.1 Kernel Performance: 5.29× over eager and 2.83× over torch.compile are achieved by RMSNorm at the largest tested size.RMSNorm reaches 2,788 GB/s, or 83% of H100 peak bandwidth; cross-entropy reaches 2,070 GB/s and softmax reaches 2,800 GB/s.
  • 7.1 Kernel Performance: 12 of 16 shown configurations beat torch.compile, while all 34 configurations pass all five verification stages with zero failures.The gains are attributed to single-pass Triton fusion that reduces HBM traffic; matmul remains below cuBLAS but beats torch.compile by 1.55× at 2048^3.
  • 7.2 Community Deployment Results: 44.086µs secured first place on the NVIDIA B200 vectorsum_v2 leaderboard, ahead of the second-place 44.249µs entry.The winning kernel was produced by an overnight iterative search over block sizes, warp-level reductions, and vectorized memory access.
  • 7.2 Community Deployment Results: 1.63× to 2.15× speedups over CUTLASS were reported for a community-generated Triton FP4 matmul kernel across multiple H100 shapes.The kernel reached up to 2,898 TFLOPS and was generated through a single agent interaction of approximately three minutes.
  • 7.3 Optimization Loop Dynamics: RMSNorm optimization typically gains 10 to 30% from block-size sweeps, 10 to 20% from memory improvements, and 5 to 10% from epilogue fusion.After 30 to 50 experiments, increasing reverts trigger the orchestrator’s move-on behavior.
  • 7.3 Optimization Loop Dynamics: Matmul requires more experiments because its search spans tile dimensions, warp counts, pipeline stages, and accumulator precision.The starter reaches 278 TFLOPS versus cuBLAS at over 800 TFLOPS, leaving a wider optimization gap.
  • 7.3 Optimization Loop Dynamics: Move-on criteria include five consecutive reverts, 90% peak utilization, a two-hour timeout, or a 2× speedup threshold.These criteria limit time spent on kernels with diminishing returns and prioritize improvements according to their end-to-end impact.

8 KernelBench Integration

AutoKernel integrates with KernelBench through problem loading, evaluation, and batch scoring components, while replacing one-shot generation with iterative optimization.

  • The KernelBench problem loader fetches datasets or local clones, analyzes operations, and generates starter ModelNew classes.The bridge.py component contains 674 lines.
  • The evaluation harness checks correctness, numerical stability, determinism, and performance using repeated trials and CUDA timing.Its checks include five correctness trials at atol = 10^-2, NaN/Inf detection, three determinism runs, and trimmed-median performance.
  • The batch scorer computes fast_p at seven thresholds ranging from 1.0× through 5.0×.
  • AutoKernel runs 50 to 300 iterative experiments per problem instead of relying on the one-shot scores commonly reported for KernelBench.

9 HuggingFace Kernels Export

The export tool packages optimized CUDA kernels for HuggingFace Hub distribution and generates build, binding, and cross-compilation metadata.

  • The 868-line export_hf.py tool packages optimized CUDA kernels for distribution through the HuggingFace Hub.
  • It generates build.toml with backend and Hub metadata, torch_binding.cpp with schema-based registration, and flake.nix for cross-compilation.
  • Users install distributed optimized kernels through get_kernel("rightnow-ai/matmul").

10 Design Rationale

AutoKernel favors a simple, auditable optimization loop: a single agent changes mutable kernel code while fixed, hardware-informed evaluation determines which experiments persist.

  • A single agent in a tight loop avoids the coordination overhead and failure modes of specialized multi-agent systems.
  • The benchmark remains immutable while the agent modifies candidate code, preventing candidates from gaming their evaluator.
  • Each experiment maps to a git commit, so kept changes advance the branch and reverted changes disappear through git reset.
  • Roofline classification and percentage of peak guide the agent toward the appropriate optimization tier after each experiment.
  • Tab-separated result files provide dependency-free, human-readable, git-friendly experiment tracking.

11 Limitations and Future Work

AutoKernel’s current scope is bounded by LLM code-generation capabilities and single-GPU individual-kernel optimization, while future work targets broader search and cross-kernel strategies.

  • Limitations: AutoKernel may struggle with software pipelining, custom PTX emission, and multi-CTA cooperative strategies because it inherits its underlying LLM’s capabilities.
  • Limitations: The current system optimizes individual kernels on a single GPU; distributed kernels and multi-device memory management remain out of scope.
  • Future Work: Future directions include population-based multi-GPU search, learned search policies, profiling-guided mutations, and cross-kernel fusion discovery.
  • Conclusion: AutoKernel combines model-level profiling, Amdahl’s-law prioritization, correctness validation, dual backends, and a six-tier playbook into an overnight autonomous process.
Loading 2603.21331v1…