KernelOPT: Dispatch-Aware Agentic Search for GPU Kernel Optimization

Aheli Poddar; Sanskar Prasad; Arindam Samanta; Subha Chakraborty; Vishal Goyal; Rohit Singh Rathaur

cs.DCSep 24, 2026

KernelOPT: Dispatch-Aware Agentic Search for GPU Kernel Optimization

Authors: Aheli Poddar, Sanskar Prasad, Arindam Samanta, Subha Chakraborty, Vishal Goyal, Rohit Singh Rathaur

Organizations: Red Hat

Abstract

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 of static validation, multi-seed correctness, model-level float64-fallback verification, and performance gating filters candidates during optimization and verifies the re-stitched model end-to-end. If no candidate passes all four gates, the system preserves the compiler baseline. The system accepts PyTorch nn.Modules, standalone Triton kernels, and Helion kernels. Evaluated on 250 KernelBench problems, KernelOPT achieves geometric mean speedups over \texttt{torch.compile} of 1.40 $\times$ (Level 1: 51/100), 1.15 $\times$ (Level 2: 31/100), and 1.07 $\times$ (Level 3: 12/50) across all problems.

Figures & tables

Figure 1: KernelOPT pipeline. A PyTorch model is traced via Inductor; sub-kernels are classified as Triton-generated (optimizable) or extern-library (preserved). Five agents iterate within a LangGraph state machine using NCU profiling feedback. Optimized kernels are re-stitched and verified at the model level.

A. Summary								B. Speedup distribution (94 opt.)
Level	Total	Opt.	M.(opt)	M.(s.f.)	Fall.	Geo.	Max	Range	L1	L2	L3	Tot.
L1	100	51	20	13	16	1.40 $\times$	88.63 $\times$	$\geq$ 10 $\times$	4	1	0	5
L2	100	31	35	0	34	1.15 $\times$	33.11 $\times$	$[5,10)$	6	4	0	10
L3	50	12	3	0	35	1.07 $\times$	4.12 $\times$	$[2,5)$	7	0	2	9
								$[1.5,2)$	2	2	1	5
								$[1.1,1.5)$	8	8	4	20

Table 1: Comprehensive KernelBench results vs. torch.compile . Panel A : Level summary where Match (opt) produced correct output at

{\sim}1.0\times

and Match (synth. fail) indicates LLM synthesis failure. Panel B : Speedup distribution. Panel C : Dominant operation breakdown. Panel D : Fallback root causes (cuBLAS/cuDNN dominance accounts for 72%).

Config	Opt	Match	Fall	Opt%
Full system	19	14	17	38%
No beam ( $B{=}1$ )	5	3	41	10%
No NCU profiling	5	4	41	10%
No memory	4	3	43	8%
No multi-iter ( $T{=}1$ )	4	4	42	8%
No perf gate	3	3	44	6%

Table 2: Ablation study on 50 weighted-random-sampled problems (

n{=}50

; “No inductor-aware” runs on L2 only,

n{=}20

). Removing any single component reduces optimization success by 74–84%. Disabling Inductor-aware synthesis eliminates all L2 optimization. A per-architecture breakdown of the L3 results is provided in Table 10 (Appendix F ).

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Stage	Gate	Tolerance	Seeds	Reference	Scope
Synthesis	Eager validation	$r{=}a{=}10^{-4}$	2: $(0,42)$	PyTorch eager model	Per synth. kernel
Gate 1	Static validation	N/A (crash check)	N/A	N/A	Per candidate
Gate 2	Multi-seed corr.	$r{=}a{=}10^{-3}$	3: $[0,42,1337]$	Synthesized Triton	Per candidate
Gate 3	Model-level (Tier 1)	$r{=}a{=}10^{-4}$	3: $[0,42,123]$	Re-stitched model	End-to-end
	Model-level (Tier 2)	$\rho\leq 10$ or abs bound		Float64 reference	End-to-end
Gate 4	Performance gate	$\gamma=1.03$ (3%)	N/A	do_bench (25/100 ms)	End-to-end

