让大模型像专家一样自动优化多场景的CUDA内核。
Making LLMs Optimize Multi-Scenario CUDA Kernels Like Experts
- 构建多智能体系统,结合硬件信息自动编译执行。
- 在多数算子上提速超35%,部分媲美闭源库性能。
- 覆盖科学计算等广泛场景,适合高性能计算研究者。
手动优化GPU内核既困难又耗时。随着大模型发展,自动化内核优化逐渐成为现实。但现有方法局限于机器学习应用(如PyTorch算子),忽视了科学计算中的稀疏矩阵运算等更广领域。为此,我们提出MSKernelBench基准,涵盖基础代数、常见LLM内核、稀疏矩阵操作和科学计算例程,支持FP32与BF16精度。基于此,我们开发了CUDAMaster系统,采用多智能体架构,利用性能分析信息自动构建完整编译与执行链。实验表明,CUDAMaster在多数算子上实现显著加速,较Astra提升约35%;在若干情况下,其性能可匹敌甚至超越cuBLAS等高度优化的闭源库。演示代码见https://hanyx2021.github.io/MSKernelBenchDemo/。
原文摘要 · Abstract (English)
Optimizing GPU kernels manually is a challenging and time-consuming task. With the rapid development of LLMs, automated GPU kernel optimization is gradually becoming a tangible reality. However, current LLM-driven automated optimization methods narrowly focus on machine learning applications, such as PyTorch operator optimization, while overlooking broader domains like sparse matrix operations in scientific computing. Extending to these broader applications brings new challenges for the benchmark and algorithm. Therefore, developing a general-purpose automated kernel optimization method becomes our primary focus. In this paper, we address the absence of systematic evaluation for multi-scenario settings by introducing MSKernelBench, which spans multiple scenarios, including fundamental algebraic operations, common LLM kernels, sparse matrix operators, and scientific computing routines, each supporting both FP32 and BF16 precision. Building on this benchmark, we introduce CUDAMaster, a multi-agent, hardware-aware system for kernel optimization that leverages profiling information and automatically constructs the full compilation and execution toolchain. Experimental results demonstrate that CUDAMaster achieves significant speedups across most operators, outperforming Astra by about 35%. In several cases, its performance matches or surpasses that of highly optimized, closed-source libraries such as cuBLAS. A demo showcasing the original and optimized code for each operator is available at https://hanyx2021.github.io/MSKernelBenchDemo/.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。