Faire en sorte que les LLM optimisent les noyaux CUDA multi-scénarios comme des experts

Making LLMs Optimize Multi-Scenario CUDA Kernels Like Experts

March 7, 2026

Auteurs: Yuxuan Han, Meng-Hao Guo, Zhengning Liu, Wenguang Chen, Shi-Min Hu

cs.AI

Résumé

L'optimisation manuelle des noyaux GPU est une tâche complexe et chronophage. Avec le développement rapide des LLM, l'optimisation automatisée des noyaux GPU devient progressivement une réalité tangible. Cependant, les méthodes d'optimisation automatisée actuelles pilotées par les LLM se concentrent étroitement sur les applications de l'apprentissage automatique, telles que l'optimisation des opérateurs PyTorch, tout en négligeant des domaines plus vastes comme les opérations sur matrices creuses dans le calcul scientifique. L'extension à ces applications élargies présente de nouveaux défis pour les benchmarks et les algorithmes. Par conséquent, le développement d'une méthode d'optimisation automatisée des noyaux à usage général devient notre objectif principal. Dans cet article, nous répondons à l'absence d'évaluation systématique pour les configurations multi-scénarios en introduisant MSKernelBench, qui couvre plusieurs scénarios, incluant les opérations algébriques fondamentales, les noyaux LLM courants, les opérateurs sur matrices creuses et les routines de calcul scientifique, chacun prenant en charge les précisions FP32 et BF16. Sur la base de ce benchmark, nous présentons CUDAMaster, un système multi-agent, matériellement conscient pour l'optimisation des noyaux qui exploite les informations de profilage et construit automatiquement la chaîne complète d'outils de compilation et d'exécution. Les résultats expérimentaux démontrent que CUDAMaster obtient des accélérations significatives sur la plupart des opérateurs, surpassant Astra d'environ 35%. Dans plusieurs cas, ses performances égalent ou dépassent celles de bibliothèques hautement optimisées et propriétaires telles que cuBLAS. Une démonstration présentant le code original et optimisé pour chaque opérateur est disponible à l'adresse 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/.