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

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. This 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 runtime-behavior clustering algorithm that leverages Lipschitz continuity. Extensive experiments on TritonBench-G with three GPU architectures and four code LLMs show that KernelBand consistently outperforms the strongest available agent baseline, achieving up to 1.91× geometric mean speedup over correctly optimized kernels with 39–140% relative improvement in Fast@1 success rate. Our code is available at https://github.com/TongmingLAIC/KernelBand.