専門家のようにマルチシナリオCUDAカーネルを最適化するLLMの実現

要旨

GPUカーネルの手動最適化は、困難で時間を要する作業である。大規模言語モデル（LLM）の急速な発展に伴い、GPUカーネルの自動最適化が現実のものとなりつつある。しかし、現在のLLM駆動の自動最適化手法は、PyTorchオペレータの最適化など機械学習アプリケーションに限定されており、科学技術計算における疎行列演算など、より広範な領域が見過ごされている。これらの広範な応用への拡張は、ベンチマークとアルゴリズムに新たな課題をもたらす。そこで我々は、汎用的な自動カーネル最適化手法の開発を主眼とした。本論文では、複数シナリオ設定に対する体系的な評価が欠如している問題に対処するため、MSKernelBenchを提案する。これは基本代数演算、一般的なLLMカーネル、疎行列演算オペレータ、科学技術計算ルーチンといった複数シナリオを網羅し、各シナリオはFP32とBF16の両精度をサポートする。このベンチマークを基盤として、プロファイリング情報を活用し、コンパイル及び実行ツールチェーン全体を自動構築する、マルチエージェントでハードウェアを意識したカーネル最適化システムCUDAMasterを導入する。実験結果により、CUDAMasterは大半のオペレータで大幅な高速化を達成し、Astraを約35%上回ることを示す。いくつかのケースでは、cuBLASのような高度に最適化されたクローズドソースライブラリの性能に匹敵あるいは凌駕する。各オペレータの元のコードと最適化されたコードを実演するデモは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/.

専門家のようにマルチシナリオCUDAカーネルを最適化するLLMの実現

Making LLMs Optimize Multi-Scenario CUDA Kernels Like Experts

要旨

Support