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

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.

Modern AI assistants such as ChatGPT require enormous amounts of computation, most of which runs on specialized chips called GPUs. The speed and cost of these services depend heavily on small pieces of low-level code, called kernels, that tell the chip how to carry out mathematical operations efficiently. Writing fast kernels usually requires deep hardware expertise, so it remains a major bottleneck for deploying large AI systems. Recent code-writing AI models can generate working programs, but making a program truly fast is much harder than making it correct. There are many possible ways to rewrite and tune a kernel, and blindly trying them wastes time and computing resources. We present KernelBand, a system that helps AI models search for faster kernels more strategically. Instead of treating every possible code variant independently, KernelBand uses measurements from real hardware to discard unpromising options early and groups similar variants so that useful lessons from one attempt can guide nearby attempts. Across three GPU families and four code-writing AI models, KernelBand generates kernels up to 1.91 times faster than the strongest prior automated approach and succeeds much more often. By making kernel optimization more automatic and efficient, this work can help reduce the cost and energy use of large-scale AI services.