Appendix

Table 3: Complete verification gate specification. Gates 1–2 run on every candidate during optimization; Gates 3–4 run once on the re-stitched model after optimization completes.

r

and

a

denote rtol and atol .

Category	Strategy	Succ.	Fail.
Matmul	SMEM bank conflict padding	2	0
	TF32 tensor cores	2	0
	Add Triton autotune	4	0
	SMEM padding (alt)	2	0
Pointwise	Add Triton autotune	130	2
	TF32 tensor cores	11	1

Appendix

Table 4: Shipped optimization guidelines by kernel category. Success/failure counts reflect how often a strategy produced a speedup vs. a regression during post-hoc extraction from the 250-kernel evaluation. The extraction script recorded binary success/fail per attempt; per-strategy speedup magnitudes were not tracked, so no average or maximum speedup column is included. Strategies with 0 failures and

{\geq}2

successes are included in the default cold-start context.

Config	Ablated Component	Flag
C0	Full system (baseline)	(default)
C1	Greedy search (no beam)	--beam-width 1
C2	No optimization memory	--no-memory
C3	No E2E performance gate	--no-e2e-perf-gate
C4	No NCU profiling context	--skip-ncu
C5	Single iteration ( $T{=}1$ )	-T 1

Appendix

Table 5: Ablation configurations. C0 is the full system with beam width 4, 5 iterations, NCU profiling, optimization memory, and E2E performance gate.

	L1 (20)			L2 (20)			L3 (10)
Config	Opt	Match	Fall	Opt	Match	Fall	Opt	Match	Fall	Geomean
C0 (full)	10	6	4	6	7	7	3	1	6	1.94 $\times$
C1 (no beam)	4	3	13	1	0	19	0	0	10	1.78 $\times$
C2 (no memory)	4	3	13	0	0	20	0	0	10	1.23 $\times$
C3 (no perf gate)	3	3	14	0	0	20	0	0	10	1.14 $\times$
C4 (no NCU)	3	4	13	2	0	18	0	0	10	3.16 $\times$

Appendix

Table 6: Ablation results summary. The C0 row shows the full-system results on the same 50-kernel subset. “Opt” = optimized (speedup

{>}1.03\times

), “Match” = within 3% of baseline, “Fall” = fell back to baseline. C6 runs only on L2. Geomean is computed over optimized kernels only.

Figure 2: Ablation study results. Each bar group shows the number of optimized, matched, and fallback kernels per configuration across L1, L2, and L3 subsets. The full system (C0) achieves 19 optimizations; removing any single component reduces this to ${\leq}5$ . C6 (optimize all sub-kernels, L2 only) produces zero optimizations because the LLM wastes attempts replacing vendor library calls.

	Kernel	Comp (%)	Mem (%)	Occ (%)	Planner direction (iter.)	Sub- $k$	E2E
L1	matmul_kernel (007) Regs/thread: 184	57	47	12.4	Iter 1: fast_accum_fp16_dot ; regs 184 $\to$ 125, occ. 12 $\to$ 24%	2.33 $\times$	2.89 $\times$
L2	fused_linear_mul_ hardtanh_gelu (053) Regs/thread: 255 (max)	37	59	12.5	Iter 1: reduce_regs (1.0 $\times$ ); Iter 3: tf32x3_tensor_cores (3.12 $\times$ )	3.12 $\times$	1.60 $\times$
L3	gemm_bias_relu (002, MLP) Regs/thread: 198	60	67	12.2	Iter 1: reduce_reg_pressure (1.02 $\times$ ); Iter 3: k_loop_peeling (1.12 $\times$ )	1.12 $\times$	4.12 $\times$

Appendix

Table 7: Representative optimization traces from each level. Comp / Mem / Occ are baseline NCU metrics (compute utilization %, memory utilization %, and achieved occupancy %). Planner direction shows selected strategy per iteration. Sub-

k

is sub-kernel speedup; E2E is model-level speedup versus compiler baseline.

