Haciendo que los LLM Optimicen Kernels CUDA Multi-escenario como Expertos
Making LLMs Optimize Multi-Scenario CUDA Kernels Like Experts
March 7, 2026
Autores: Yuxuan Han, Meng-Hao Guo, Zhengning Liu, Wenguang Chen, Shi-Min Hu
cs.AI
Resumen
La optimización manual de kernels de GPU es una tarea compleja y que consume mucho tiempo. Con el rápido desarrollo de los LLM, la optimización automatizada de kernels de GPU se está convirtiendo gradualmente en una realidad tangible. Sin embargo, los métodos actuales de optimización automatizada impulsados por LLM se centran exclusivamente en aplicaciones de aprendizaje automático, como la optimización de operadores de PyTorch, mientras pasan por alto dominios más amplios, como las operaciones con matrices dispersas en la computación científica. Extender estos métodos a aplicaciones más amplias presenta nuevos desafíos para los benchmarks y algoritmos. Por lo tanto, desarrollar un método de optimización de kernels automatizado y de propósito general se convierte en nuestro objetivo principal. En este artículo, abordamos la ausencia de evaluación sistemática para entornos multiescenario mediante la introducción de MSKernelBench, que abarca múltiples escenarios, incluyendo operaciones algebraicas fundamentales, kernels comunes de LLM, operadores de matrices dispersas y rutinas de computación científica, cada uno compatible con precisiones FP32 y BF16. Sobre la base de este benchmark, presentamos CUDAMaster, un sistema multiagente y consciente del hardware para la optimización de kernels, que aprovecha información de profiling y construye automáticamente la cadena completa de herramientas de compilación y ejecución. Los resultados experimentales demuestran que CUDAMaster logra mejoras significativas de velocidad en la mayoría de los operadores, superando a Astra en aproximadamente un 35%. En varios casos, su rendimiento iguala o supera al de bibliotecas altamente optimizadas y de código cerrado, como cuBLAS. Una demostración que muestra el código original y optimizado para cada operador está disponible en https://hanyx2021.github.io/MSKernelBenchDemo/.
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/.