Leveraging Large-Scale Data for Efficient Low-Bit CUTLASS GEMM Optimization via Neural Networks

Optimizing GEneral Matrix Multiplication (GEMM) on GPU platforms is becoming increasingly critical to meet the growing computational demands of modern deep neural network research. While significant progress has been made in accelerating high-precision GEMM, the optimization of low-bit GEMM remains a challenging open problem. The CUTLASS library provides highly optimized low-bit GEMM templates leveraging Tensor Cores; however, performance varies considerably depending on tile and pipeline configurations across different GPU architectures. In this work, we propose a novel auto-tuning framework for low-bit CUTLASS GEMM, utilizing a neural network model to predict optimal GEMM template parameters for target GPUs. Our model is trained on a synthetic dataset with up to 116100 unique samples, encompassing diverse matrix sizes across various Ampere GPUs, and is thoroughly evaluated on these hardware platforms. Experimental results show that our method achieves an accuracy of up to 95.11% on the validation dataset. Furthermore, real-time evaluations of low-bit data types on the A100 GPU demonstrate speedups of up to 1.99× for GEMM operations and 1.28× for the linear layer, compared to the default CUTLASS templates.