Kernel	Variant	Comp %	Mem %	Regs	Occ %	Duration	Sub- $k$
L1-007 (matmul)	Baseline	57.1	46.7	184	12.4	5.95 ms	—
E2E 2.89 $\times$	Optimized (fp16 dot)	33.8	75.4	125	24.1	1.82 ms	3.27 $\times$
L1-012 (diag matmul)	Baseline (synth.)	14.0	69.0	18	81.2	32.83 $\mu$ s	—
E2E 88.63 $\times$ ; speedup from algorithmic rewrite ( $O(N^{3})\!\to\!O(N^{2})$ ), not NCU-guided micro-optimization.
L2-018 (fused chain)	Baseline (synth.)	21.8	56.1	29	45.9	12.45 $\mu$ s	—
E2E 33.11 $\times$ ; speedup from fusing 5 operations with algebraic simplification.

Appendix

Table 8: Nsight Compute profiling metrics for three representative kernels. Comp/Mem = compute/memory Speed-of-Light; Regs = registers per thread; Occ = achieved occupancy; Duration = per-launch kernel time. L1-007 metrics match the main paper (Table 7 : 57%, 47%, 12.4%).

Figure 3: Speedup distribution histograms by KernelBench level. L1 (100 kernels): 51 optimized, 34 matched, 15 fallback, geomean 1.40 $\times$ . L2 (100 kernels): 31 optimized, 35 matched, 34 fallback, geomean 1.15 $\times$ . L3 (50 kernels): 14 optimized, 4 matched, 32 fallback, geomean 1.07 $\times$ . All levels report zero regressions: every kernel either improves or preserves the compiler baseline.

Lvl	K#	Speed	Cat	Ops	Rsn	Lvl	K#	Speed	Cat	Ops	Rsn
L1	001	5.77 $\times$	OPT	Matmul	TO	L1	002	6.06 $\times$	OPT	Matmul
L1	003	5.11 $\times$	OPT	Matmul		L1	004	1.30 $\times$	OPT	Matmul
L1	005	1.01 $\times$	OPT			L1	006	0.98 $\times$	MAT	Matmul
L1	007	2.89 $\times$	OPT	Matmul		L1	008	2.63 $\times$	OPT	Matmul	TO
L1	009	2.07 $\times$	OPT	Matmul		L1	010	5.79 $\times$	OPT	Matmul
L1	011	1.05 $\times$	OPT	Einsum		L1	012	88.63 $\times$	OPT	Matmul

Appendix

Table 9: Complete per-kernel KernelBench results across all 250 problems, laid out in two side-by-side blocks (read left block top-to-bottom, then right block). Speed is E2E model speedup vs. torch.compile ; “–” denotes a fallback (baseline preserved).

Figure 4: Root cause distribution for fallback kernels across L1–L3. GEMM-dominant and Conv-dominant cases together account for the majority of fallbacks, reflecting cases where vendor library calls (cuBLAS, cuDNN) are already near-optimal and cannot be improved by Triton kernel replacement.

Architecture	$n$	Sub- $k$	Opt	Speedup	Fall.
MLP	3	3–4	1	4.12 $\times$	GEMM
ResNet	3	8–52	1	1.55 $\times$	Conv
DenseNet	4	3–337	2	1.06–1.15 $\times$	Conv
MobileNet/EffNet	6	7–60	2	1.03–2.33 $\times$	Conv
ViT / Swin / CViT	5	1–99	2	1.10–1.11 $\times$	GEMM
U-Net/NetVLAD/SSM	5	1–16	4	1.06–1.17 $\times$	Perf

Appendix

Table 10: L3 results by model architecture (all 50 kernels). Sub-

k

: Inductor-generated Triton sub-kernels per model (range across the family). Opt: number of models optimized. Fall.: dominant fallback cause within each family.

Figure 5: L3 KernelBench results by architecture family. Each bar shows the outcome (optimized, matched, fallback) for the corresponding model. Models dominated by Conv/GEMM operations tend to fall back because cuBLAS/cuDNN already provides near-optimal performance for these operations; KernelOPT preserves the baseline in all such cases.

