arXivDaily arXiv每日学术速递 周一至周五更新
​通知:联邦节假日调休:哥伦布日 / 原住民日(Columbus Day / Indigenous Peoples' Day),今日(13日)无更新。
arXiv 2603.29010cs.LGcs.AI

利用领域特定语言和光速指导提升GPU内核优化代理的效率

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

NVIDIA(英伟达)

Siva Kumar Sastry Hari, Vignesh Balaji, Sana Damani, Qijing Huang, Christos Kozyrakis

AI总结:

本文提出通过领域特定语言和光速指导提升GPU内核优化代理效率,展示通过DSL和SOL指导显著提升优化效率,并发现弱模型在低token成本下可超越强基线代理。

摘要:

使用LLM代理优化GPU内核是一个在大规模设计空间中迭代的过程。每个候选方案必须生成、编译、验证和分析,因此减少试验次数可节省运行时间和成本。我们有两个关键观察:首先,代理所处的抽象级别很重要。如果太低,LLM会浪费在低影响细节上;如果太高,可能会错过重要优化选择。其次,代理难以判断是否达到边际收益递减点,继续搜索会浪费资源。这些观察推动了两个设计原则:(1) 一种紧凑的领域特定语言(DSL),可以在上下文中学习,使模型在更高层次推理的同时保留重要优化杠杆;(2) 光速指导(SOL)利用基本原理性能界限来引导和预算搜索。我们实现了这些原则在μCUTLASS中,这是一个支持CUTLASS后端GPU内核的DSL及其编译器,涵盖内核配置、后缀融合和多阶段流水线。我们使用SOL指导来估计余量并引导优化试验,优先处理接近SOL的问题,并标记可能在基准测试中作弊的内核。在59个KernelBench问题上,使用GPT-5-mini生成DSL代码将几何平均回归从0.40x提升到1.27x,添加SOL指导后达到1.56x。在不同模型层级上,μCUTLASS + SOL指导使较弱模型在较低token成本下超越强基线代理。SOL指导的预算节省19-43%的token,同时保留至少95%的几何平均速度提升,最佳策略达到1.68x的效率提升。最后,SOL分析有助于检测基准测试作弊情况,其中内核可能在表面上快速,但未能完成预期计算。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