用强化学习自动优化矩阵乘法核，性能超越cuBLAS
CUDA-L2: Surpassing cuBLAS Performance for Matrix Multiplication through Reinforcement Learning
- 用大模型+强化学习搜索1000种配置，自动调优半精度矩阵乘法
- 离线模式下比torch.matmul快22.0%，比cuBLAS快19.2%
- 适合追求极致算子性能的深度学习系统开发者
本文提出CUDA-L2，一个结合大语言模型与强化学习的系统，用于自动优化半精度通用矩阵乘法（HGEMM）的CUDA内核。以CUDA执行速度为强化学习奖励，CUDA-L2在1000种配置中自动优化HGEMM内核。该系统在离线模式下平均比torch.matmul快22.0%，比cuBLAS快19.2%（使用最优布局），比cuBLASLt-heuristic快16.8%，比cuBLASLt-AutoTuning快11.4%；在模拟实时推理的服务器模式下，提速进一步提升至28.7%、26.0%、22.4%和15.9%。结果表明，即使高度优化的HGEMM内核也可通过大模型引导的强化学习自动化实现显著性能提升，系统性探索人类难以触及的配置空间。项目与代码见github.com/ornith-ai/CUDA-L2。
原文摘要 · Abstract (English)
In this paper, we propose CUDA-L2, a system that combines large language models (LLMs) and reinforcement learning (RL) to automatically optimize Half-precision General Matrix Multiply (HGEMM) CUDA kernels. Using CUDA execution speed as the RL reward, CUDA-L2 automatically optimizes HGEMM kernels across 1,000 configurations. CUDA-L2 systematically outperforms major matmul baselines to date, from the widely-used torch.matmul to state-of-the-art Nvidia's closed-source libraries, i.e., cuBLAS, cuBLASLt. In offline mode, where kernels are executed consecutively without time intervals, CUDA-L2 yields +22.0% over torch.matmul on average; +19.2% over cuBLAS using the optimal layout configuration (normal-normal NN and transposed-normal TN); +16.8% over cuBLASLt-heuristic, which queries cuBLASLt library and selects the algorithm based on the heuristic's suggestion; and +11.4% over the most competitive cuBLASLt-AutoTuning model, which selects the fastest algorithm from up to 100 candidates from cuBLASLt's suggestions. In server mode, where kernels are executed at random intervals simulating real-time inference, the speedups further increase to +28.7%, +26.0%, +22.4%, and +15.9% for torch.matmul, cuBLAS, cuBLASLt-heuristic, and cuBLASLt-AutoTuning respectively. CUDA-L2 shows that even the most performance-critical, heavily-optimized kernels like HGEMM can be improved through LLM-guided RL automation by systematically exploring configuration spaces at scales impractical for humans. Project and code can be found at github.com/ornith-ai/CUDA-L2
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。