KernelBand: Steering LLM-based Kernel Optimization via Hardware-Aware Multi-Armed Bandits

TL;DR

KernelBand introduces a hardware-aware multi-armed bandit framework for optimizing GPU kernels, effectively balancing exploration and exploitation to improve performance in LLM serving tasks across diverse hardware.

Contribution

It formulates kernel optimization as a multi-armed bandit problem and designs novel pruning and clustering mechanisms for efficient exploration.

Findings

Achieves over 33% average performance improvement.

Outperforms state-of-the-art methods across multiple GPU architectures.

Provides theoretical guarantees on sample efficiency.

Abstract

High-performance GPU kernels are critical for efficient LLM serving, yet their optimization remains a bottleneck requiring deep system expertise. While code LLMs show promise in generating functionally correct code, kernel optimization is intrinsically a search problem over a vast optimization space. The fundamental mismatch prevents existing LLM agents from efficiently exploring the optimization space for diverse hardware and compute patterns. To bridge the gap, we present KernelBand, a framework that formulates kernel optimization as a Multi-Armed Bandit (MAB) problem, explicitly balancing exploration and exploitation to unlock the potential of code LLMs. To navigate the infinite arm space of optimization strategies applied to candidate kernels, we design two key mechanisms: a hardware-aware pruning strategy via profiling bounds and a trace-driven clustering algorithm that leverages Lipschitz continuity. Theoretically, we prove that KernelBand reduces the regret bound to depend on the compact covering number of runtime clusters, ensuring sample-efficient discovery of high-performance kernels. Extensive experiments on TritonBench-G with three GPU architectures and four code LLMs show that KernelBand consistently and substantially outperforms state-of-the-art methods with over 33% average improvement.