KernelOPT: Dispatch-Aware Agentic Search for GPU Kernel Optimization

Problem

This paper addresses the gap in GPU kernel efficiency between modern compilers and expert-written implementations. The authors highlight that existing compilation techniques often fail to achieve the performance levels of hand-optimized kernels, necessitating a more effective optimization approach. The work is presented as a preprint and has not undergone peer review.

Method

The proposed system, KernelOPT, employs a multi-agent architecture consisting of five profiling-guided large language model (LLM) agents. Key components of KernelOPT include:

• Preservation of vendor library calls, specifically targeting cuBLAS and cuDNN.
• Optimization of generated Triton sub-kernels.
• A four-gate verification cascade that includes:
  • Static validation to ensure syntactic correctness.
  • Multi-seed correctness checks to validate output consistency across different random seeds.
  • Model-level float64-fallback verification to ensure numerical stability.
  • Performance gating to filter out suboptimal solutions.

KernelOPT accepts various input types, including PyTorch nn.Modules, standalone Triton kernels, and Helion kernels. The evaluation dataset comprises 250 KernelBench problems, which serve as benchmarks for assessing the optimization performance.

Results

KernelOPT demonstrates notable performance improvements over the existing torch.compile method across multiple levels of optimization:

• A geometric mean speedup of 1.40× at Level 1 (51/100).
• A geometric mean speedup of 1.15× at Level 2 (31/100).
• A geometric mean speedup of 1.07× at Level 3 (12/50). These results indicate that KernelOPT significantly enhances kernel execution speed compared to the baseline compiler, showcasing its effectiveness in optimizing GPU kernels.

Limitations

The authors do not report any limitations in their work. However, the absence of reported limitations may suggest a lack of comprehensive evaluation across diverse GPU architectures or workloads, which could affect generalizability.

Why it matters

The implications of this work are substantial for the field of GPU programming and optimization. By leveraging a multi-agent system for kernel optimization, KernelOPT could facilitate the development of more efficient GPU applications, potentially leading to performance gains in various domains such as deep learning and scientific computing. This approach may inspire further research into agent-based optimization techniques and their application in other areas of machine learning and high-performance computing.