Figure 6: Optimization trajectory for L1-007 (matmul). Each point represents a candidate kernel evaluated during beam search. The planner identifies register pressure (184 regs, 12.4% occupancy) as the primary bottleneck from NCU metrics. Iteration 1, beam B applies FP16 dot-product accumulation, reducing registers to 125 and achieving a 3.27 $\times$ sub-kernel speedup (2.89 $\times$ E2E after model-level overhead).

Component	Specification
GPU	NVIDIA H200 SXM (Hopper, CC 9.0)
SMs	132
HBM3	141 GB, 4.8 TB/s bandwidth
L2 Cache	50 MB
CUDA	12.8
PyTorch	2.7.0

Appendix

Table 11: Hardware, software, and hyperparameter configuration used for all experiments.

Explore similar work

AgentKernelArena: Generalization-Aware Benchmarking of GPU Kernel Optimization Agents
May 16, 2026Sharareh Younesian, Wenwen Ouyang, Sina Rafati +11Graphics Processing Unit Kernels Agentic Optimization
AMD
Xe-Forge: Multi-Stage LLM-Powered Kernel Optimization for Intel GPU
Apr 16, 2026Marcin Spoczynski, Daniel Fleischer, Moshe Berchansky +5Graphics Processing Unit Kernels Cuda
Intel Corporation
FastKernels: Benchmarking GPU Kernel Generation in Production
May 22, 2026Gabriele Oliaro, Yichao Fu, May Jiang +5Graphics Processing Unit Kernels
1Snowflake AI Research · 2CMU · 3UCSD +1

May 16, 2026, cs.CLJ/K move · Enter open · S save

AgentKernelArena: Generalization-Aware Benchmarking of GPU Kernel Optimization Agents

Sharareh Younesian, Wenwen Ouyang, Sina Rafati, Mehdi Rezagholizadeh+10

AMD

Graphics Processing Unit Kernels Agentic Optimization Pytorch Chatbot Arena

Open paper AskPDF arXiv

GPU kernel optimization is increasingly critical for efficient deep learning systems, but writing high-performance kernels still requires substantial low-level expertise. Recent AI coding agents can iteratively read code, invoke compilers and profilers, and refine implementations, yet existing kernel benchmarks evaluate single LLM calls rather than full agent workflows, and none include both kernel-to-kernel optimization and unseen-configuration generalization testing. We present AgentKernelArena, an open-source benchmark for measuring AI coding agents on GPU kernel optimization. The benchmark contains 196 tasks spanning HIP-to-HIP optimization, Triton-to-Triton optimization, and PyTorch-to-HIP translation, and evaluates complete agent workflows in isolated workspaces using gated compilation, correctness, and performance checks, centralized scoring and an unseen-configuration generalization protocol that tests whether optimizations transfer to input configurations the agent never observed. Across production agents including Cursor Agent, Claude Code, and Codex Agent, we find near-perfect compilation and high correctness rates on most task categories, with the strongest configurations achieving mean speedups of up to 6.89x on PyTorch-to-HIP, 6.69x on HIP-to-HIP, and 2.13x on Triton-to-Triton tasks. Our unseen-configuration evaluation shows that HIP-to-HIP and Triton-to-Triton optimizations largely transfer to unseen input shapes, while PyTorch-to-HIP exhibits substantial correctness drops, indicating that agents generating kernels from scratch frequently hardcode shape-specific assumptions. AgentKernelArena is designed as a modular, extensible framework for rigorous evaluation of agentic GPU kernel optimization across agents, tasks, and hardware targets.

Abstract

Explore similar work

AgentKernelArena: Generalization-Aware Benchmarking of GPU Kernel Optimization Agents

Xe-Forge: Multi-Stage LLM-Powered Kernel Optimization for Intel GPU

FastKernels: Benchmarking GPU Kernel Generation in Production