# Fast Inference with Kronecker-Sparse Matrices

arXiv: https://arxiv.org/abs/2405.15013 · HTML: https://mlsanity.com/paper/2405.15013/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 15 of its 44 references that are in the map.
- First posted on arXiv 23 May 2024.

## Builds on (most cited first, 15 of 15)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding](https://mlsanity.com/paper/1510.00149/) (2015)
- [Scaling Vision Transformers](https://mlsanity.com/paper/2106.04560/) (2021)
- [Monarch: Expressive Structured Matrices for Efficient and Accurate Training](https://mlsanity.com/paper/2204.00595/) (2022)
- [Learning Fast Algorithms for Linear Transforms Using Butterfly Factorizations](https://mlsanity.com/paper/1903.05895/) (2019)
- [Pixelated Butterfly: Simple and Efficient Sparse training for Neural Network Models](https://mlsanity.com/paper/2112.00029/) (2021)
- [TorchSparse: Efficient Point Cloud Inference Engine](https://mlsanity.com/paper/2204.10319/) (2022)
- [Monarch Mixer: A Simple Sub-Quadratic GEMM-Based Architecture](https://mlsanity.com/paper/2310.12109/) (2023)
- [SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning](https://mlsanity.com/paper/2207.04606/) (2022)
- [Kaleidoscope: An Efficient, Learnable Representation For All Structured Linear Maps](https://mlsanity.com/paper/2012.14966/) (2020)
- [Butterfly Transform: An Efficient FFT Based Neural Architecture Design](https://mlsanity.com/paper/1906.02256/) (2019)
- [Deformable Butterfly: A Highly Structured and Sparse Linear Transform](https://mlsanity.com/paper/2203.13556/) (2022)
- [Efficient Identification of Butterfly Sparse Matrix Factorizations](https://mlsanity.com/paper/2110.01230/) (2021)
- [Does a sparse ReLU network training problem always admit an optimum?](https://mlsanity.com/paper/2306.02666/) (2023)

## Abstract (arXiv)

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.