A resource-aware workload scheduling method for unbalanced GEMMs on GPUs

Abstract GEMM (General Matrix Multiplication) serves as a fundamental operator for deep learning computations. Especially in attention-based deep learning models, such as Bert, GPT, and SAM, the sizes of matrices involved in GEMMs exhibit an unbalanced distribution due to the variable input, resulting in the low utilization of hardware resources. To address the issue, this paper proposes inserting a novel GEMM processing layer into the deep learning inference stack and using an adaptive load balancing method to partition and schedule GEMM computation tasks. The method is implemented with hardware runtime resource information, such as the occupancy of computing units, etc. Experiment results show the remarkable performance of our method in unbalanced input GEMM scenarios, achieving an average performance improvement of 2.3x. The method also performs well in attention-based models (GPT-2 and SAM), achieving an average inference speed improvement of 1.1x. These findings highlight the effectiveness of resource-aware algorithm optimization, especially for computation task scheduling.