Fast Kronecker Matrix-Matrix Multiplication on GPUs

Abstract

Kronecker Matrix-Matrix Multiplication (Kron-Matmul) is the multiplication of a matrix with the Kronecker Product of several smaller matrices. Kron-Matmul is a core operation for many scientific and machine learning computations. State-of-the-art Kron-Matmul implementations utilize existing tensor algebra operations, such as matrix multiplication, transpose, and tensor matrix multiplication. However, this design choice prevents several Kron-Matmul specific optimizations, thus, leaving significant performance on the table. To address this issue, we present FastKron, an efficient technique for Kron-Matmul on single and multiple GPUs. FastKron is independent of linear algebra operations enabling several new optimizations for Kron-Matmul. Thus, it performs up to 40.7x and 7.85x faster than existing implementations on 1 and 16 GPUs respectively.

arXiv commentAccepted at PPoPP 2024

References in corpus (10)

• GPyTorch: Blackbox Matrix-Matrix Gaussian Process Inference with GPU Acceleration
• Kernel Interpolation for Scalable Structured Gaussian Processes (KISS-GP)
• High-Performance Tensor Contraction without Transposition
• Thoughts on Massively Scalable Gaussian Processes
• Tensor Contractions with Extended BLAS Kernels on CPU and GPU
• Product Kernel Interpolation for Scalable Gaussian Processes
• TSM2X: High-Performance Tall-and-Skinny Matrix-Matrix Multiplication on GPUs
• DISTAL: The Distributed Tensor Algebra Compiler
• TTC: A Tensor Transposition Compiler for Multiple Architectures
• Generalized vec trick for fast learning of pairwise kernel models