SpTMM: Multiplying Ternary Matrix Without Multiplication for Transformers

Abstract

1.58-bit quantization methods, which use ternary weight matrices composed of {-1, 0, 1} in Transformer architectures, have attracted attention for their low latency, reduced memory cost, and high energy efficiency. Ternary Matrix Multiplication (TMM) has the important property of eliminating multiplication instructions, and the presence of 0s in the weight matrix allows us to skip the corresponding computations. By co-designing these properties with the NVIDIA GPU architecture, we propose SpTMM, a novel TMM method. We demonstrate that SpTMM outperforms both cuBLAS and bitnet.cpp at sparsity levels of about 33% and above. Our CUDA kernel source code is publicly available at https://github.com/Ogiwara-CostlierRain464/SpTMM.