用对比强化学习自动优化CUDA代码,速度提升最高达120倍。
CUDA-L1: Improving CUDA Optimization via Contrastive Reinforcement Learning
- 基于对比强化学习,通过加速奖励信号自动优化CUDA代码。
- 在250个内核上平均提速3.12倍,峰值达120倍,超越多种现有工具。
- 能发现隐藏瓶颈,识别无效优化,适合需要高效GPU计算的开发者。
GPU计算资源需求激增,亟需自动化CUDA优化策略。尽管大模型在代码生成方面展现潜力,但当前最优模型在提升CUDA性能方面成功率仍低。本文提出CUDA-L1,一种基于新型对比强化学习的自动化CUDA优化框架。该框架在A100上训练,对KernelBench中的250个CUDA内核实现平均3.12倍、中位1.42倍的加速,峰值达120倍。相比默认基线,其性能优于Torch Compile(x2.77)、Torch Compile带减少开销(x2.88)、CUDA Graph(x2.81)及cuDNN库(x7.72)。此外,模型具备跨GPU架构的可移植性。CUDA-L1不仅能发现多种优化技术并策略性组合,揭示优化的乘法特性,还能识别非明显性能瓶颈,并拒绝实际有害的优化建议。结果表明,仅通过加速奖励信号,强化学习即可将初始表现差的大模型转化为有效优化器,无需人工经验,为自动化优化开辟新路径,显著提升GPU效率,缓解计算资源压力。
原文摘要 · Abstract (English)
The exponential growth in demand for GPU computing resources has created an urgent need for automated CUDA optimization strategies. While recent advances in LLMs show promise for code generation, current SOTA models achieve low success rates in improving CUDA speed. In this paper, we introduce CUDA-L1, an automated reinforcement learning framework for CUDA optimization that employs a novel contrastive RL algorithm. CUDA-L1 achieves significant performance improvements on the CUDA optimization task: trained on A100, it delivers an average speedup of x3.12 with a median speedup of x1.42 against default baselines over across all 250 CUDA kernels of KernelBench, with peak speedups reaching x120. In addition to the default baseline provided by KernelBench, CUDA-L1 demonstrates x2.77 over Torch Compile, x2.88 over Torch Compile with reduce overhead, x2.81 over CUDA Graph implementations, and x7.72 over cuDNN libraries. Furthermore, the model also demonstrates portability across different GPU architectures. Beyond these benchmark results, CUDA-L1 demonstrates several properties: it 1) discovers a variety of CUDA optimization techniques and learns to combine them strategically to achieve optimal performance; 2) uncovers fundamental principles of CUDA optimization, such as the multiplicative nature of optimizations; 3) identifies non-obvious performance bottlenecks and rejects seemingly beneficial optimizations that actually harm performance. The capabilities demonstrate that, RL can transform an initially poor-performing LLM into an effective CUDA optimizer through speedup-based reward signals alone, without human expertise or domain knowledge. This paradigm opens possibilities for automated optimization of CUDA operations, and holds promise to substantially promote GPU efficiency and alleviate the rising pressure on GPU computing resources.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。