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

Jaber Jaber; Osama Jaber

cs.LGMar 22, 2026

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

Authors: Jaber Jaber, Osama Jaber

Organizations: RightNow AI

Abstract

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.

Figures & tables

System	Method	Backend	Model	Iterative	Open
			prof.		src
KernelBench ( Ouyang et al., 2025 )	Benchmark	Multi-DSL	✗	✗	✓
GEAK ( Wang et al., 2025 )	Multi-agent	Triton	✗	✓	✓
CudaForge ( Zhang et al., 2025 )	Coder+Judge	CUDA	✗	✓	✓
KernelFoundry ( Wiedemann et al., 2026 )	Evolutionary	CUDA/SYCL	✗	✓	✓
Kernel Scientist ( Andrews and Witteveen, 2025 )	Evolutionary	CUDA	✗	✓	✓

Table 1: Comparison of LLM-based GPU kernel optimization systems. AutoKernel is the only system that combines model-level profiling with dual Triton/CUDA C++ backend support.

Panel of Figure 1 — Figure 1: AutoKernel architecture. Phase A profiles the model and extracts bottleneck kernels. Phase B runs the autonomous optimization loop: the agent edits kernel.py , the benchmark verifies through five correctness stages, and the orchestrator decides to keep, revert, or move to the next kernel. Phase C verifies end-to-end correctness and speedup.

Panel of Figure 2 — Figure 2: Five-stage correctness pipeline. Any failure immediately rejects the candidate. Throughput is only measured after all five stages pass. Each stage catches a distinct class of bugs.

Name	M	N	K	Purpose
tiny	128	128	128	Minimum viable
small	512	512	512	Small workload
large	2048	2048	2048	Production-like
xlarge	4096	4096	4096	Stress test
tall	8192	1024	1024	Non-square
deep_k	1024	1024	8192	Large reduction dim

Table 2: Shape sweep for matmul. Other kernels have analogous sweeps with 8 to 10 configurations covering tiny, standard, large, non-square, and workload-specific dimensions.

Kernel	Regime	Metric	Key Starter Techniques
matmul	Compute	TFLOPS	128 $\times$ 128 tiles, WMMA, double-buffered smem
flash_attn	Compute	TFLOPS	Tiled online softmax, causal masking
fused_mlp	Compute	TFLOPS	SwiGLU gate-up fusion
softmax	Memory	GB/s	Warp-shuffle reductions, half2 loads
layernorm	Memory	GB/s	Welford single-pass, float4 loads
rmsnorm	Memory	GB/s	Warp shuffle cascade, fast rsqrt

Table 3: Supported kernel types with performance regime and starter implementation techniques.

Kernel	Size	Eager	Compiled	Ours	vs Eager	vs Compiled	Thru.
		( $\mu$ s)	( $\mu$ s)	( $\mu$ s)
matmul	2048 3	28.1	101.2	65.3	0.43 $\times$	1.55 $\times$	263 TF/s
	4096 3	182.8	257.8	494.2	0.37 $\times$	0.52 $\times$	278 TF/s
	8192 3	1679.5	1916.1	5773.1	0.29 $\times$	0.33 $\times$	190 TF/s
softmax	4096 2	58.9	96.1	40.1	1.47 $\times$	2.40 $\times$	1675 GB/s
softmax	8192 2	270.4	330.0	95.9	2.82 $\times$	3.44 $\times$	2800 GB/s

Table 4: Kernel performance on H100 (FP16), selected configurations. Three baselines: PyTorch eager (cuBLAS/ATen), torch.compile ( max-autotune ), and AutoKernel Triton starters. 16 representative configurations shown from the full 34 tested (all passing correctness). Full evaluation completes in under 10 minutes.

Shape (M $\times$ N $\times$ K)	Triton	CUTLASS	Speedup
	(TF/s, ms)	(TF/s, ms)
