KernelOPT: Dispatch-Aware Agentic Search for GPU Kernel Optimization
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 (Level 1: 51/100), 1.15 (Level 2: 31/100), and 1.07 (Level 3: 12/50) across all problems.
Figures & tables
| 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 | 88.63 | 10 | 4 | 1 | 0 | 5 |
| L2 | 100 | 31 | 35 | 0 | 34 | 1.15 | 33.11 | 6 | 4 | 0 | 10 | |
| L3 | 50 | 12 | 3 | 0 | 35 | 1.07 | 4.12 | 7 | 0 | 2 | 9 | |
| 2 | 2 | 1 | 5 | |||||||||
| 8 | 8 | 4 | 20 | |||||||||
| Config | Opt | Match | Fall | Opt% |
|---|---|---|---|---|
| Full system | 19 | 14 | 17 | 38% |
| No beam ( ) | 5 | 3 | 41 | 10% |
| No NCU profiling | 5 | 4 | 41 | 10% |
| No memory | 4 | 3 | 43 | 8% |
| No multi-iter ( ) | 4 | 4 | 42 | 8% |
| No perf gate | 3 | 3 | 44 | 6% |
Appendix figures & tables14 assets
Supplementary material from the paper’s appendix.
Appendix
| Stage | Gate | Tolerance | Seeds | Reference | Scope |
|---|---|---|---|---|---|
| Synthesis | Eager validation | 2: | 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. | 3: | Synthesized Triton | Per candidate | |
| Gate 3 | Model-level (Tier 1) | 3: | Re-stitched model | End-to-end | |
| Model-level (Tier 2) | or abs bound | Float64 reference | End-to-end | ||
| Gate 4 | Performance gate | (3%) | N/A | do_bench (25/100 ms) | End-to-end |
| 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 |
| 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 |
| 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 |
| C1 (no beam) | 4 | 3 | 13 | 1 | 0 | 19 | 0 | 0 | 10 | 1.78 |
| C2 (no memory) | 4 | 3 | 13 | 0 | 0 | 20 | 0 | 0 | 10 | 1.23 |
| C3 (no perf gate) | 3 | 3 | 14 | 0 | 0 | 20 | 0 | 0 | 10 | 1.14 |
| C4 (no NCU) | 3 | 4 | 13 | 2 | 0 | 18 | 0 | 0 | 10 | 3.16 |
| Kernel | Comp (%) | Mem (%) | Occ (%) | Planner direction (iter.) | Sub- | E2E | |
|---|---|---|---|---|---|---|---|
| L1 | matmul_kernel (007) Regs/thread: 184 | 57 | 47 | 12.4 | Iter 1: fast_accum_fp16_dot ; regs 184 125, occ. 12 24% | 2.33 | 2.89 |
| L2 | fused_linear_mul_ hardtanh_gelu (053) Regs/thread: 255 (max) | 37 | 59 | 12.5 | Iter 1: reduce_regs (1.0 ); Iter 3: tf32x3_tensor_cores (3.12 ) | 3.12 | 1.60 |
| L3 | gemm_bias_relu (002, MLP) Regs/thread: 198 | 60 | 67 | 12.2 | Iter 1: reduce_reg_pressure (1.02 ); Iter 3: k_loop_peeling (1.12 ) | 1.12 | 4.12 |
| Kernel | Variant | Comp % | Mem % | Regs | Occ % | Duration | Sub- |
| L1-007 (matmul) | Baseline | 57.1 | 46.7 | 184 | 12.4 | 5.95 ms | — |
| E2E 2.89 | Optimized (fp16 dot) | 33.8 | 75.4 | 125 | 24.1 | 1.82 ms | 3.27 |
| L1-012 (diag matmul) | Baseline (synth.) | 14.0 | 69.0 | 18 | 81.2 | 32.83 s | — |
| E2E 88.63 ; speedup from algorithmic rewrite ( ), not NCU-guided micro-optimization. | |||||||
| L2-018 (fused chain) | Baseline (synth.) | 21.8 | 56.1 | 29 | 45.9 | 12.45 s | — |
| E2E 33.11 ; speedup from fusing 5 operations with algebraic simplification. | |||||||
| Lvl | K# | Speed | Cat | Ops | Rsn | Lvl | K# | Speed | Cat | Ops | Rsn |
|---|---|---|---|---|---|---|---|---|---|---|---|
| L1 | 001 | 5.77 | OPT | Matmul | TO | L1 | 002 | 6.06 | OPT | Matmul | |
| L1 | 003 | 5.11 | OPT | Matmul | L1 | 004 | 1.30 | OPT | Matmul | ||
| L1 | 005 | 1.01 | OPT | L1 | 006 | 0.98 | MAT | Matmul | |||
| L1 | 007 | 2.89 | OPT | Matmul | L1 | 008 | 2.63 | OPT | Matmul | TO | |
| L1 | 009 | 2.07 | OPT | Matmul | L1 | 010 | 5.79 | OPT | Matmul | ||
| L1 | 011 | 1.05 | OPT | Einsum | L1 | 012 | 88.63 | OPT | Matmul |
| Architecture | Sub- | Opt | Speedup | Fall. | |
| MLP | 3 | 3–4 | 1 | 4.12 | GEMM |
| ResNet | 3 | 8–52 | 1 | 1.55 | Conv |
| DenseNet | 4 | 3–337 | 2 | 1.06–1.15 | Conv |
| MobileNet/EffNet | 6 | 7–60 | 2 | 1.03–2.33 | Conv |
| ViT / Swin / CViT | 5 | 1–99 | 2 | 1.10–1.11 | GEMM |
| U-Net/NetVLAD/SSM | 5 | 1–16 | 4 | 1.06–1.17 | Perf |
| 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 |