KAMI: Communication-Avoiding General Matrix Multiplication within a Single GPU

Efficient general matrix-matrix multiplication (GEMM) has attracted significant research attention in HPC and AI workloads. While large-scale GEMM has nearly achieved the peak floatingpoint performance of GPUs, substantial opportunities for optimization remain in small and batched GEMM operations.

We in this paper propose KAMI, a set of 1D, 2D, and 3D GEMM algorithms that extend the theory of communication-avoiding (CA) techniques within a single GPU. KAMI optimizes thread blocklevel GEMM by utilizing tensor cores as computational units, lowlatency thread registers as local memory, and high-latency onchip shared memory as a communication medium. We provide a theoretical analysis of CA performance from the perspective of GPU clock cycles, rather than the traditional execution time. Also, we implement sparse-dense matrix-matrix multiplication (SpMM) and sparse general matrix-matrix multiplication (SpGEMM) with this compute-communication pattern. Experimental results for general, low-rank, batched, and sparse multiplication on NVIDIA, AMD, and Intel GPUs show significant performance improvements over existing libraries cuBLAS, cuBLASDx, CUTLASS, MAGMA, and SYCL-Bench.