128 $\times$ 3072 $\times$ 3072	186 TF, 0.013	100 TF, 0.024	1.86 $\times$
128 $\times$ 18432 $\times$ 3072	1105 TF, 0.013	550 TF, 0.026	2.01 $\times$
1024 $\times$ 3072 $\times$ 3072	1477 TF, 0.013	686 TF, 0.028	2.15 $\times$
1024 $\times$ 18432 $\times$ 3072	2777 TF, 0.042	1609 TF, 0.072	1.73 $\times$
2048 $\times$ 3072 $\times$ 3072	1662 TF, 0.023	964 TF, 0.040	1.72 $\times$

Table 5: Community-reported FP4 matmul results on H100: AutoKernel Triton kernel vs CUTLASS fused baseline. The Triton kernel was generated with a single agent prompt (

\sim

3 minutes).

GPU Kernel Optimization

GPU kernel optimization improves the speed and memory efficiency of computations by tailoring code to GPU hardware and workloads.

See all 202 papers

KernelOPT: Dispatch-Aware Agentic Search for GPU Kernel Optimization
Sep 24, 2026Aheli Poddar, Sanskar Prasad, Arindam Samanta +3GPU Kernel Optimization Agentic Search
Red Hat
KernelClaw: A Multi-Agent Framework for GPU Kernel Optimization with Dual-Memory Mechanisms
Mar 10, 2026Qitong Sun, Jun Han, Tianlin Li +6GPU Kernel Optimization GPU Kernel Generation
School of Computer Science and Engineering, Beihang University, China · Zhejiang Lab, China · Nanyang Technological University, Singapore
HIERA: Workload-Aware Planning Across Implementation Spaces for GPU Kernel Optimization
Aug 21, 2026Jinghao Wang, Qiqi Gu, Chenpeng Wu +3GPU Kernel Optimization GPU Kernel Generation
Shanghai Key Laboratory of Scalable Computing and Systems · School of Computer Science, Shanghai Jiao Tong University

Sep 24, 2026J/K move · Enter open · S save

KernelOPT: Dispatch-Aware Agentic Search for GPU Kernel Optimization

Aheli Poddar, Sanskar Prasad, Arindam Samanta, Subha Chakraborty+2

Red Hat

GPU Kernel Optimization Agentic Search Large Language Model-Guided Optimization Compiler Optimization Neural Network Hardware Acceleration

Open paper AskPDF arXiv

Deep learning inference and training performance depends critically on GPU kernel efficiency. Modern compilers such as PyTorch Inductor automatically generate GPU kernels from high-level model code, but frequently underperform expert-written implementations by wide margins. Recent LLM-assisted kernel optimizers can close this gap for standalone kernels, yet treat compiled models as black boxes, generally optimizing individual standalone kernels without respecting the compiler's structural decisions or verifying the model end-to-end. We present KernelOPT, a multi-agent system that treats compiled models as structured artifacts. It preserves vendor library calls (cuBLAS, cuDNN) and exclusively targets generated Triton sub-kernels using five profiling-guided LLM agents. A four-gate verification cascade applies static validation, multi-seed correctness checking, model-level float64-fallback verification, and performance gating (

γ{=}1.03

) to filter candidates and verify the re-stitched model end-to-end. When candidates fail verification, the system preserves the compiler baseline. The system accepts PyTorch nn Modules, standalone Triton kernels, and Helion kernels. Evaluated on 250 KernelBench problems (100 Level 1, 100 Level 2 and 50 Level 3) on NVIDIA H200, KernelOPT achieves geometric mean speedups over torch compile of 1.40

\times

(L1), 1.15

\times

(L2), and 1.07

\times

(L3) across all kernels, including fallback cases. Optimized-only geomeans (excluding cases where verification gates preserve the compiler baseline) are substantially higher: 2.54

\times

(L1: 36/100), 1.84

\times

(L2: 23/100), and 1.37

\times

(L3: 11/50), reflecting where the optimizer achieves meaningful leverage.

Abstract

GPU Kernel Optimization

Explore similar work

KernelOPT: Dispatch-Aware Agentic Search for GPU Kernel Optimization

KernelClaw: A Multi-Agent Framework for GPU Kernel Optimization with Dual-Memory Mechanisms

HIERA: Workload-Aware Planning Across Implementation Spaces for GPU Kernel Optimization