TLAIC Logo

    KernelBand：通过硬件感知多臂老虎机引导LLM核优化

    1. 引言

           大规模部署大型语言模型（LLM）需要针对注意力机制和矩阵乘法等基本操作进行高度优化的GPU内核。尽管像Triton这样的领域特定语言简化了内核开发，但由于配置空间巨大以及需要专业的硬件知识，实现峰值性能仍然具有挑战性。最近的方法尝试利用代码LLM进行内核优化，但这些方法面临一个根本性挑战：LLM擅长生成功能正确的代码，而内核优化本质上是一个搜索问题，需要仔细探索不连续的性能景观。

    KERNELBAND通过将内核优化表述为多臂赌博机（MAB）问题来解决这种不匹配，提供了一个原则性框架，在使用LLM进行性能优化时平衡探索和利用。该方法引入了硬件感知剪枝和迹驱动聚类来管理实际上无限的优化策略空间，从而能够以理论保证高效地发现高性能内核。

    2. 问题表述与理论基础

           KERNELBAND将内核优化形式化为找到一个最优内核 k*，该内核在保持功能正确性的同时，最大限度地减少硬件平台 H 的延迟 L(k, H)。优化过程被建模为遍历一个有向图，其中节点代表有效的内核，边对应于通过代码LLM应用优化策略。

    其核心创新在于将其表述为一个具有不断扩展的动作空间的迭代上下文赌博机问题。在每个时间步 t，该框架操作于：

    该框架依赖于两个关键的结构假设。首先，**硬件感知收益有界性（Hardware-Aware Gain Boundedness）**指出，任何策略的预期改进都受硬件限制，这使得能够剪枝那些因资源饱和而无法产生显著收益的策略。其次，**行为空间中的利普希茨连续性（Lipschitz Continuity in Behavior Space）**假设具有相似运行时特征的内核将对相同的优化做出相似的响应，从而允许在具有可比执行模式的内核之间进行泛化。

    3. 系统架构和核心组件

    KERNELBAND的架构集成了三个主要组件，以解决LLM引导的内核优化挑战。

    运行时行为特征化

    每个内核通过两种互补的表示进行特征化：

    4. 通过动态聚类进行结构化探索

           为了管理不断扩展的动作空间，KERNELBAND采用动态聚类内核前沿 Pt。每隔 τ 次迭代，内核使用K-Means在其行为特征向量上被划分为 K 个簇。这种方法提供了几个优势：

    • 泛化能力：利用Lipschitz连续性假设，根据聚类代表估计未探索的内核-策略组合的奖励。
    • 效率：通过仅使用聚类质心来描述硬件约束，减少了性能分析开销。
    • 适应性：周期性地重新聚类以适应随着新内核发现而演变的优化前沿。

    5. 硬件约束的多臂老虎机策略

    决策过程实现了一种掩码多臂老虎机方法，其中策略选择受硬件可行性约束：

    6. 理论分析与保证

    7. 实验评估与结果

           评估是在TritonBench-G上进行的，该基准包含了13个功能类别中的183个GPU Triton内核，使用了三种NVIDIA架构（RTX 4090、H20、A100）和多个代码大语言模型（LLMs）。

    a. 性能优势

           KERNELBAND 在所有测试配置中均持续实现了卓越的性能。在A100上，它实现了1.91倍的几何平均加速比和79.8%的正确率，在加速比上比GEAK基线高出42.5%，在成功率上高出66.2%。在其他平台上也观察到了类似的优势，其加速比在1.45倍到1.91倍之间，而基线为1.06倍到1.44倍。

    b. 硬件适应性

           该框架通过平台特定的策略选择模式展示了有效的硬件适应性。在RTX 4090上，KERNELBAND 为融合策略分配了更多的探索预算（18.5%），而H20则为12.8%；同时，H20 对平铺优化表现出更高的关注（10.0% vs. 7.6%）。这种差异反映了硬件感知剪枝机制根据设备特定瓶颈校准探索的能力。

    c. 扩展性和鲁棒性

           KERNELBAND 展示了卓越的扩展性行为，在增加迭代预算的情况下性能持续提升，而基线则更早达到饱和。该框架在不同的LLM后端（DeepSeek-V3.2、GPT-5、Claude Opus 4.5、Gemini 3 Flash）上均保持了鲁棒的性能，表明其作为代码生成模型的通用放大器是有效的。

    8. 意义与贡献

           KERNELBAND 在LLMs和系统优化交叉领域做出了几项关键贡献。将内核优化表述为硬件约束的多臂老虎机（MAB）问题，提供了一种原则性方法来解决LLM生成能力与性能搜索要求之间的根本不匹配。引入硬件感知剪枝和轨迹驱动聚类，提供了有效探索广阔优化空间的实用机制。

           理论分析为收敛性提供了正式保证，确立了该框架能够实现接近最优的解决方案，其样本复杂度与问题的内在维度成比例，而非整个配置空间。全面的实验评估表明，该框架在不同的硬件平台和LLM后端上均表现出持续改进，这预示着其广泛的适用性。

           这项工作对高性能内核开发的普及具有重要意义，有可能降低GPU优化的入门门槛，并提高LLM服务基础设施的效率。将LLM原则性地整合到复杂的搜索问题中，为将类似方法推广到其他需要高效探索大型配置空间的领域奠定了基础。