用人类专家经验指导大模型自动优化CUDA内核,兼顾速度与能效。
Optimizing CUDA like a Human: Micro-Profiling Tools as Expert Surrogates for LLM-Based GPU Kernel Optimization
- 融合大模型生成与硬件剖析工具,通过可插拔微剖析器提供可操作建议。
- 在不同难度基准上实现最高2.42倍至5.30倍加速,优于手工调优。
- 首个同时优化性能与能效的自动生成内核系统,适合深度优化开发者使用。
我们提出KernelPro,一个闭环多智能体系统,通过整合大语言模型(LLM)代码生成、硬件剖析反馈及可插拔瓶颈检测工具,自动生成、剖析并迭代优化GPU内核代码。KernelPro包含四项贡献:(1) 语义反馈算子,将专家启发式规则编码为可插拔微剖析工具,把原始硬件指标转化为可操作的自然语言指导;(2) 两阶段工具调用架构,基于屋顶图的瓶颈分类筛选执行专用分析工具,融合内核级(ncu)、指令级(SASS)和系统级(nsys)剖析;(3) 领域适配的蒙特卡洛树搜索(MCTS),含渐进式扩展、非对称分支、日志奖励校准、死端剪枝和跨迭代学习记忆;(4) 通过自主代码搜索在CUTLASS/CuTe库中直接生成源码级CuTe代码。在KernelBench上,KernelPro在级别1/2/3分别实现2.42x/4.69x/5.30x几何平均加速,覆盖所有难度等级的最先进表现。在VeOmni专家优化的MoE训练内核上,相比手工调优的Triton代码,提升1.23倍,生成了从零开始的原始CUDA+CuTe Hopper WGMMA内核。消融实验表明:每个组件独立显著提升优化质量(微剖析工具,p<0.0001;MCTS搜索,比贪心策略高26%,p=0.004;主动工具调度,提升23%,p=0.035)。此外,KernelPro是首个在保持速度前提下优化能效的自动生成内核代理,实测在相同速度下能耗降低11.6%。
原文摘要 · Abstract (English)
We present KernelPro, a closed-loop multi-agent system that automatically generates, profiles, and iteratively optimizes GPU kernel code by integrating large language model (LLM) code generation with hardware profiler feedback and pluggable bottleneck detection tools. KernelPro introduces four contributions: (1) a semantic feedback operator that encodes expert heuristics as pluggable micro-profiling tools, transforming raw hardware metrics into actionable natural language guidance; (2) a two-stage tool invocation architecture where roofline-based bottleneck classification filters which specialized analysis tools execute, combining kernel-level (ncu), instruction-level (SASS), and system-level (nsys) profiling; (3) a domain-adapted MCTS with progressive widening, asymmetric branching, log-reward calibration, dead-end pruning, and search memory for cross-iteration learning; and (4) direct CuTe source-level code generation via autonomous code search over the CUTLASS/CuTe codebase. On KernelBench, KernelPro achieves geometric mean speedups of 2.42x/4.69x/5.30x on Levels 1/2/3, establishing state-of-the-art performance across all difficulty levels. On VeOmni's expert-optimized MoE training kernels, KernelPro achieves 1.23x over hand-tuned Triton by generating a from-scratch raw-CUDA+CuTe Hopper WGMMA kernel. Ablation studies demonstrate that each design component independently and significantly improves optimization quality: micro-profiling tools (p < 0.0001 vs raw metrics), MCTS search (26% higher geometric mean vs greedy, p = 0.004), and proactive tool orchestration (23% improvement, p = 0.035). Finally, KernelPro is the first CUDA kernel coding agent to optimize energy efficiency beyond the speed-only focus of prior systems, demonstrating an 11.6% measured energy reduction at matched speed.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。