用AI自动优化CUDA代码性能,让编程更高效。
OptiML: An End-to-End Framework for Program Synthesis and CUDA Kernel Optimization
- 将代码优化转化为带验证的搜索问题,分两阶段完成。
- 在多个测试中超越主流大模型,提升运行速度并减少瓶颈。
- 适合需要高性能计算的开发者与研究者使用。
生成高性能CUDA内核仍具挑战性,因需在大量低级变换组合中探索,并依赖昂贵且噪声较大的硬件反馈。尽管大语言模型可生成功能正确的CUDA代码,但要达到优异性能仍需系统性地探索和验证优化选择。我们提出OptiML,一个端到端框架,将自然语言意图或输入的CUDA代码转化为性能优化的CUDA内核,将内核优化建模为带验证的搜索过程。OptiML包含两个解耦阶段:当输入为自然语言时,基于思维混合的生成器(OptiML-G)作为策略,生成初始可执行程序;随后,基于蒙特卡洛树搜索的优化器(OptiML-X)通过LLM驱动的修改,对合成或用户提供的内核进行迭代优化,由来自Nsight Compute的硬件感知奖励引导。每个候选变换均被编译、验证并使用Nsight Compute分析,评估指标结合运行时间、硬件瓶颈代理和防止性能退化的约束。我们在多样化的CUDA内核上评估了OptiML在合成-优化与仅优化两种场景下的表现。结果表明,OptiML持续发现经过验证的性能改进,优于强基线大模型,并产生基于剖析器证据的可解释优化轨迹。
原文摘要 · Abstract (English)
Generating high-performance CUDA kernels remains challenging due to the need to navigate a combinatorial space of low-level transformations under noisy and expensive hardware feedback. Although large language models can synthesize functionally correct CUDA code, achieving competitive performance requires systematic exploration and verification of optimization choices. We present OptiML, an end-to-end framework that maps either natural-language intent or input CUDA code to performance-optimized CUDA kernels by formulating kernel optimization as search under verification. OptiML consists of two decoupled stages. When the input is natural language, a Mixture-of-Thoughts generator (OptiML-G) acts as a proposal policy over kernel implementation strategies, producing an initial executable program. A search-based optimizer (OptiML-X) then refines either synthesized or user-provided kernels using Monte Carlo Tree Search over LLM-driven edits, guided by a hardware-aware reward derived from profiler feedback. Each candidate transformation is compiled, verified, and profiled with Nsight Compute, and evaluated by a composite objective that combines runtime with hardware bottleneck proxies and guardrails against regressions. We evaluate OptiML in both synthesis-and-optimize and optimization-only settings on a diverse suite of CUDA kernels. Results show that OptiML consistently discovers verified performance improvements over strong LLM baselines and produces interpretable optimization trajectories grounded in profiler evidence.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。