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

arXiv: https://arxiv.org/abs/2511.18868 · HTML: https://mlsanity.com/paper/2511.18868/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Serving systems and request scheduling, one of 12 directions in Efficient LLM inference and serving.
- Builds on 32 of its 70 references that are in the map.
- First posted on arXiv 24 Nov 2025.

## Placed in

- Efficient LLM inference and serving / [Serving systems and request scheduling](https://mlsanity.com/inference/serving-systems-and-request-scheduling/)

## Builds on (most cited first, 30 of 32)

- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Qwen Technical Report](https://mlsanity.com/paper/2309.16609/) (2023)
- [A Survey of Large Language Models](https://mlsanity.com/paper/2303.18223/) (2023)
- [Qwen2.5-Coder Technical Report](https://mlsanity.com/paper/2409.12186/) (2024)
- [DeepSeek-Coder: When the Large Language Model Meets Programming -- The Rise of Code Intelligence](https://mlsanity.com/paper/2401.14196/) (2024)
- [A Comprehensive Overview of Large Language Models](https://mlsanity.com/paper/2307.06435/) (2023)
- [Kimi K2: Open Agentic Intelligence](https://mlsanity.com/paper/2507.20534/) (2025)
- [Introduction to Multi-Armed Bandits](https://mlsanity.com/paper/1904.07272/) (2019)
- [FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving](https://mlsanity.com/paper/2501.01005/) (2025)
- [Ansor: Generating High-Performance Tensor Programs for Deep Learning](https://mlsanity.com/paper/2006.06762/) (2020)
- [Flex Attention: A Programming Model for Generating Optimized Attention Kernels](https://mlsanity.com/paper/2412.05496/) (2024)
- [Learning to Optimize Tensor Programs](https://mlsanity.com/paper/1805.08166/) (2018)
- [A Survey on Code Generation with LLM-based Agents](https://mlsanity.com/paper/2508.00083/) (2025)
- [Online Clustering of Bandits](https://mlsanity.com/paper/1401.8257/) (2014)
- [Towards Efficient Generative Large Language Model Serving: A Survey from Algorithms to Systems](https://mlsanity.com/paper/2312.15234/) (2023)
- [Kevin: Multi-Turn RL for Generating CUDA Kernels](https://mlsanity.com/paper/2507.11948/) (2025)
- [TritonBench: Benchmarking Large Language Model Capabilities for Generating Triton Operators](https://mlsanity.com/paper/2502.14752/) (2025)
- [Infinite-LLM: Efficient LLM Service for Long Context with DistAttention and Distributed KVCache](https://mlsanity.com/paper/2401.02669/) (2024)
- [TurboTransformers: An Efficient GPU Serving System For Transformer Models](https://mlsanity.com/paper/2010.05680/) (2020)
- [Algorithms for multi-armed bandit problems](https://mlsanity.com/paper/1402.6028/) (2014)
- [Towards Robust Agentic CUDA Kernel Benchmarking, Verification, and Optimization](https://mlsanity.com/paper/2509.14279/) (2025)
- [CudaForge: An Agent Framework with Hardware Feedback for CUDA Kernel Optimization](https://mlsanity.com/paper/2511.01884/) (2025)
- [TileLang: A Composable Tiled Programming Model for AI Systems](https://mlsanity.com/paper/2504.17577/) (2025)
- [Geak: Introducing Triton Kernel AI Agent & Evaluation Benchmarks](https://mlsanity.com/paper/2507.23194/) (2025)
- [TritonRL: Training LLMs to Think and Code Triton Without Cheating](https://mlsanity.com/paper/2510.17891/) (2025)
- [A survey on multi-player bandits](https://mlsanity.com/paper/2211.16275/) (2022)
- [ConCuR: Conciseness Makes State-of-the-Art Kernel Generation](https://mlsanity.com/paper/2510.07356/) (2025)
- [SwizzlePerf: Hardware-Aware LLMs for GPU Kernel Performance Optimization](https://mlsanity.com/paper/2508.20258/) (2025)
- [Optimizing CUDA Code By Kernel Fusion---Application on BLAS](https://mlsanity.com/paper/1305.1183/) (2013)

## Built on by (most cited first, 2 of 2)

- [The Correctness Illusion in LLM-Generated GPU Kernels](https://mlsanity.com/paper/2606.20128/) (2026)
- [Test-Input Generation for Tensor Programs: What Actually Finds Kernel Bugs](https://mlsanity.com/paper/2606.27396/) (2026)

## Abstract (arXiv)

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.
