arXiv:2603.29010cs.LGcs.AI2026-03被引 6

用专用语言和即时性能指引,让AI更高效优化GPU代码。

Improving Efficiency of GPU Kernel Optimization Agents using a Domain-Specific Language and Speed-of-Light Guidance

  • 设计轻量级领域语言,让AI在更高层思考优化策略。
  • 引入光速指引,提前识别优化瓶颈,减少无效尝试。
  • 弱模型用此方法可超越强模型,且节省19%-43%计算资源。

使用LLM代理优化GPU内核是一个在巨大设计空间中迭代的过程。每个候选方案都需要生成、编译、验证和性能分析,因此减少尝试次数能显著降低时间和成本。我们提出两个关键观察:第一,代理操作的抽象层级至关重要——过低则浪费推理于无关细节,过高则可能遗漏重要优化选择;第二,代理难以判断是否已到边际效益递减点,导致持续搜索浪费资源。据此提出两项设计原则:(1) 一种紧凑的领域特定语言(DSL),可在上下文中快速学习,使模型在保持关键优化杠杆的前提下进行高层推理;(2) 光速(SOL)指引,利用基本原理性能上限来引导并预算搜索。我们在μCUTLASS中实现这些原则,该工具基于CUTLASS支持的GPU内核,涵盖内核配置、后处理融合及多阶段流水线。SOL指引用于估计性能余量,指导优化试验,降低接近极限问题的优先级,并标记基准测试作弊情况。在59个KernelBench问题上,相同迭代预算下,从生成底层代码转为使用DSL的GPT-5-mini,将几何平均回归效率从0.40x提升至1.27x(相对于PyTorch)。加入SOL引导后进一步提升至1.56x。跨模型层级,μCUTLASS + SOL指引使较弱模型在更低的令牌成本下超越更强基线代理。SOL引导预算可节省19%-43%的令牌,同时保留至少95%的几何平均加速比,最优策略达到1.68x效率增益。此外,SOL分析有助于检测基准测试作弊案例,即内核看似快速但未执行实际计算。

原文摘要 · Abstract (English)

Optimizing GPU kernels with LLM agents is an iterative process over a large design space. Every candidate must be generated, compiled, validated, and profiled, so fewer trials will save both runtime and cost. We make two key observations. First, the abstraction level that agents operate at is important. If it is too low, the LLM wastes reasoning on low-impact details. If it is too high, it may miss important optimization choices. Second, agents cannot easily tell when they reach the point of diminishing returns, wasting resources as they continue searching. These observations motivate two design principles to improve efficiency: (1) a compact domain-specific language (DSL) that can be learned in context and lets the model reason at a higher level while preserving important optimization levers, and (2) Speed-of-Light (SOL) guidance that uses first-principles performance bounds to steer and budget search. We implement these principles in $μ$CUTLASS, a DSL with a compiler for CUTLASS-backed GPU kernels that covers kernel configuration, epilogue fusion, and multi-stage pipelines. We use SOL guidance to estimate headroom and guide optimization trials, deprioritize problems that are near SOL, and flag kernels that game the benchmark. On 59 KernelBench problems with the same iteration budgets, switching from generating low-level code to DSL code using GPT-5-mini turns a 0.40x geomean regression into a 1.27x speedup over PyTorch. Adding SOL-guided steering raises this to 1.56x. Across model tiers, $μ$CUTLASS + SOL-guidance lets weaker models outperform stronger baseline agents at lower token cost. SOL-guided budgeting saves 19-43% of tokens while retaining at least 95% of geomean speedup, with the best policy reaching a 1.68x efficiency gain. Lastly, SOL analysis helps detect benchmark-gaming cases, where kernels may appear fast while failing to perform the intended computation.

GPU优化LLM代理领域语言

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