Making LLMs Optimize Multi-Scenario CUDA Kernels Like Experts

TL;DR

This paper addresses 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.

Abstract

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/.

Figures (9)

• Figure 1: Overview of MSKernelBench and CUDAMaster.Left: MSKernelBench is a comprehensive multi-scenario CUDA optimization benchmark covering dense, sparse, LLM, and scientific kernels, etc. Top right: CUDAMaster classifies tasks according to performance bottlenecks, filters profile information in a targeted manner, and uses a multi-agent system based on various tools to achieve CUDA kernel optimization. Bottom right: On a macro level, the framework successfully balances performance improvements with operator correctness, increasing performance demands across all tested scenarios. On a micro level, the framework can generate kernels surpassing hand-tuned implementations like closed-source libraries across multiple domains.
• Figure 2: Composition of the 50 Operators by Task Category.
• Figure 3: CUDAMaster Framework for Automated Kernel Optimization Pipeline.
• Figure 4: Cumulative success rate of all tasks versus speedup threshold($\tau$) for OpenAI o4-mini and DeepSeek-V3.2 models under FP32 and BF16 precision.
• Figure 5: Performance comparison of our agent-optimized kernels against closed-source libraries (cuBLAS, cuDNN, cuSPARSE) and existing optimization frameworks (Astra) across six representative operators and varying problem scales. Subfigures (a), (b), and (c) show performance advantages over closed-source libraries in three key operations: (a) sparse matrix-vector multiplication (SpMV) over cuSPARSE, (b) 2D convolution over cuDNN, and (c) dot product over cuBLAS. Subfigures (d), (e), and (f) show competitive performance against Astra on LLM operators (RMSNorm, SiLU&Mul and Merge Attention States).