Fast Inference with Kronecker-Sparse Matrices
This work addresses performance bottlenecks for researchers and practitioners using structured matrices in deep learning, offering incremental improvements in efficiency.
The paper tackled the problem of inefficient GPU kernels for Kronecker-sparse matrix multiplication, which suffer from high data movement costs, and achieved a median speedup of x1.4 and energy reduction of 15% across 600 patterns, with end-to-end latency reductions up to 22% in models like ViT-S/16.
Kronecker-sparse (KS) matrices -- whose supports are Kronecker products of identity and all-ones blocks -- underpin the structure of Butterfly and Monarch matrices and offer the promise of more efficient models. However, existing GPU kernels for KS matrix multiplication suffer from high data movement costs, with up to 50% of time spent on memory-bound tensor permutations. We propose a fused, output-stationary GPU kernel that eliminates these overheads, reducing global memory traffic threefold. Across 600 KS patterns, our kernel achieves in FP32 a median speedup of x1.4 and lowers energy consumption by 15%. A simple heuristic based on KS pattern parameters predicts when our method outperforms existing ones. We release all code at github.com/PascalCarrivain/ksmm, including a PyTorch-compatible KSLinear layer, and demonstrate in FP32 end-to-end latency reductions of up to 22% in ViT-S/16 and 16% in GPT-2 medium.