arXiv:2511.18868cs.LGcs.AI2025-11被引 2

用强化学习方法自动优化GPU核函数,提升大模型推理速度。

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

  • 将内核优化建模为多臂老虎机问题,平衡探索与利用。
  • 在三种GPU上平均提速超33%,优于现有方法。
  • 适合需要高效部署大模型的系统工程师使用。

高性能GPU内核对大模型服务至关重要,但其优化仍需深厚系统知识,是主要瓶颈。尽管代码大模型能生成功能正确的代码,但内核优化本质上是在巨大优化空间中的搜索问题,现有模型难以高效探索。为此,我们提出KernelBand,将内核优化建模为多臂老虎机(MAB)问题,显式平衡探索与利用,释放代码大模型潜力。为应对候选内核优化策略的无限臂空间,设计了两种机制:基于性能分析边界的硬件感知剪枝策略,以及利用Lipschitz连续性的轨迹驱动聚类算法。理论上证明,KernelBand的损失界依赖于运行时簇的紧凑覆盖数,保证样本高效发现高性能内核。在TritonBench-G数据集上,针对三种GPU架构和四种代码大模型的实验表明,KernelBand consistently and substantially outperforms state-of-the-art methods,平均提速超过33%。

原文摘要 · Abstract (English)

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.

LLM优化GPU加速强化学习内核生成

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。