KernelOPT improves GPU kernel speed by optimizing compiler outputs
KernelOPT: Dispatch-Aware Agentic Search for GPU Kernel Optimization
Distributed, Parallel, and Cluster ComputingArtificial IntelligenceMachine Learning
Summary
Deep learning run times strongly depend on how efficiently GPU code is written, but automatically generated GPU programs are often slower than expert ones. The authors present KernelOPT, a new tool that safely improves GPU code by focusing on specific parts while keeping trusted library calls intact. It uses multiple AI helpers to propose improvements and checks that the overall program still works well before adopting changes. Tests on many problems show KernelOPT can make GPU code run up to 40% faster than current compiler outputs.
What this means in practice
- •For deep learning engineers: Improve inference and training speed by automatically optimizing GPU kernels in compiled PyTorch models while ensuring correctness.
- •For gpu software developers: Use KernelOPT to refine performance critical GPU kernels after automatic compilation without rewrites of entire models.
Authors
Aheli Poddar, Sanskar Prasad, Arindam Samanta, Subha Chakraborty, Vishal Goyal, Rohit Singh Rathaur
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.