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

This article focuses on improving the efficiency of General Matrix Multiplication (GEMM) operations in deep learning inference, particularly under unbalanced input scenarios common in attention-based models like GPT-2 and SAM. It proposes an adaptive load balancing (ALB) method that inserts a GEMM processing layer into the inference stack to partition and schedule GEMM tasks using hardware runtime resource information, such as GPU occupancy. Experimental results demonstrate that ALB achieves up to 2.3 times speedup over existing methods in unbalanced GEMM workloads and improves inference speed by about 1.1 times in GPT-2 and SAM models by enhancing GPU resource utilization. The approach integrates hardware feedback into software optimization, employing matrix partitioning and a dynamic programming-based scheduling algorithm to optimize task execution under hardware constraints.