<figure><img src="image_1.png" /></figure>

HAL
open science

# Make Inference Faster: Efficient GPU Memory Management for Butterfly Sparse Matrix Multiplication

Antoine Gonon, Léon Zheng, Pascal Carrivain, Quoc-Tung Le

<figure><img src="image_2.png" /></figure>

## To cite this version:

Antoine Gonon, Léon Zheng, Pascal Carrivain, Quoc-Tung Le. Make Inference Faster: Efficient GPU Memory Management for Butterfly Sparse Matrix Multiplication. 2024. hal-04584450v2

HAL Id: hal-04584450

https://hal.science/hal-04584450v2

Preprint submitted on 23 May 2024

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.
# Fast Inference with Kronecker-Sparse Matrices

Antoine Gonon, Léon Zheng, Pascal Carrivain, Quoc-Tung Le

► To cite this version:

Antoine Gonon, Léon Zheng, Pascal Carrivain, Quoc-Tung Le. Fast Inference with Kronecker-Sparse Matrices. International Conference on Machine Learning, Jul 2025, Vancouver (BC), Canada. hal-04584450v5

HAL Id: hal-04584450

https://hal.science/hal-04584450v5

Submitted on 13 Jun 2025

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_1.png" /></figure>

Distributed under a Creative Commons Attribution 4.0 International License
Make Inference Faster:
Efficient GPU Memory Management for Butterfly
Sparse Matrix Multiplication

Antoine Gonon*

Univ Lyon, EnsL, UCBL,
CNRS, Inria, LIP

Léon Zheng*

valeo.ai,
Univ Lyon, EnsL, UCBL,
CNRS, Inria, LIP

Pascal Carrivain*

Univ Lyon, EnsL, UCBL,
CNRS, Inria, LIP

Quoc-Tung Le

Univ Lyon, EnsL, UCBL,
CNRS, Inria, LIP

Abstract

This paper is the first to assess the state of existing sparse matrix multiplication algorithms on GPU for the *butterfly* structure, a promising form of sparsity. This is achieved through a comprehensive benchmark that can be easily modified to add a new implementation. The goal is to provide a simple tool for users to select the optimal implementation based on their settings. Using this benchmark, we find that existing implementations spend up to 50% of their total runtime on memory rewriting operations. We show that these memory operations can be optimized by introducing a new CUDA kernel that minimizes the transfers between the different levels of GPU memory, achieving a median speed-up factor of ×1.4 while also reducing energy consumption (median of ×0.85). We also demonstrate the broader significance of our results by showing how the new kernel can speed up the inference of neural networks.

# 1 Introduction

Accelerating the inference and training of deep neural networks is a major challenge given their constantly growing resource requirements. At the very heart of this is the acceleration of matrix multiplication on GPU, which is one of the main operation during training and inference. For instance, in a forward pass of vision transformers (ViTs) [Dosovitskiy et al., 2020], between 30% and 60% of the total time is spent in linear layers (see Appendix B.6 for details). One key approach that aims to accelerate computations is by enforcing *sparsity* constraints on certain weight matrices in the model.

The butterfly matrices have emerged as a promising form of sparse matrices: approximating a matrix by a product of butterfly factors can be done using an efficient algorithm that outperforms gradient descent, with some guarantees of reconstruction if the target matrix admits exactly or approximately a butterfly factorization [Le et al., 2022, Zheng et al., 2023, Le, 2023]; butterfly matrices can be quantized more efficiently than by naive rounding [Gribonval et al., 2023]; and butterfly matrices have a nearly linear *theoretical* complexity for matrix-vector multiplication. The latter is precisely why many linear operators, such as the Discrete Fourier Transform (DFT) or the Hadamard Transform, have fast algorithms: they can be written as butterfly matrices, allowing for fast implementations of these operators. This nearly linear *theoretical* complexity for matrix-vector multiplication comes

<footnote>*Equal contribution</footnote>
Fast Inference with Kronecker-Sparse Matrices

Antoine Gonon*¹² Léon Zheng*¹³⁴ Pascal Carrivain*¹ Quoc-Tung Le¹⁵

**Abstract**

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 ×1.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.

# 1. Introduction

Matrix multiplications dominate the *time* and *energy* budgets of modern-day networks involving large linear layers. For instance, a single forward of a Vision Transformer (ViT-H/14) spends ~60% of its runtime in linear layers (Appendix E.7). Using structured sparse factors in place of dense weights—either learned from scratch or introduced after training—offers a compelling way to reduce memory usage and inference FLOPs (Dao et al., 2019; 2020;

2022a;b). Among the many structures explored, the *Butterfly/Monarch family* has gained traction in both vision and language models (Dao et al., 2019; 2022b; Fu et al., 2023). Its building blocks are matrices whose *support* is a Kronecker product of identities and all-ones blocks. We call them **Kronecker-sparse (KS) matrices** (Definition 2.1), and denote by $(a, b, c, d)$ the pattern of a KS matrix $K$ with support $I_a \otimes I_{b \times c} \otimes I_d$, where $\otimes$ is the Kronecker product, $I_n$ is the $n \times n$ identity matrix, and $I_{b \times c}$ is matrix of size $b \times c$ full of ones. Factorizing a dense matrix $W$ into KS matrices $K_1 \cdots K_L$ can reduce inference cost, as the sequential matrix-vector multiplications $x \mapsto K_1 \cdots K_L x$ can be more efficient than computing the dense multiplication $x \mapsto Wx$. The canonical example is the fast Fourier transform (FFT): the $N \times N$ discrete Fourier transform (DFT) matrix can be factored into $L = \log_2(N)$ KS factors of size $N \times N$, with at most $O(N)$ nonzero entries per factor (Dao et al., 2019). This reduces the overall cost of computing the DFT from $O(N^2)$ (via naive multiplication with the original matrix $W$) to $O(N \log_2 N)$ (through sequential multiplication with the KS factors $K_\ell$); see Figure 1.

<figure><img src="image_2.png" /><figcaption>**Figure 1. FFT as a fast Kronecker-sparse transform.** The dense DFT matrix $W$ can be factored as $W = K_1 \cdots K_L$ (up to a column permutation), where each KS factor $K_\ell$ has support $I_{2^\ell-1} \otimes I_{2 \times 2} \otimes I_{2^L-\ell}$. This factorization underlies the $O(N \log N)$ efficiency of the FFT algorithm.</figcaption></figure>

**Scope of This Work: Focus on Inference.** This work focuses on efficiently computing $Y = XK^\top$, where $K \in \mathbb{R}^{M \times N}$ is a fixed Kronecker-sparse matrix and $X \in \mathbb{R}^{B \times N}$ is a dense matrix representing a batch of inputs. We assume $K$ can be fully preprocessed offline, but inputs and outputs must follow a fixed, dense layout—PyTorch’s default in our case—which cannot be altered. This reflects standard inference scenarios, where models are already trained and deployed, and data arrives in a prescribed format. Optimizing this phase is essential, as inference accounts for over 90% of large-scale ML costs (HPCwire, 2019; Barr, 2019).

*Equal contribution ¹ENS de Lyon, CNRS, Inria, Université Claude Bernard Lyon 1, LIP, UMR 5668, 69342, Lyon cedex 07, France ²Institute of Mathematics, EPFL, Lausanne, Switzerland ³valeo.ai, Paris, France ⁴Huawei Lagrange Mathematics and Computing Research Center, Paris, France ⁵Toulouse School of Economics, Toulouse, France. Correspondence to: Antoine Gonon <antoine.gonon@epfl.ch>.

*Proceedings of the 42nd International Conference on Machine Learning*, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

¹Also integrated into the signal-processing oriented package lazylinop (Appendix D.1).
from the definition of a butterfly matrix $\mathbf{W}$: it must admits a factorization $\mathbf{W} = \mathbf{B}_1 \dots \mathbf{B}_L$, called *butterfly factorization*, where each factor $\mathbf{B}_\ell$ has some specific *structured sparsity pattern* (support of the matrix) with an associated small theoretical multiplication complexity. In general, the support of a butterfly factor is of the form of $\mathbf{S}_\pi = \mathbf{I}_a \otimes \mathbf{1}_{b \times c} \otimes \mathbf{I}_d$ (Figure 2) for some tuple of integers $\pi = (a, b, c, d)$, where $\otimes$ denotes the Kronecker product, $\mathbf{1}_{m \times n}$ is the matrix of size $m \times n$ full of ones, and $\mathbf{I}_n$ is the identity matrix of size $n$, see Definition 2.1 below [Lin et al., 2021, Le, 2023]. A concrete example is given in Figure 1.

<figure><img src="image_3.png" /><figcaption>Figure 1: Example of butterfly factorization $\mathbf{W} = \mathbf{B}_1 \dots \mathbf{B}_L$, for $L = 4$. Here, the factor $\mathbf{B}_\ell \in \mathbb{R}^{N \times N}$ (with $N = 2^L$) has support $\mathbf{S}_\ell = \mathbf{I}_{2^\ell-1} \otimes \mathbf{1}_{2 \times 2} \otimes \mathbf{I}_{2^{L-\ell}}$. This corresponds to the butterfly factorization of the Discrete Fourier Transform matrix $\mathbf{W}$, up to a permutation of its column indices. In practice, the goal is to replace weight matrices of neural networks by butterfly matrices $\mathbf{W} = \mathbf{B}_1 \dots \mathbf{B}_L$ while having (i) at least the same accuracy for the learning task at hand, (ii) less parameters to store, and (iii) an accelerated inference and training phase. To the best of our knowledge, previous works mostly focused on (i) and (ii) [Vahid et al., 2020, Lin et al., 2021, Dao et al., 2022a,b]. This paper is the first<fnref n="2" /> to extensively study the *time efficiency* of these butterfly networks.</figcaption></figure>

**Main contributions.** (i) The first contribution is to assess for the *first time* the efficiency of many baseline PyTorch GPU implementations for multiplying a batch of vectors with a butterfly matrix, including those that rely on existing efficient routines for batch GEMM<fnref n="3" />, block-sparse matrix multiplication and tensor contraction. The goal is to provide a benchmark that can be easily adapted to include a new implementation, and that can be used to select the best implementation for given settings.

(ii) The existing implementations call high-performance libraries from Python. However, Python lacks routines to explicitly control the GPU memory transfers, and because of that, we find with our benchmark that existing implementations spend up to 50% of their total runtime on GPU memory rewriting operations. To address this, we release a new open-source CUDA kernel that minimizes the transfers between the different levels of GPU memory, achieving a median speed-up factor of ×1.4 in *float-precision* while also improving energy efficiency with a median reduction factor of ×0.85. We further show that the new kernel gets more and more advantageous when the relative number of memory rewrites increases. We also demonstrate the broader significance of our results by showing how the new kernel can accelerate the inference of neural networks.

**Outline.** Section 2 introduces the framework to study butterfly matrix multiplication, and describes baseline GPU implementations on PyTorch. Section 3 assesses for the first time the cost of GPU memory access in these baselines. Section 4 explains how the new CUDA kernel reduces the memory transfer compared to previous existing implementations. Section 5 benchmarks the execution time and the energy consumption of baseline GPU implementations on PyTorch, and the new kernel, for the multiplication with a single butterfly factor. Section 6 concretely illustrates broader implications of this work: the new kernel can be used to speed up the inference of neural networks.

## 2 Background on butterfly factorization

We adopt the definition of butterfly matrices stemming from Lin et al. [2021] and mathematically formalized in Le [2023]. To the best of our knowledge, it captures all the variants of butterfly factorizations that have been empirically tested for deep neural networks in the literature [Dao et al., 2019, 2022a,b, Vahid et al., 2020, Lin et al., 2021, Fu et al., 2023]. Details about this unification are given for the curious readers in Appendix C.2.

<footnote n="2">²See Appendix A for the only numerical reports we found in the literature.</footnote>

<footnote n="3">³GEMM stands for General Matrix Multiplication.</footnote>
**Cost of Memory Operations.** The fastest current GPU implementations for Kronecker-sparse multiplication—namely, Monarch’s implementation (Dao et al., 2022b), which we refer to as BMM (as it relies on batched general matrix multiplication), and the block-sparse format method (BSR)—follow a shared strategy: (i) permute the input to ensure that specific elements are stored contiguously in memory; (ii) invoke a high-performance multiplication routine, such as a general matrix multiplication (GEMM); and (iii) apply an inverse permutation to restore the output to its original layout.

However, our profiling (in Section 3 below) reveals that these two permutations—before and after the multiplication—along with the associated memory reads and writes, can account for up to 50% of the total runtime (Figure 4). Notably, we find that the overhead from these data movement operations increases with the ratio $h(b, c) = \frac{b+c}{bc}$, a simple expression that we motivate analytically and validate empirically across 600 KS patterns in this paper.

**Contributions.** We show that the data movement costs caused by the two permutations are not unavoidable. By introducing a new *output-stationary* tiling scheme for KS matrix multiplication, we fuse the three GPU kernels—permute, GEMM, and inverse permute—into a *single* kernel, eliminating two global memory passes. Our main contributions are as follows:

(i) **Time & Energy Benchmark.** We introduce the first large-scale public benchmark of KS matrix multiplication on GPU, spanning 6 orders of magnitude in matrix sizes, several sparsity levels, and supporting both single (FP32) and half (FP16) floating-point formats. Our results confirm that the proportion of the runtime spent at memory rewrites grows with $h(b, c) = \frac{b+c}{bc}$.

(ii) **Single-Kernel Implementation.** We propose an *output-stationary* tiling that reduces the cost of memory operations. We release publicly our kernel implementation in CUDA and OpenCL, plus a drop-in *KSLinear* layer for PyTorch. Our kernel achieves a median speedup of ×1.4, and a median energy reduction of 15% in FP32.

(iii) **Design Heuristic.** For a KS pattern $(a, b, c, d)$, we empirically show that the ratio $h(b, c) = \frac{b+c}{bc}$ predicts latency, while $d \times h(b, c)$ predicts energy, giving practitioners a one-line rule for pattern selection.

(iv) **End-to-End Impact.** We demonstrate that injecting KS sparsity in linear layers of Transformers such as ViT-S/16 in vision and GPT-2 medium in language cuts wall-clock inference in FP32 by 22% and 16% respectively, compared to the original dense implementation.

**Outline.** Section 2 formalizes KS matrices and reviews existing GPU algorithms on PyTorch. Section 3 quantifies

the time they spend on memory permutations. Section 4 presents the fused kernel and its analysis. Section 5 benchmarks time and energy. Section 6 demonstrates network-level gains. Section 7 concludes with open directions.

## 2. Kronecker-Sparse Matrices in a Nutshell

A *Kronecker-sparse (KS)* matrix imposes constraints on the *locations* of its nonzero entries (i.e., its support), but places no restrictions on the values of those entries. The support is defined by a Kronecker product of three simple matrices—identity, all-ones, and identity—which forms the common structural core of the Butterfly, Monarch, Kaleidoscope, and related layers (Dao et al., 2019; 2022a;b; Lin et al., 2021; Fu et al., 2023).

**Definition 2.1 (KS pattern and matrix).** A KS pattern is an integer 4-tuple $\pi = (a, b, c, d)$. Its corresponding support is $S_\pi = I_a \otimes I_{b \times c} \otimes I_d$ (see Figure 2). A matrix $K \in \mathbb{R}^{abd \times acd}$ is $\pi$-Kronecker-sparse if $\text{supp}(K) \subseteq \text{supp}(S_\pi)$, where $\text{supp}(\cdot)$ is the set of matrix indices corresponding to nonzero entries. We write $K \in \mathcal{K}_\pi$.

<figure><img src="image_3.png" /><figcaption>Figure 2. A $\pi$-Kronecker-sparse matrix with $\pi = (a, b, c, d)$ is a block-diagonal matrix with $a$ blocks, where each block itself is a block matrix composed by $b \times c$ diagonal matrices of size $d \times d$. The colored cells correspond to the nonzeros. We color the cells with different colors to indicate that the corresponding weights are free to take different values.</figcaption></figure>

The mask $S_\pi$ has shape $abd \times acd$ and $abcd$ nonzeros, so the density (of nonzero entries) is $1/(ad)$. Increasing either $a$ or $d$ makes the factor sparser. See Appendix B for:

* a discussion on where KS fits in the wider sparsity landscape, with in particular Table 4 showing how popular structured layers in neural networks map to concrete KS patterns;

* links between KS matrix multiplication, sparse 3D convolutions, and sparse-tensor compilers;

* a brief overview of techniques with provable guarantees for approximating a dense matrix $W$ by a KS product $K_1 \cdots K_L$. Note that this approximation problem lies outside the scope of this paper—we assume a fixed KS matrix $K$ is given and focus on efficient multiplication with dense input batches $X$; see Section 1.
**Definition 2.1** (Butterfly factor, architecture and matrix). Given a tuple $\pi := (a, b, c, d) \in (\mathbb{N}_{>0})^4$, a matrix $\mathbf{B} \in \mathbb{R}^{abcd \times acd}$ is called a $\pi$-butterfly factor (or simply factor when $\pi$ is clear from the context) if $\text{supp}(\mathbf{B}) \subseteq \text{supp}(\mathbf{S}_{\pi})$, where $\mathbf{S}_{\pi} := \mathbf{I}_a \otimes \mathbf{1}_{b \times c} \otimes \mathbf{I}_d$ (see Figure 2) and where $\text{supp}(\mathbf{M})$ denotes the support of a matrix $\mathbf{M}$, i.e., the subset of indices $(i, j)$ for which the entries of $\mathbf{M}$ at $(i, j)$ is nonzero. The set of $\pi$-butterfly factors is denoted $\Sigma^{\pi}$.

<figure><img src="image_4.png" /><figcaption>Figure 2: A $\pi$-butterfly factor with $\pi = (a, b, c, d)$ is a block-diagonal matrix with $a$ blocks, where each block itself is a block matrix composed by $b \times c$ diagonal matrices of size $d \times d$.</figcaption></figure>

A butterfly architecture of depth $L \in \mathbb{N}_{>0}$ is a sequence $\beta = (\pi_{\ell})_{\ell=1}^L$ of patterns $\pi_{\ell} = (a_{\ell}, b_{\ell}, c_{\ell}, d_{\ell}) \in (\mathbb{N}_{>0})^4$ satisfying the following compatibility condition$^4$: $a_{\ell}c_{\ell}d_{\ell} = a_{\ell+1}b_{\ell+1}d_{\ell+1}$ for $\ell \in \{1, \dots, L-1\}$. A matrix $\mathbf{W}$ is called a butterfly matrix if there exists a butterfly architecture $\beta = (\pi_{\ell})_{\ell=1}^L$ with associated $\pi_{\ell}$-butterfly factors $\mathbf{B}_{\ell}$ such that $\mathbf{W} = \mathbf{B}_1 \dots \mathbf{B}_L$.

Therefore, a $\pi$-butterfly factor is *sparse and structured*. It has at most $abcd$ nonzero entries when $\pi = (a, b, c, d)$, which yields a sparsity ratio $\frac{abcd}{a^2 bcd^2} = \frac{1}{ad}$, since it is of size $abd \times acd$.

## 2.1 Generic algorithm for butterfly sparse matrix multiplication

Algorithm 1 is a generic algorithm tailored to the butterfly sparsity, allowing for the multiplication of a batch of vectors with a single $\pi$-butterfly factor. It generalizes to general butterfly patterns $\pi = (a, b, c, d)$ the one suggested by Dao et al. [2022b] in the specific cases $a = 1$ or $d = 1$.

**Algorithm 1** Butterfly sparse matrix multiplication

**Input:** $\pi = (a, b, c, d)$, $\mathbf{B} \in \Sigma^\pi$, $\mathbf{X} \in \mathbb{R}^{K \times N}$ ($N := acd$)

**Output:** $\mathbf{Y} = \mathbf{X}\mathbf{B}^\top \in \mathbb{R}^{K \times M}$ ($M := abd$)

1: $\mathbf{Y} \leftarrow \mathbf{0}_{K \times M}$

2: **for** $(i, j) \in [0, a-1] \times [0, d-1]$ **do**

3:     $\mathbf{col} \leftarrow \{i \frac{N}{a} + j + \ell d \mid \ell \in [0, c-1]\}$

4:     $\mathbf{row} \leftarrow \{i \frac{M}{a} + j + kd \mid k \in [0, b-1]\}$

5:     $\mathbf{Y}[:, \mathbf{row}] \leftarrow \mathbf{X}[:, \mathbf{col}]\mathbf{B}^\top[\mathbf{col}, \mathbf{row}]$

6: **end for**

**Algorithm 2** Equivalent formulation

**Input:** $\pi, \mathbf{X}, \tilde{\mathbf{B}} := \mathbf{P}^\top\mathbf{B}\mathbf{Q}^\top$ with $\mathbf{B} \in \Sigma^\pi, \mathbf{P} := (\mathbf{I}_a \otimes \mathbf{P}_{b,d}), \mathbf{Q} := (\mathbf{I}_a \otimes \mathbf{P}_{c,d})^\top$ cf. (1)

**Output:** $\mathbf{Y} = \mathbf{X}\mathbf{B}^\top \in \mathbb{R}^{K \times M}$

1: $\tilde{\mathbf{X}} \leftarrow \mathbf{X}\mathbf{Q}^\top$

2: $\tilde{\mathbf{Y}} \leftarrow \tilde{\mathbf{X}}\tilde{\mathbf{B}}^\top$

3: $\mathbf{Y} \leftarrow \tilde{\mathbf{Y}}\mathbf{P}^\top$

**Notations.** $\mathbf{X} \in \mathbb{R}^{K \times N}$ is the input matrix (batch size $K$, input dimension $N$). $\Sigma^\pi$ is the set of matrices with butterfly pattern $\pi = (a, b, c, d)$ (Definition 2.1). $\mathbf{0}_{m \times n}$ is the $m \times n$ matrix filled with zeros. For integers $a \le b$, $[a, b] := \{a, a+1, \dots, b\}$. For a matrix $\mathbf{M}$, $\mathbf{M}[I,:]$ is the submatrix restricted to rows $I$, and $\mathbf{M}[I, J]$ is the restriction to rows $I$ and columns $J$. Matrix transposition is represented by $^\top$. Matrix indices start at zero.

**Theoretical complexity.** It is not hard to see that the theoretical complexity of Algorithm 1, defined as the number of scalar multiplications, is $Kabcd$, for a batch size $K$ and a pattern $\pi = (a, b, c, d)$.

**On Algorithm 1, and the equivalent Algorithm 2.** When $d = 1$, the butterfly factor $\mathbf{B}$ is block-diagonal with $a$ dense blocks, as can be seen from Figure 2. In this special case, Algorithm 1 loops over each of these blocks, given by $\mathbf{B}[\mathbf{row}, \mathbf{col}]$, where the subsets $\mathbf{row}$ and $\mathbf{col}$ are indexed by $i \in [0, a-1]$ in Algorithm 1, and performs the matrix multiplication with the corresponding submatrix of $\mathbf{X}$. The general case $d \ge 1$ is similar: the butterfly factor $\mathbf{B}$ is, up to permutation

<footnote>$^4$The compatibility condition ensures that the output dimension of $\mathbf{B}_{\ell+1}$ matches the input dimension of $\mathbf{B}_{\ell}$ so that the product $\mathbf{B}_1 \dots \mathbf{B}_L$ is well-defined.</footnote>
## 2.1. Previous PyTorch GPU Implementations for KS

The fastest publicly available PyTorch implementations optimized for the KS structure follow the three-step procedure described in Algorithm 1: (1) permute the inputs; (2) apply a high-performance GEMM routine to $\tilde{K}$, a pre-permuted, block-diagonal version of $K$; (3) permute the output back to its original layout.

<caption>**Algorithm 1** Permutation-based KS matmul</caption>

**Require:** KS pattern $\pi = (a, b, c, d)$, inputs $X \in \mathbb{R}^{B \times acd}$, permutation matrices $P, Q$ (as in Appendix C.1) pre-permuted KS matrix $\tilde{K} := P^\top K Q^\top$
**Ensure:** outputs $Y := X K^\top \in \mathbb{R}^{B \times abd}$
1: $\tilde{X} \leftarrow X Q^\top$ {permute columns}
2: $\tilde{Y} \leftarrow X \tilde{K}^\top$ {matrix multiply}
3: $Y \leftarrow \tilde{Y} P^\top$ {permute columns back}

Any pair of permutation matrices $P$ and $Q$ in Algorithm 1 yields the correct KS matrix multiplication result. Indeed, any permutation matrices $P$ and $Q$ satisfy $PP^\top = Q^\top Q = I$, so this ensures that the multiplication of a batch of inputs $X \in \mathbb{R}^{B \times acd}$ with a KS factor $K \in \mathbb{R}^{abd \times acd}$ is equal to:

$$Y = X K^{\top} = \underbrace{X Q^{\top}}_{:= \tilde{X}} \underbrace{Q K^{\top}}_{:= \tilde{K}^{\top}} P P^{\top} = \tilde{X} \tilde{K}^{\top} P^{\top}.$$

Existing implementations use a specific pair of permutations $(P, Q)$ to achieve high efficiency. These are chosen so that the permuted matrix $\tilde{K} := P^\top K Q^\top$ becomes block-diagonal with dense sub-blocks (details in Appendix C.1), enabling fast parallel multiplication in Step (2) using standard routines on each dense sub-block.

The main differences between existing implementations of Algorithm 1 lie in how they perform Step (2): either via batched GEMM (BMM) (Dao et al., 2022b) or block-sparse multiplication (BSR, for Block Compressed Sparse Row); see Table 1. We will also compare these to a direct tensor contraction baseline tailored to KS structure—EINSUM—inspired by Dao et al. (2022b). Implementation details are provided in Appendix A.

<table><caption>**Table 1.** Two PyTorch realizations of Algorithm 1.</caption><thead><tr><th></th><th>BMM</th><th>BSR</th></tr></thead><tbody><tr><td>K̃ storage</td><td>(ad, b, c) tensor</td><td>(abd, acd) in BSR</td></tr><tr><td>Permute Q (line 1)</td><td>torch.reshape</td><td></td></tr><tr><td>Multiply step (line 2)</td><td>torch.bmm</td><td>linear</td></tr><tr><td>Permute P (line 3)</td><td>torch.reshape</td><td></td></tr></tbody></table>

The convenience of reducing the problem to a matrix multiplication with the block-diagonal matrix $\tilde{K}$ comes at a

cost: it requires two full-tensor permutations in Steps (1) and (3), which we find to consume a non-negligible portion of the runtime (Section 3). To address this, we introduce a new mathematically equivalent algorithm that avoids these permutations entirely (Section 4).

## 2.2. Reminder on Memory-Layout

PyTorch uses row-major storage for tensors. Placing the batch dimension first (batch-size-first<fnref n="2" />) ensures that each sample is stored contiguously in memory; placing it last (batch-size-last) instead makes the features contiguous. The historical default in most machine learning pipelines (including PyTorch) is batch-size-first. While the primary focus of this paper is to compare kernel implementations, we also study the effect of memory layout and find that batch-size-last performs better for KS matrices. To support both conventions, we provide a PyTorch KSLinear layer compatible with both layouts.

# 3. Cost of the Permutations in Baseline Implementations

As we show in Section 5, BMM is the fastest among existing implementations of KS matrix multiplication. It implements Algorithm 1 directly, with two explicit permutations—lines 1 and 3—that move data between global memory and registers before and after the multiplication with the block-diagonal matrix $\tilde{K}$. Figure 3 (left) illustrates the data flow in this baseline.

In this section, we empirically demonstrate that these two permutations can account for up to 50% of the total runtime, motivating the design of a new tiling strategy that avoids this overhead (Section 4).

## 3.1. Why Data Transfers Matter

GPU memory management plays a critical role in performance optimization. Memory on a GPU is organized hierarchically: global memory is large but slow, registers are small but fast, and shared memory sits in between (NVIDIA, 2024, Section 2.3). By default, data resides in global memory. When a kernel is executed, each GPU thread typically loads data from global memory into registers, performs register-level computations, and writes results back to global memory (Figure 3). Because memory access often becomes a performance bottleneck, minimizing data movement between global memory, shared memory, and registers is essential for an efficient implementation (NVIDIA, 2024, Section 5.3).

<footnote n="2">²We name the two memory layouts batch-size-first and batch-size-last, by analogy with PyTorch’s channels-last optimization, which moves the channels dimension to the end for convolutions.</footnote>
<figure><img src="image_5.png" /><figcaption>Figure 3: Illustration of Algorithm 1 for sparsity pattern $\pi = (2, 3, 2, 3)$ and batch size $K = 8$. The subsets of rows and columns ($row_1, col_1$) are associated with the values $(i, j) = (0, 1)$ in the "for" loop of Algorithm 1, whereas ($row_2, col_2$) are associated with $(i, j) = (1, 1)$.</figcaption></figure>

operations, block-diagonal with ad dense blocks, and Algorithm 1 loops over each of these dense blocks, given by $B[\text{row}, \text{col}]$ with row and col defined in lines 3 and 4. See Figure 3 for an illustration. More precisely, for any $\pi = (a, b, c, d)$, denoting $\tilde{\pi} = (ad, b, c, 1)$, we have:

$$ S_{\pi} = \underbrace{(I_a \otimes P_{b,d})}_{:=P} \underbrace{(I_{ad} \otimes 1_{b \times c})}_{:=S_{\pi}} \underbrace{(I_a \otimes P_{c,d})^{\top}}_{:=Q} = PS_{\pi}Q, \quad (1) $$

where $P_{p,q}$ for two integers $p, q$ is the so-called (p, q) perfect shuffle permutation matrix of size $pq \times pq$ [Van Loan, 2000] (see Appendix C.1 for details). Therefore, for any $B \in \Sigma^{\pi}$, we have $B = P\tilde{B}Q$ with $\tilde{B} := P^{\top}BQ^{\top} \in \Sigma^{\pi}$, i.e., $\tilde{B}$ is block-diagonal with ad dense blocks of size $b \times c$. Algorithm 1 loops over each of the ad dense submatrices $B[\text{row}, \text{col}]$ and accumulates the result in $Y$. Many concrete implementations of Algorithm 1 presented below are based on the equivalent formulation $Y = XB^{\top} = XQ^{\top}\tilde{B}^{\top}P^{\top}$: they directly store $\tilde{B}^{\top}$ instead of $B^{\top}$, permute the inputs with $Q$, multiply with $\tilde{B}^{\top}$, and repermute with $P$, which corresponds to Algorithm 2.

## 2.2 Baseline GPU implementations

We now describe concrete baseline GPU implementations of Algorithms 1 and 2. The exact codes are given in Appendix D.1.

**bmm and bsr implementations.** We consider the bmm implementation from Dao et al. [2022b], and we also propose a new implementation bsr. Note that the original bmm implementation from Dao et al. [2022b] only works for a pattern $\pi = (a, b, c, d)$ satisfying $a = 1$ or $d = 1$. We extend it to the general case. Both bmm and bsr implement Algorithm 2 as specified by Table 1. For the multiplication with $\tilde{B}$ (line 2 in Algorithm 2), bmm relies on batched GEMM NVIDIA routines called through `torch.bmm`, while bsr relies on the PyTorch block-sparse library.

<table><caption>Table 1: Differences in the implementation of Algorithm 2 between bmm and bsr.</caption><thead><tr><th></th><th>bmm</th><th>bsr</th></tr></thead><tbody><tr><td>Storage format for <span style="text-decoration: overline;">B</span></td><td>3D-tensor of shape (ad, b, c)</td><td>2D-tensor of shape (abd, acd) stored in BSR<sup>3</sup> format</td></tr><tr><td>Line 1 of Algorithm 2</td><td></td><td>torch reshape</td></tr><tr><td>Line 2 of Algorithm 2</td><td>torch.bmm</td><td>torch nn functional linear</td></tr><tr><td>Line 3 of Algorithm 2</td><td></td><td>torch reshape</td></tr></tbody></table>

**einsum implementation.** We propose a new baseline implementing Algorithm 1 with tensor contractions, using the einops library [Rogozhnikov, 2021]. It stores the nonzero entries of $B \in \Sigma^{\pi}$ with a 4D-tensor $B_{\text{einsum}}$ of shape $(a, b, c, d)$, in such a way that the slice $B_{\text{einsum}[i, \cdot, \cdot, j]}$ for $(i, j) \in [0, a-1] \times [0, d-1]$ stores the entries of $B[\text{row}, \text{col}]$ where row, col are defined in lines 3 and 4 of Algorithm 1. The batched matrix multiplication operations at line 5 are then implemented using Einstein summation between this 4D-tensor and a reshaped input tensor.

The above implementations (bmm, bsr, einsum) can be compared to the two following generic implementations (dense and sparse) that ignore the butterfly sparsity.

**dense implementation.** This ignores the sparsity of the butterfly factor $B$, by storing all its entries, including zeros, in a tensor of shape $(M, N)$. The multiplication is done with `torch nn functional linear`, the default PyTorch implementation for linear layers.
<figure><img src="image_4.png" /><figcaption>Figure 3. Data flow in BMM (three global passes) (Dao et al., 2022b) versus our fused KERNEL (one pass).</figcaption></figure>

## 3.2. Cost of Data Transfers in Baseline Implementations

We focus on BMM as we will find out it is the fastest existing baseline. To isolate the cost of the memory rewriting operations corresponding to the permutations in Algorithm 1 (lines 1 and 3), we compare the runtime of BMM with a modified version in which these permutations are simply removed. Figure 4 (y-axis) shows that across 600 KS patterns, permutations consume up to 45% of the time<fnref n="3" />, and the share grows with $h(b, c) = \frac{b+c}{bc}$; we provide more explanation about this ratio in Section 4.

<figure><img src="image_5.png" /><figcaption>Figure 4. Share of runtime spent only on permutations in BMM (FP32, BSF layout). Each box aggregates patterns with the same $h(b, c) = \frac{b+c}{bc}$.</figcaption></figure>

**Take-Away:** Reducing the cost of memory transfers is a promising direction for improving the performance of Kronecker-sparse matrix multiplication. This is the main objective of our new implementation, described in the following section.

## 4. Novel Tiling for KS Multiplication with Reduced Memory Transfers

Baseline implementations call three separate GPU kernels (PERMUTE-GEMM-PERMUTE); every call round-trips through global memory (Figure 3, left).

A naive approach to prevent such back and forths would be to directly try to merge these three different kernels into a single one. However, this is not as straightforward as it seems since a thread can only access data processed by other threads in the *same block* (no communication/synchronization across thread blocks). Thus, if a thread is assigned tiles (submatrices) on which to apply the three operations PERMUTE-MULTIPLY-PERMUTE, its workload cannot depend on the results obtained by a thread in *another* block. This requires rearranging the tiles assigned to each thread to ensure both *efficiency* and that no thread waits for the result of another thread in *another* block.

To this end, we introduce in Section 4.1 a novel, mathematically equivalent reformulation of the permutation-based algorithm (Algorithm 1), formalized in Algorithm 2. This reformulation leads to a new tiling strategy that enables us to implement the multiplication in a single GPU kernel, as described in Section 4.2. We then perform an analytical comparison of memory operations between our implementation and existing baselines (Section 4.3), and derive a heuristic that predicts when our kernel is expected to reduce the cost of the permutation-related memory operations of Algorithm 1. We validate this prediction empirically in Section 5.

### 4.1. New Tiling Strategy

We introduce a new tiling strategy for KS matrix multiplication, to reduce input/output (I/O) transfers between the different memory levels. Tiling consists of partitioning the

<footnote n="3">³Regardless of the memory layout convention, batch-size-first or batch-size-last as shown in Appendix E.3.</footnote>
**sparse implementation.** This exploits the sparsity of the butterfly factor $\mathbf{B}$ but not its structure (re- call that the sparsity pattern is not arbitrary, but structured as Kronecker products, see Definition 2.1). The nonzero entries of the factor $\mathbf{B}$ are saved in a tensor stored in the Compressed Sparse Row (CSR) format, and the matrix multiplication is done with `torch.nn(functional.linear)`.

**Batch-size-first vs. batch-size-last.** The entries of the input $\mathbf{X} \in \mathbb{R}^{K \times N}$ can be stored either in a PyTorch tensor $\mathbf{X}_{bsf}$ of shape $(K, N)$, or in a PyTorch tensor $\mathbf{X}_{bsl}$ of shape $(N, K)$, in such a way that the entries of the row $\mathbf{X}[k, :]$ are stored in the slices $\mathbf{X}_{bsf}[k, :]$ and $\mathbf{X}_{bsl}[k, :]$. Because of PyTorch's row-major convention, the tensor $\mathbf{X}_{bsf}$ stores in contiguous memory the entries of each row $\mathbf{X}[k,:]$, as opposed to $\mathbf{X}_{bsl}$ that store contiguously the entries of each column $\mathbf{X}[:, i]$. These two different memory layouts are called *batch-size-first* and *batch-size-last*<fnref n="6" /> in this paper. Note that the tensor saving the output $\mathbf{Y} = \mathbf{X}\mathbf{B}^\top$ will always be in the same memory layout as the input tensor. All the implementations above can be implemented in both ways. While the main point of the paper is to compare the implementations, we will also study the effect of this memory layout convention.

## 3 Memory accesses in baseline implementations

While it is not clear what the tensor contraction *einsum* implementation is exactly doing underneath, the implementations *bmm* and *bsr* explicitly perform permutation operations corresponding to **P** and **Q** (lines 1 and 3 in Algorithm 2) to be able to use high-performance multiplication routines for the multiplication with $\tilde{\mathbf{B}}$ (line 2 in Algorithm 2). This paper assesses for the first time the cost of these memory operations in practice, as we now discuss.

**Importance of data transfers.** GPU memory management plays a critical role in optimizing performance. Memory in a GPU is organized hierarchically, with global memory being the largest and slowest, followed by shared memory, and finally registers, which are the smallest and fastest [NVIDIA, 2024, Sec. 2.3]. By default, data resides in the global memory of the GPU. Each thread of the GPU runs a kernel that reads data from global memory into registers, performs register- level computations, and writes the results back to global memory. Therefore, when operations are bottlenecked by memory accesses, it is critical to minimize data transfers between global memory, shared memory, and registers to obtain an efficient GPU implementation [NVIDIA, 2024, Sec. 5.3].

**Data transfers in baseline implementations.** In this paper, we argue that the baseline *bmm*, *bsr* and *einsum* implementations for butterfly multiplication require performing several passes between global memory and registers that can account for a large proportion of the total runtime in practice. This suggests that there is room for improvement in the memory accesses of these implementations.

<figure><img src="image_6.png" /><figcaption>Figure 4: Data flow between the different levels of GPU memory for the *bmm* implementation (Section 2.2) from Dao et al. [2022b] and the new *kernel* (Section 4).</figcaption></figure>

<footnote n="6">⁶By analogy with the recent PyTorch optimization channels last that moves the channels dimension to the last position for convolutional layers.</footnote>
<figure><img src="image_6.png" /><figcaption>Figure 5. Illustration of Algorithm 2 for sparsity pattern $\pi = (2, 3, 2, 3)$ and batch-size $B = 8$. The rows and columns $(row_1, col_1)$ are associated with $(i, j) = (0, 1)$ in the “for” loop of Algorithm 2, whereas $(row_2, col_2)$ is associated with $(i, j) = (1, 1)$.</figcaption></figure>

matrices into smaller submatrices called tiles, which are processed in parallel, with the final result obtained by accumulating intermediate results from each tile (NVIDIA, 2024; 2023a). Our strategy arises from a mathematically equivalent reformulation of Algorithm 1, formalized as Algorithm 2, which we now explain.

**Algorithm 2** New mathematically equivalent tiling of Algorithm 1 (no global memory permutations), see Figure 5.

**Require:** pattern $\pi = (a, b, c, d)$, inputs $X \in \mathbb{R}^{B \times acd}$, KS factor $K \in \mathcal{K}_{\pi}$

**Ensure:** output $Y := XK^{\top} \in \mathbb{R}^{B \times abd}$

1: $Y \leftarrow 0$ {in global memory}

2: **for** $i = 0$ to $a-1$, $j = 0$ to $d-1$ **in parallel do**

3:     $row \leftarrow \{i \frac{M}{a} + j + kd : 0 \le k < b\}$

4:     $col \leftarrow \{i \frac{N}{a} + j + ld : 0 \le l < c\}$

5:     $Y[:, row] \leftarrow Y[:, row] + X[:, col] K^{\top}[col, row]$

6: **end for**

**Why Algorithm 2 is Mathematically Equivalent to Algorithm 1?** Given a KS pattern $\pi = (a, b, c, d)$, the corresponding support $S_{\pi} := I_a \otimes 1_{b \times c} \otimes I_d$ has shape $abd \times acd$. We partition its row indices as $[[0, abd - 1]] := \{0, \dots, abd - 1\} = \bigcup_{i=0}^{a-1} \bigcup_{j=0}^{d-1} row_{i,j}$ where:

$$row_{i,j} := \left\{ i \frac{M}{a} + j + kd : 0 \le k < b \right\}. \quad (1)$$

Due to the structure of $S_{\pi}$, all $b$ rows in each $row_{i,j}$ share the same support, which we denote $col_{i,j}$:

$$col_{i,j} := \left\{ i \frac{N}{a} + j + ld : 0 \le l < c \right\}. \quad (2)$$

The family $\{col_{i,j}\}_{(i,j)}$ forms a partition of the column indices $[[0, acd - 1]]$. This means the matrix multiplication $Y = XK^{\top}$, when restricted to the rows in $row_{i,j}$, can be expressed as:

$$ \begin{aligned} Y[:, row_{i,j}] &= X (K^{\top}[:, row_{i,j}]) \\ &= X[:, col_{i,j}] K^{\top}[col_{i,j}, row_{i,j}], \end{aligned} \quad (3) $$

where $M[:, J]$ denotes matrix M restricted to columns $J$, and $M[I, J]$ denotes the submatrix of M with rows $I$ and columns $J$. By grouping rows that share the same support, the full KS matrix multiplication reduces to a set of smaller, independent dense multiplications—one per $row_{i,j}$—that can be processed in parallel. This is illustrated in Figure 5. Hence, Algorithm 2 is mathematically equivalent to Algorithm 1.

## 4.2. New Kernel Implementation

We now introduce our fused implementation of Algorithm 2, denoted ` KERNEL ` in the remainder. A compact pseudocode sketch is given in Algorithm 3.

**` KERNEL ` versus the permutation-based baseline.** While both Algorithms 1 and 2 ultimately compute the same set of dense sub-products—one per tile $row_{i,j}$—they differ in how they reach them, and hence in their natural GPU implementations:

* **Algorithm 1.** This baseline performs two explicit permutations of $X$ and $Y$ to map the entries of $X[:, col_{i,j}]$ and $Y[:, row_{i,j}]$—whose columns are spaced by $d$—to contiguous regions in the permuted tensors $\tilde{X}$ and $\tilde{Y}$. This enables the use of standard dense GEMM kernels, which require contiguity in global memory. In practice, this strategy is implemented in BMM via two additional `PERMUTE` kernels, which rewrite each tile contiguously in *global memory*.

* **Algorithm 2.** This variant avoids global permutations altogether. It computes each dense sub-product directly, even when the corresponding rows and columns in $row_{i,j}$ and $col_{i,j}$ are not contiguous in global memory. Our fused implementation `KERNEL` follows this approach: each thread block directly loads $X[:, col_{i,j}]$ and $K^{\top}[col_{i,j}, row_{i,j}]$ from global to shared memory, performs the multiplication, and writes the result into $Y[:, row_{i,j}]$, thus bypassing the two extra permutation passes.
Let us focus on bmm, as we will find it to be faster than einsum and bsr. The data flow of bmm is illustrated in Figure 4. There is one pass between the global memory and the registers to perform the permutation with P (line 3 in Algorithm 2), one for the multiplication with $\tilde{B}$ (line 2), and another one for the permutation with Q (line 1).

Doing at least three passes between the global memory and the registers is necessary for any implementation that performs all the multiplications $X[\cdot, col]B^T[col, row]$ (line 5 in Algorithm 1, or equivalently line 2 in Algorithm 2) by calling high-performance libraries from common Python interfaces (PyTorch, Tensorflow). Indeed, calling multiplication routines from Python requires in general having the entries of $X[\cdot, col]$ and $B^T[col, row]$ stored contiguously in global memory. This is not the case of the entries of $X[\cdot, col]$ when $X \in \mathbb{R}^{K \times N}$ is directly stored contiguously in a 2D-tensor of shape $(K, N)$ or $(N, K)$, and when $d > 1$, because the indices in col are equally spaced by $d$ (see line 4 in Algorithm 2). And this is the only memory layout that we can assume for the entries of $X$ when the matrix multiplication is a part of a larger pipeline, such as in a neural network. Therefore, a first pass between the global memory and the registers is required to write the entries of $X[\cdot, col]$ contiguously in global memory, see Figure 4. Similarly, the result of the multiplication $Y[\cdot, row] = X[\cdot, col]B^T[col, row]$, as returned by an efficient routine called from Python, will in general be stored contiguously in global memory for each pair $(row, col)$. But indices in row are equally spaced by $d$: this requires another pass to rewrite them equally spaced by $d$ in the final output 2D-tensor storing all the entries of $Y$ contiguously in memory. With the pass implied by the actual multiplication routine, this results in three passes between the global memory and the registers (Figure 4).

**Estimated time for memory rewrites in bmm.** To the best of our knowledge, this paper is the first to discuss the practical cost of memory rewrites in baseline implementations such as bmm. For input and output dimensions $N$ and $M$, and a batch size $K$, the memory rewrites of $X[\cdot, col]$ and $Y[\cdot, row]$ for all pairs $(row, col)$ (Algorithm 1) concern $KN + KM$ coefficients to be moved in memory: each coefficient of the input and output tensors is moved exactly once. The total number of scalar multiplications performed when computing the products $Y[\cdot, row] = X[\cdot, col]B^T[col, row]$ for all $(row, col)$ (Algorithm 1) is equal to $K \times \#nnz$ (batch size times the number of nonzero in $B$).

Figure 5 reports time estimates for memory rewrites in the bmm implementation (see Appendix B.2 for more details on the experimental protocol). It shows that the relative time spent on memory rewrites increases with the ratio

$$ \frac{\text{number of memory rewrites}}{\text{number of scalar multiplications}} = \frac{KN + KM}{K \times \#nnz} = \frac{b+c}{bc} \qquad (2) $$

where the last equality holds in the case of a butterfly factor with sparsity pattern $\pi = (a, b, c, d)$, since $N = acd$, $M = abd$ and $\#nnz = abcd$ in this case. Figure 5 shows that these memory rewrites can take up to 45% of the total runtime.<fnref n="7" /> In conclusion, it is crucial to optimize the data transfers between the different levels of GPU memory to improve current implementations.

<figure><img src="image_7.png" /><figcaption>Figure 5: Estimated relative time spent on memory rewrites in bmm for the multiplication with $B \in \Sigma^\pi$, for several $\pi = (a, b, c, d)$. We regroup patterns by their value of $(b+c)/(bc)$, and plot a boxplot to summarize the corresponding measurements.</figcaption></figure>

<footnote n="7">⁷Regardless of the memory layout convention, batch-size-first or batch-size-last.</footnote>
**Thread-block tiling.** We assign to each thread block a single dense sub-product $X[;, col_{i,j}] \cdot K^T[col_{i,j}, row_{i,j}]$. We make threads within a block:

(i) cooperatively load the two tiles from global memory to shared memory,

(ii) move their assigned subtiles to registers,

(iii) perform multiply-and-accumulate operations,

(iv) store intermediate results to shared memory, and

(v) write cooperatively the final output tile back to global memory.

Standard CUDA optimizations such as double buffering, vectorized memory access, and warp-level parallelism are applied (see Appendix D.3 for implementation details).

**Read/write efficiency.** Replacing two permutation kernels with direct memory access leads to different tradeoffs depending on the layout:

* **KS matrix tiles.** As in BMM, the factor $K$ is preprocessed and stored so that every tile $K^T[col_{i,j}, row_{i,j}]$ is already contiguous in global memory. This setup cost is paid once and amortized across inference runs with different inputs $X$, so we exclude it from our analysis (Section 1).

* **Input/output tiles.** Contiguity of memory access now depends on the tensor layout. In the feature-major layout (batch-size-last in row-major), each column of $X[;, col_{i,j}]$ and $Y[;, row_{i,j}]$ is stored contiguously in global memory, enabling efficient coalesced access within columns. In this setting, KERNEL benefits from both eliminating permutations and maintaining memory locality.

In contrast, under the batch-major layout (batch-size-first in row-major), entries within the same column of $X[;, col_{i,j}]$ or $Y[;, row_{i,j}]$ are not contiguous in global memory, as the indices in $col_{i,j}$ and $row_{i,j}$ are spaced by $d$. This leads to strided, non-coalesced access patterns. While KERNEL still avoids the overhead of explicit permutations, this benefit has to be weighted against more fragmented memory access.

In summary, KERNEL eliminates two global-memory round trips compared to permutation-based implementations. It is often faster in the batch-size-last layout, where it benefits from both avoiding explicit permutations and preserving coalesced memory access. In the batch-size-first layout, this advantage must be weighed against the cost of non-coalesced memory access due to strided reads and writes; nevertheless, KERNEL remains competitive, and even outperforms baselines in several cases, as observed in Section 5.

**Algorithm 3 Sketch of the fused output-stationary kernel (one tile ($row_{i,j}, col_{i,j}$) assigned to each thread block).**

**Require:** pattern $\pi = (a, b, c, d)$, inputs $X \in \mathbb{R}^{B \times acd}$, KS factor $K \in K_\pi$

**Ensure:** output $Y := XK^T \in \mathbb{R}^{B \times abd}$

1: identify $(i, j)$ tile assigned to current thread block

2: compute $row, col$ as in Algorithm 2

3: **global** → **shared**: load $X[;, col], K^T[col, row]$

4: **for** stride = 0 to $c-1$ by size-subtile **do**

5:     identify subtile assigned to current thread

6:     **shared**→**registers**: load subtiles of $X[;, col]$ and $K[\text{row}, col]$

7:     Multiply and accumulate into registers

8:     Prefetch next subtile in parallel (double buffering)

9:     **registers**→**global**: store $Y[;, row]$

### 4.3. How Much Traffic Does the Fuse Save?

**Against Specialized Baselines.** Recall $\#nz = abcd$ is the number of nonzeros of a KS matrix with pattern $(a, b, c, d)$, $B$ is the batch size, and $N = acd$, $M = abd$ are the input and output dimensions. Specialised baselines call three kernels (PERMUTE-GEMM-PERMUTE). The first PERMUTE involves $2BN$ global memory operations (read $X \in \mathbb{R}^{B \times N}$ once, write its permuted version $\tilde{X}$ once). The GEMM involves $BN + BM$ global memory operations (read once the permuted inputs $\tilde{X}$, write once the permuted outputs $\tilde{Y}$). The last PERMUTE involves $2BM$ global memory operations of the input and output tensors (read $\tilde{Y}$ once, write $\tilde{Y}$ once). The total number of global reads and writes on the input and output tensors is $3B(N+M)$. Our kernel moves each entry exactly once, saving the extra $2B(N+M)$ global reads/writes of the two permutations. The ratio of wasted global memory traffic to useful multiplies is therefore twice

$$ \frac{B(N + M)}{B \#nz} = \frac{b+c}{bc} = h(b, c), \quad (4) $$

which yields the heuristic that correlated with the time spent on permutations in Figure 4. Since our new kernel precisely reduces the cost of the permutations, we expect better gains for patterns with large $h(b, c)$, which is indeed what we observe in Section 5.

**Against Generic Baselines.** Generic dense/CSR PyTorch baselines, which ignore the KS structure, are noted DENSE and SPARSE. Compared with DENSE, we skip all zero entries of $K$, avoiding unnecessary computations. Compared with SPARSE, we match the number of global reads and writes but enjoy better coalescing thanks to knowing the KS support.

**Take-Away:** The fused, output-stationary KERNEL reduces global reads and writes of the input and output tensors by up to a factor three, directly translating into the speedups and energy cuts in practice, as we now report the next section.
## 4 New CUDA kernel with reduced memory transfers

For better management of memory accesses, we need to go to a lower level than Python. This is where the new CUDA kernel comes in. It has the minimum possible number of back and forth between global memory and registers thanks to custom read and write phases between global and shared memory, tailored to *butterfly sparsity*. These new read and write phases allow to perform a single back and forth between the different levels of memory, as shown in Figure 4.

**Implementation.** The reading phase loads the entries of $X[\cdot, col]$ directly from the global memory by accessing non-contiguous columns of $X$ equally spaced by $d$ (see Figure 3), and stores them contiguously in shared memory. Since we are still loading entire columns of $X$, the reading phase is expected to be more efficient if the entries in the same column are contiguous in memory, in order to ensure that most of the memory accesses remain contiguous and efficient. This will only be the case when $X$ is stored with the *batch-size-last* memory layout (as defined in Section 2.2). When in shared memory, the kernel implements the classic tile matrix multiplication algorithm<fnref n="8" /> [Li et al., 2019, Boehm, 2022, NVIDIA, 2023a,b, 2024] to compute each product $X[\cdot, col]B^T[col, row]$ for each subsets row, col in Algorithm 1. Once these products are computed, the kernel again has a custom writing phase from shared to global memory, rewriting the results that are stored contiguously in shared memory, to non-contiguous locations in global memory, because each submatrix $Y[\cdot, row]$ corresponds to non-consecutive columns of the output $Y$ (see Figure 3). This writing phase is also expected to be more efficient in the *batch-size-last* memory layout, ensuring the entries of a same column of $Y[\cdot, row]$, and thus the write operations, to be contiguous in memory.

**Comparison with other baseline implementations.** This new kernel has fewer global memory accesses compared to the baselines einsum, bsr and bmm, because it only reads each coefficient of $X$ and $B$ once and writes the result of the multiplication $Y$ once, while those baseline implementations read $X$ and $Y$ twice and rewrite them once (to permute them). It also has fewer global memory accesses compared to the dense implementation (Section 2.2), since the dense implementation also reads the zero entries of $B$ and the corresponding coefficients of $X$, while our kernel does not. Compared to the generic sparse implementation, we have the same number of memory accesses, but the sparse implementation is agnostic to the sparsity structure of $B$, so it is expected to be less efficient since the memory accesses are not tailored to the known location of the nonzero entries.

## 5 Benchmarking the multiplication by a single butterfly factor

We perform the first benchmark of different implementations for butterfly matrix multiplication, at the finest granularity, where we measure the time for the multiplication with a single butterfly factor $B$. In particular, we validate numerically the benefits of reducing the memory transfers in the kernel implementation, compared to the baselines einsum, bsr and bmm.

**Protocol.** The benchmark is run in *float-precision* on a subset of 600 sparsity patterns $\pi = (a, b, c, d)$ in $\alpha \times \beta \times \beta \times \alpha$, with $\alpha := \{1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, 96, 128\}$, $\beta := \{48, 64, 96, 128, 192, 256, 384, 512, 768, 1024\}$, such that $b = c$ or $b = 4c$ or $c = 4b$. These patterns correspond to dimensions of butterfly factors $B \in \mathbb{R}^{M \times N}$ with $(M, N) = (abd, acd)$ that could be used in neural networks. We choose as batch size $K = 128 \times 196 = 25088$, which is standard as it corresponds to the effective batch size for linear layers in ViTs where the number of sequences per batch is 128, and the number of tokens per sequence is 196. Further details are given in Appendix B.1.

**Implementations specialized to butterfly sparsity improves over generic implementations.** The first line of Table 2 shows that at least one of the implementations specialized to the butterfly structure among kernel, bmm, einsum and bsr improves over the generic dense and sparse implementation, which do not take into account the butterfly sparsity. The speedup increases with the matrix size, see Appendix B.3.

**The baseline bmm is faster than the other baselines einsum and bsr.** This is shown in the second line of Table 2, where the bmm implementation improves over min(einsum, bsr) in 93% of the tested cases. The speedup increases with the matrix size, see Appendix B.4. Therefore, when comparing

<footnote n="8">⁸Classical optimizations are used, as detailed in Appendix D.2.</footnote>
<table><caption>Table 2. Across the 600 patterns, how often is an algorithm faster and by how much (median speedup in parentheses), in FP32.</caption><thead><tr><th>Comparison</th><th>Win rate</th><th>Median × faster</th></tr></thead><tbody><tr><td>min{KABEL, BMM, EINSUM, BSR} &lt; min{DENSE, SPARSE}</td><td>98.1%</td><td>×6.5</td></tr><tr><td>BMM &lt; min{EINSUM, BSR, DENSE, SPARSE}</td><td>90.0%</td><td>×1.36</td></tr><tr><td>KABEL &lt; min{BMM, EINSUM, BSR, DENSE, SPARSE}</td><td>85.5%</td><td>×1.39</td></tr></tbody></table>

## 5. Benchmark: 600 KS patterns, Time & Energy

We now quantify the benefit of *KABEL* over the fastest public baselines—BMM, BSR, and EINSUM—on a single Nvidia A100 (40 GB) for time and on a single Nvidia V100 (32 GB) for energy<fnref n="4" />.

**Benchmarked KS Patterns.** We explore patterns of the form $\pi = (a, b, c, d)$, sweeping over the space $\alpha \times \beta \times \beta \times \alpha$, with $\alpha = \{1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, 96, 128\}$ and $\beta = \{48, 64, 96, 128, 192, 256, 384, 512, 768, 1024\}$. We constrain the shapes by enforcing $b = c$, $b = 4c$, or $c = 4b$, which reflect common configurations in Transformer architectures.<fnref n="5" /> This results in over 600 distinct KS matrix configurations, with sizes ranging from 102 × 102 up to 131072 × 131072—spanning six orders of magnitude and broadly covering<fnref n="6" /> the hidden dimensions found in vision and language Transformers. The input batch size is fixed to $B = 128 \times 196 = 25088$, corresponding to a batch of 128 sequences with a context length of 196 tokens, which is a typical value encountered in Transformer inference.<fnref n="7" /> All results below take, for every implementation, the most efficient (in time or energy) of *batch-size-first* and *batch-size-last* layouts; the full split is in Appendix E.6.

**Covered Sparsity Levels.** The 600 KS patterns selected for this benchmark are skewed toward high sparsity: the *median sparsity* is 97.9%, and 75% of the patterns have sparsity $\ge 91.7\%$. Additional details on the sparsity distribution are provided in Appendix E.2.

### 5.1. Time: *KABEL* Wins on 85% of Patterns in FP32

Table 2 highlights three key findings: (i) KS-aware implementations (*KABEL*, BMM, EINSUM, BSR) are more than ×6 faster than generic baselines that ignore the KS structure (*DENSE*, *SPARSE*); (ii) among existing methods, BMM is the fastest baseline (excluding our new *KABEL* from the comparison); (iii) our fused, output-stationary *KABEL* surpasses all baselines on 85% of the benchmarked grid in FP32, achieving a median speedup of ×1.4.

Because the benchmark grid includes many highly sparse KS patterns, we also compare performance at fixed sparsity levels in Appendix E.2. This analysis confirms that *KABEL* provides consistent speedups across all sparsity regimes—not just in the extremely sparse or dense cases.

### 5.2. What Drives the Speedup? Mostly the Heuristic

$$h(b, c) = \frac{b+c}{bc}$$

To understand what drives the speedup achieved by *KABEL*, we perform a multiple linear regression<fnref n="8" /> in log-log scale using two predictors: the density<fnref n="9" /> $\frac{1}{ad}$ and the heuristic $h(b, c) = \frac{b+c}{bc}$. Before fitting the regression, we verified that the two predictors—density and the heuristic $h(b, c)$—are only weakly correlated across the 600 KS patterns (Pearson correlation: -0.28). This low correlation ensures that the regression is not affected by multicollinearity, and both predictors can be interpreted as contributing distinct information to the model.

Letting speedup := $\frac{\min \text{ time(BMM, BSR, EINSUM)}}{\text{time(KABEL)}}$, the regression yields:

$$\log(\text{speedup}) \approx 1.69 - 0.031 \log(\text{density}) + 0.325 \log(h).$$

This model achieves an adjusted $R^2$ of 0.697 and passes the Jarque-Bera normality test<fnref n="10" /> on residuals ($p = 0.12$). Both predictors are statistically significant ($p < 0.001$).

Notably, $h(b, c)$ is approximately ×10 more influential than density in explaining the speedup. To further isolate the impact of the heuristic $h(b, c)$ from sparsity, we fix the sparsity and examine how the speedup varies with $h(b, c)$.

<footnote n="4">⁴*pyJoules* is not yet compatible with Nvidia A100 GPUs so we benchmarked the energy consumption on V100 GPUs. Note that we observed the same *relative* time efficiency ranking on the two GPUs.</footnote>

<footnote n="5">⁵Recall that for a KS pattern $(a, b, c, d)$, the input and output dimensions of the matrix are given by $N = acd$ and $M = abd$. When $b = c$, we have $N = M$, yielding a square matrix suitable for weight matrices in the self-attention block. When $b = 4c$ (resp. $c = 4b$), the output dimension becomes $M = 4N$ (resp. $M = N/4$), matching the shape of up-projection (resp. down-projection) layers typically found in feed-forward blocks.</footnote>

<footnote n="6">⁶The largest LLaMA 3 model (405B) includes FFN layers with matrices up to $53248 \times 53248$.</footnote>

<footnote n="7">⁷e.g., in ViT with a patch size of 16 × 16 for image resolution of $224 \times 224$.</footnote>

<footnote n="8">⁸We excluded fully dense patterns from this regression, i.e., those with density = 1.</footnote>

<footnote n="9">⁹We chose density over sparsity here, as it gave better regression results.</footnote>

<footnote n="10">¹⁰We were able to pass the normality test in log-log scale, but not in linear-linear or log-linear settings.</footnote>
<table><caption>Table 2: Percentage out of 600 patterns (a, b, c, d) where algo1 is faster than the algo2 (denoted by time(algo1) < time(algo2)), and the median acceleration factor in such cases (that is, the median ratio time of algo2 / time of algo1). For each implementation, we take the minimum time between the batch-size-first and the batch-size-last memory layout.</caption><thead><tr><th>min time<br>(kernel)<br>bmm<br>einsum<br>bsr</th><th>&lt; min time (dense)<br>sparse</th><th>time(bmm) &lt; min time<br>(einsum)<br>bsr<br>dense<br>sparse</th><th>time kernel) &lt; min time<br>(bmm<br>einsum<br>bsr<br>dense<br>sparse)</th></tr></thead><tbody><tr><td>99.67% (×6.57)</td><td></td><td>92.66% (×1.37)</td><td>88.10% (×1.39)</td></tr></tbody></table>

the new kernel implementation to other baselines, we will mainly focus on the comparison between bmm and kernel.

**The new kernel implementation is faster than existing baselines.** The third row of Table 2 shows that kernel is faster than all other baselines in 88% of the tested patterns. This empirically validates the benefits of the reduced memory transfer in the kernel implementation. In the following, we provide further details on the influence of the memory layout (*batch-size-first* vs. *batch-size-last*) on this improvement. Additionally, we analyze the patterns $\pi = (a, b, c, d)$ for which the kernel outperforms baseline implementations.

<figure><img src="image_8.png" /><figcaption>Figure 6: Time speedup factor of kernel compared to min(bmm, einsum, bsr). For each implementation, we take the minimum time between the *batch-size-first* and *batch-size-last* memory layouts. We regroup the patterns by their value of $(b+c)/(bc)$, and plot a boxplot to summarize the corresponding measurements.</figcaption></figure>

<figure><img src="image_9.png" /><figcaption>Figure 7: Energy consumed by kernel compared to the minimum consumed by bmm, einsum and bsr. For each implementation, we take the minimum energy consumed between the *batch-size-first* and *batch-size-last* memory layouts. We regroup patterns by their value of $d(b+c)/(bc)$.</figcaption></figure>

**Impact of the memory layout.** For baseline implementations, switching to *batch-size-last* yields a high systematic speedup for *sparse*, high variability in the speedup of *bsr*, and essentially no impact to negative impact for the other methods, see Appendix B.5 for numerical results. The important part is that it has no impact on bmm, and since bmm is the fastest baseline implementation (Table 2), switching to *batch-size-last* has no impact on the best of the baseline implementations. However, it yields a systematic speedup (about ×2) for the kernel implementation. This acceleration is expected, since the *batch-size-last* memory layout allows for more efficient memory accesses in the kernel implementation, as detailed in Section 4.

**Analyzing the cases where kernel outperforms baselines.** As seen in Section 4, the kernel has an improved memory access design compared to the rest of the baselines. Figure 6 confirms this experimentally: **the kernel implementation becomes increasingly time-efficient compared to the baseline implementations as the relative number of memory accesses increases**, i.e. when the
<figure><img src="image_7.png" /><figcaption>Figure 6. Comparison of kernel in terms of (a) times speed up factor and (b) energy consumption, relative to min(BMM, EINSUM, BSR), on all the 600 KS patterns in FP32. For each implementation, we take the minimum time/energy spent between the *batch-size-first* and *batch-size-last* memory layouts. We regroup patterns by their value of $h(b, c) = \frac{b+c}{bc}$ for(a) and $d \times h(b, c)$ for (b).</figcaption></figure>

For each sparsity bin in Appendix E.2, we observe a strong and consistent correlation in log-log scale (between 0.76 and 0.91), confirming that $h(b, c)$ is a reliable predictor of performance regardless of sparsity; see Appendix E.2 for details.

In conclusion, maximizing $h(b, c)$ is a simple and effective design principle for selecting KS patterns that benefit from KERNEL's speedup over prior baselines (see Figure 6a). In contrast, density has only a weak and inconsistent influence once $h$ is controlled for.

### 5.3. Energy: 15% Median Saving in FP32, Predicted by $d \times h(b, c)$

Because each column gathered by a permutation is spaced by $d$, the distance travelled by memory requests grows with $d$, and $d \times h(b, c)$ is a good proxy for the proportion of energy spent on memory operations. Figure 6b matches theory: the higher that proxy, the more expensive the rewrites and the larger the energy saving of KERNEL. Overall, KERNEL reduces energy consumption compared to existing baselines on 72% of the patterns, achieving a median saving of 15% in FP32.

### 5.4. Additional Remarks

**BSF vs. BSL.** Switching baselines to *batch-size-last* (BSL) does not change the runtime of BMM<fnref n="11" /> However, KERNEL gains a stable ×2 when switching from *batch-size-first*

(BSF) to *batch-size-last*, thanks to fully coalesced writes as we saw in Section 4.2; full numbers are in Appendix E.6.

**Impact of Size.** Speedups increase with matrix size $M \times N = abd \times acd$ (Figure 12); our largest test case ($131k \times 131k$) already exceeds Llama 3-405B's hidden layers.

**Portability.** All results above use CUDA. For broader hardware support, we also provide an OpenCL version, enabling benchmarks on non-NVIDIA devices like AMD GPUs and CPUs.

### 5.5. Take-Away

Exploiting the KS structure is essential: ×6 vs. dense/CSR generic baselines. Removing the two permutations gives in float-precision an additional ×1.4 in time and 15% energy reduction. The simple rule “maximise $h(b, c)$” predicts where the fused kernel shines; sparsity alone does not.

## 6. Broader Impact: End-to-End Gains on ViT and GPT Models

Industry reports estimate that inference accounts for over 90% of the total cost of machine learning at scale (HPCwire, 2019; Barr, 2019). To test whether the layer-wise speedups of KERNEL from Section 5 translate into application-level gains, we inject KS sparsity into a Vision Transformer (Dosovitskiy et al., 2020) and a GPT-style language model (Radford et al., 2019), and measure inference time.

**Target Architecture.** Fully connected (FC) layers are natural candidates for KS-based acceleration. In our

<footnote n="11">¹¹Note that BMM remains the best baselines compared to EINSUM, BSR both in BSF and BSL layouts (Appendix E.6 for details).</footnote>
following ratio increases (introduced in (2))

$$ \frac{\text{number of memory rewritings}}{\text{number of scalar multiplications}} = (b+c)/(bc). $$

**The kernel improves on energy efficiency.** For sparsity patterns $\pi = (a, b, c, d)$, the integer $d$ corresponds to the distance between two columns of the butterfly factor with the same set of nonzero entries. When $d$ is large, the memory rewritings made by `bmm` are likely to be more expensive as the columns in the same set `col` of Algorithm 1 are further away from each other. In practice, we observe the kernel to be more and more energy-efficient as the relative number of memory rewritings $(b+c)/bc$ times the cost $d$ of these memory rewritings increases, as shown in Figure 7. Overall, the median energy reduction factor is $\times 0.85$, and the new kernel improves the energy consumption in 72% of the tested cases. See Appendix B.1 for details about the measurements. This is particularly noteworthy as it demonstrates that the kernel not only achieves higher time efficiency but also reduces energy consumption compared to other baselines. This dual advantage makes the kernel an effective solution for improving both performance and sustainability.

# 6 Broader implications for neural networks: accelerating inference

The inference of neural networks is claimed to represent 90% of the cost of machine learning at scale according to independent reports from both NVIDIA [HPCwire, 2019] and Amazon Web Services [Jeff Barr, 2019]. We now investigate whether replacing fully-connected layers by butterfly matrices accelerates the inference. While the same could also apply to other architectures, we will consider Vision Transformers (ViTs) [Dosovitskiy et al., 2020]. We find that the computational cost of fully-connected layers is significant in such architectures: depending on the size of the ViT, from 30% to 60% of the total time in a forward pass is spent in fully-connected layers (see Appendix B.6 for details).

**Protocol.** We benchmark in *float-precision* various components of a ViT-S/16 architecture: a linear layer with bias, an MLP with non-linear activation and/or normalization layers, a multi-head attention module, etc. As in *Dao et al.* [2022b], the matrices that are replaced by a butterfly matrix are the weight matrices of linear layers in feed-forward network modules, and the projection matrices for keys, queries and values in multi-head attention modules. We focus on *batch-size-first* as it is the default convention in PyTorch<fnref n="9" />. Details and some additional results are given in Appendix B.7.

**Results.** We denote by `time(fully-connected)` the inference time without butterfly matrices (and therefore, with the standard PyTorch implementation). Table 3 shows that `time kernel` < `time(bmm)` < `time(fully-connected)` over all the different submodules. **This concretely shows that using butterfly matrices and the kernel implementation accelerates the inference of standard neural networks.**

<table><caption>Table 3: Acceleration of submodules of a ViT-S/16 using butterfly matrices.</caption><thead><tr><th></th><th>time(bmm)<br>time(fully-connected)</th><th>time kernel<br>time(fully-connected)</th></tr></thead><tbody><tr><td>Linear N × N</td><td>0.82</td><td><strong>0.50</strong></td></tr><tr><td>Feed-forward network</td><td>0.91</td><td><strong>0.77</strong></td></tr><tr><td>Multi-head attention</td><td>0.87</td><td><strong>0.79</strong></td></tr><tr><td>Block</td><td>0.90</td><td><strong>0.78</strong></td></tr><tr><td>Butterfly ViT-S/16</td><td>0.89</td><td><strong>0.78</strong></td></tr></tbody></table>

# 7 Conclusion

This work is the first to evaluate the efficiency of existing butterfly sparse matrix multiplication algorithms on the GPU. Our benchmark shows that baseline implementations require costly memory rewrites in global memory, which are not negligible in practice. The proposed new CUDA kernel significantly reduces the cost of global memory accesses. Specifically, this new kernel is faster than

<footnote n="9">⁹The insertion of butterfly matrices in the *batch-size-last* memory layout would *a priori* require a careful implementation of the rest of the operations, that are for now optimized in *batch-size-first* in PyTorch.</footnote>
experiments, we simply replace the standard PyTorch `Linear` module with our publicly available `KSLinear` module.<fnref n="12" /> Transformer architectures such as ViT and GPT offer a suitable testbed, as they contain numerous FC layers—accounting for 30–60% of total inference time depending on model size (Appendix E.7).<fnref n="13" /> We evaluate our method on ViT-S/16 (21M parameters) and GPT-2 Medium (345M), which both fit comfortably on a single GPU.

**Setup.** Following the protocol of Dao et al. (2022b), we replace specific weight matrices of linear layers in ViT and GPT with a *product of two* KS factors, using patterns detailed in the appendices. All other components of the Transformer forward pass use the default PyTorch kernels in the default *batch-size-first* layout. Complete setup details are provided in Appendix E.8.

**Results.** Table 3 shows that the fused kernel reduces end-to-end latency by **22%** on ViT-S/16 in FP32, corresponding to a ×1.28 speedup over the original dense implementation. For comparison, the best public baseline achieves an 11% reduction (×1.12 speedup). On GPT-2 Medium, the fused kernel yields a **16%** latency reduction (×1.19 speedup), while the best baseline achieves 12% (×1.14 speedup). Additional experimental results are provided in Appendix E.8.

<table><caption>Table 3. Latency ratio versus the dense baseline (fully-connected). Lower is better.</caption><thead><tr><th></th><th>Sparsity level</th><th>time(BMM)<br>time(DENSE)</th><th>time(KERNEL)<br>time(DENSE)</th></tr></thead><tbody><tr><td>ViT-S/16</td><td>57.1 %</td><td>0.89</td><td><strong>0.78</strong></td></tr><tr><td>GPT-2 Medium</td><td>24.8 %</td><td>0.88</td><td><strong>0.84</strong></td></tr></tbody></table>

**Discussion.** In our experiments, inference in ViT and GPT models is run using the *batch-size-first* layout. However, as discussed in Section 5.4, switching to *batch-size-last* can yield up to a ×2 speedup within KS layers. Extending *batch-size-last* support to other core components of neural network inference—such as activation functions, softmax, and LayerNorm—could unlock additional gains, provided these operations can be implemented with similar efficiency in *batch-size-last*. Assessing the end-to-end impact of *batch-size-last* on full-model performance remains an open engineering challenge.

**Take-Away:** Replacing dense weights by KS factors and using the new `KERNEL` speedups the end-to-end inference of Transformers in FP32, suggesting that the proposed tiling strategy can bring practical benefits in real-world settings.

# 7. Conclusion and Discussion

We revisited GPU multiplication with Kronecker-sparse (KS) matrices, identified that permutation overheads dominate existing kernels, and proposed an output-stationary tiling strategy that fuses both permutations with the GEMM into a single CUDA/OpenCL kernel. Across 600 KS patterns, our kernel is ×1.4 faster (median, FP32) than the best public baseline and reduces energy consumption by 15%. When integrated into ViT and GPT models, it yields practical inference-time speedups in FP32.

## 7.1. Key Take-Aways

**Fused Tiling Strategy.** Our output-stationary tiling reduces memory traffic by fusing both permutations with the GEMM in a single kernel.

**Simple Design Rule.** A one-line heuristic $h(b, c) = \frac{b+c}{bc}$ helps identify patterns well-suited for this kernel.

We release a KSLinear PyTorch module supporting backend selection and both *batch-size-first* and *batch-size-last* memory layouts, and which can be used in place of standard PyTorch linear layers.

## 7.2. Limitations and Next Steps

(i) Gains obtained in FP32 are not as large in FP16/FP8 where tensor cores become available (see Appendix E.9 for the analysis of FP16). This calls for further investigations: is it our kernel that can be further improved, or the other implementations that have been particularly optimized in these regimes?

(ii) *Batch-size-last* layout outperforms *batch-size-first* in our `KERNEL` (Figure 14), but PyTorch defaults to `BSF`. Using `BSL` in `BSF` pipelines requires costly transpositions that may offset the gains. Our `KSLinear` layer supports both: `BSF` for drop-in use, and `BSL` for higher performance when compatible. More broadly, the benefits of `BSL` raise an open question: how does batch dimension placement affect neural network efficiency in general?

(iii) An interesting direction is to extend end-to-end latency experiments to larger models than those in Section 6, and to study the trade-off between inference gains and accuracy when training/finetuning with `KSLinear` layers.

(iv) Cache behavior was not analyzed in depth; future work may reveal strengths or limitations of our tiling strategy and suggest further improvements.

(v) Speedups depend on batch size, KS pattern, and a few kernel hyperparameters (e.g., thread count, tile shape). Tuning is setup-specific and tedious, so we provide pre-sets for the 600 patterns we benchmarked and leave general tuning strategies to future work.

(vi) Beyond NVIDIA hardware, the OpenCL port opens the door to studies on AMD GPUs and other platforms.

<footnote n="12">¹²Available at github.com/PascalCarrivain/ksmm.</footnote>

<footnote n="13">¹³See Appendix E.8 for why applying KS sparsity to convolutional networks poses additional challenges.</footnote>
previous baseline implementations, with a median reduction factor of ×0.85. We also show how it
can be used to accelerate the inference of neural networks.

**Perspectives.** While we have focused on optimizing memory management, the multiplication part of our kernel may still have room for improvement, especially in *half-precision*. We hope this will encourage work in that direction. The new kernel is particularly performant in *batch-size-last*. We hope this will lead to further efforts in the *batch-size-first* setting and encourage revisiting other common operations in neural networks within the *batch-size-last* configuration. A translation of our kernel into OpenCL could enable it to run on AMD hardware and other platforms. We also hope that our benchmark will serve as a baseline for comparing butterfly implementations on other hardware: CPU, Intelligence Processing Unit, FPGA, etc.

Acknowledgments

This work was supported in part by the AllegroAssai ANR-19-CHIA-0009, by the NuSCAP ANR-20-CE48-0014 projects of the French Agence Nationale de la Recherche and by the SHARP ANR project ANR-23-PEIA-0008 in the context of the France 2030 program.

The authors thank the Blaise Pascal Center for the computational means. It uses the SIDUS [Quemener and Corvellec, 2013] solution developed by Emmanuel Quemener.

We would also like to thank Patrick Pérez, Gilles Puy, Elisa Riccietti, Nicolas Brisebarre, and Rémi Gribonval for their useful feedback, and Emmanuel Quemener for reserving computing resources for us while we ran our experiments.

References

Simon Boehm. How to optimize a CUDA matmul kernel for cuBLAS-like performance: A worklog, 2022. https://siboehm.com/articles/22/CUDA-MMM [Accessed: April 2024].

Tri Dao, Albert Gu, Matthew Eichhorn, Atri Rudra, and Christopher Ré. Learning fast algorithms for linear transforms using butterfly factorizations. In *ICML*, 2019.

Tri Dao, Beidi Chen, Kaizhao Liang, Jiaming Yang, Zhao Song, Atri Rudra, and Christopher Re. Pixelated butterfly: Simple and efficient sparse training for neural network models. In *ICLR*, 2022a.

Tri Dao, Beidi Chen, Nimit S Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré. Monarch: Expressive structured matrices for efficient and accurate training. In *ICML*, 2022b.

Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. *Advances in Neural Information Processing Systems*, 35:16344-16359, 2022c.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In *ICLR*, 2020.

Daniel Y Fu, Simran Arora, Jessica Grogan, Isys Johnson, Sabri Eyuboglu, Armin W Thomas, Benjamin Spector, Michael Poli, Atri Rudra, and Christopher Ré. Monarch mixer: A simple sub-quadratic GEMM-based architecture. In *NeurIPS*, 2023.

Rémi Gribonval, Theo Mary, and Elisa Riccietti. Optimal quantization of rank-one matrices in floating-point arithmetic—with applications to butterfly factorizations. preprint, 2023. URL https://inria.hal.science/hal-04125381.

HPCwire. AWS Upgrades its GPU-Backed AI Inference Platform. https://www.hpcwire.com/2019/03/19/aws-upgrades-its-gpu-backed-ai-inference-platform/, March 2019. Accessed: [April 2024].

Jeff Barr. Amazon EC2 Update – Inf1 Instances with AWS Inferentia Chips for High Performance Cost-Effective Inferencing. **aws.amazon.com/blogsaws/amazon-ec2-update-inf1-ins**
tances-with-aws-inferentia-chips-for-high-performance-cost-effective-inf
erencing, 2019. Accessed: [April 2024].
## Acknowledgments

This work was supported in part by the AllegroAssai ANR-19-CHIA-0009, by the NuSCAP ANR-20-CE48-0014 projects of the French Agence Nationale de la Recherche and by the SHARP ANR project ANR-23-PEIA-0008 in the context of the France 2030 program.

Tung Le was supported by AI Interdisciplinary Institute ANITI funding, through the French “Investments for the Future – PIA3” program under the grant agreement ANR-19-PI3A-0004, and Air Force Office of Scientific Research, Air Force Material Command, USAF, under grant numbers FA8655-22-1-7012.

The authors thank the Blaise Pascal Center for the computational means. It uses the SIDUS (Quemener & Corvellec, 2013) solution developed by Emmanuel Quemener. We would also like to thank Patrick Pérez, Gilles Puy, Elisa Riccietti, Nicolas Brisebarre, and Rémi Gribonval for their useful feedback, and Emmanuel Quemener for reserving computing resources for us while we ran our experiments.

## Impact Statement

This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.

## References

Barr, J. Amazon EC2 Update – Infl Instances with AWS Inferentia Chips for High Performance Cost-Effective Inferencing, 2019. URL https://aws.amazon.com/blogsaws/amazon-ec2-update-infl-instances-with-aws-inferentia-chips-for-high-performance-cost-effective-inferencing/. Accessed: [April 2025].

Boehm, S. How to optimize a CUDA matmul kernel for cuBLAS-like performance: A worklog, 2022. http://siboehm.com/articles/22/CUDA-MMM [Accessed: April 2024].

Dao, T., Gu, A., Eichhorn, M., Rudra, A., and Ré, C. Learning fast algorithms for linear transforms using butterfly factorizations. In ICML, 2019.

Dao, T., Sohoni, N. S., Gu, A., Eichhorn, M., Blonder, A., Leszczynski, M., Rudra, A., and Ré, C. Kaleidoscope: An efficient, learnable representation for all structured linear maps. In ICLR, 2020.

Dao, T., Chen, B., Liang, K., Yang, J., Song, Z., Rudra, A., and Re, C. Pixelated butterfly: Simple and efficient sparse training for neural network models. In ICLR, 2022a.

Dao, T., Chen, B., Sohoni, N. S., Desai, A. D., Poli, M., Grogan, J., Liu, A., Rao, A., Rudra, A., and Ré, C. Monarch: Expressive structured matrices for efficient and accurate training. In ICML, 2022b.

Dao, T., Fu, D., Ermon, S., Rudra, A., and Ré, C. Flashattention: Fast and memory-efficient exact attention with io-awareness. In NeurIPS, 2022c.

Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR, 2020.

Fu, D. Y., Arora, S., Grogan, J., Johnson, I., Eyuboglu, S., Thomas, A. W., Spector, B., Poli, M., Rudra, A., and Ré, C. Monarch mixer: A simple sub-quadratic GEMM-based architecture. In NeurIPS, 2023.

Guan, Y., Yu, C., Zhou, Y., Leng, J., Li, C., and Guo, M. Fractal: Joint multi-level sparse pattern tuning of accuracy and performance for DNN pruning. In ACM International Conference on Architectural Support for Programming Languages and Operating Systems, 2024.

HPCwire. AWS Upgrades its GPU-Backed AI Inference Platform. https://www.hpcwire.com/2019/03/19/aws-upgrades-its-gpu-backed-a-i-inference-platform/, March 2019. Accessed: [April 2025].

Le, Q.-T., Zheng, L., Riccietti, E., and Gribonval, R. Fast learning of fast transforms, with guarantees. In ICASSP, 2022.

Le, Q.-T., Riccietti, E., and Gribonval, R. Does a sparse relu network training problem always admit an optimum? In NeurIPS, 2023.

Le, Q.-T., Zheng, L., Riccietti, E., and Gribonval, R. Butterfly factorization with error guarantees. arXiv preprint arXiv:2411.04506, 2024.

Li, X., Liang, Y., Yan, S., Jia, L., and Li, Y. A coordinated tiling and batching framework for efficient GEMM on GPUs. In Symposium on Principles and Practice of Parallel Programming, 2019.

Lin, R., Ran, J., Chiu, K. H., Chesi, G., and Wong, N. Deformable butterfly: A highly structured and sparse linear transform. In NeurIPS, 2021.

NVIDIA. Efficient GEMM in CUDA: documentation, 2023a. https://github.com/NVIDIA/cutl-ass/blob/main/media/docs/efficient_gemm.md [Accessed: April 2024].
Quoc-Tung Le. *Algorithmic and theoretical aspects of sparse deep neural networks*. PhD thesis, 2023. URL https://inria.hal.science/tel-04329531.

Quoc-Tung Le, Léon Zheng, Elisa Riccietti, and Rémi Gribonval. Fast learning of fast transforms, with guarantees. In *ICASSP*, 2022.

Xiuhong Li, Yun Liang, Shengen Yan, Liancheng Jia, and Yinghan Li. A coordinated tiling and batching framework for efficient GEMM on GPUs. In *Proceedings of the 24th Symposium on Principles and Practice of Parallel Programming*, 2019.

Rui Lin, Jie Ran, King Hung Chiu, Graziano Chesi, and Ngai Wong. Deformable butterfly: A highly structured and sparse linear transform. In *NeurIPS*, 2021.

NVIDIA. Efficient GEMM in CUDA: documentation, 2023a. https://github.com/NVIDIA/cuTlass/blob/main/media/docs/efficient_gemm.md [Accessed: April 2024].

NVIDIA. Matrix multiplication background user's guide, 2023b. https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html [Accessed: April 2024].

NVIDIA. CUDA C++ programming guide, 2024. https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html [Accessed: April 2024].

E. Quemener and M. Corvellec. SIDUS—the Solution for Extreme Deduplication of an Operating System. *Linux Journal*, 2013.

Alex Rogozhnikov. Einops: Clear and reliable tensor manipulations with einstein-like notation. In *ICLR*, 2021.

Keivan Alizadeh Vahid, Anish Prabhu, Ali Farhadi, and Mohammad Rastegari. Butterfly transform: An efficient FFT based neural architecture design. In *CVPR*, 2020.

Charles F Van Loan. The ubiquitous kronecker product. *Journal of computational and applied mathematics*, 123(1-2):85-100, 2000.

Phil Wang. Scaled dot-product attention implementation, 2024a. https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html [Accessed: April 2024].

Phil Wang. Simple ViT implementation, 2024b. https://github.com/lucidrains/vit-pytorch/blob/main/vit PyTorch/Simple ViT.py [Accessed: April 2024].

Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In *CVPR*, 2022.

Léon Zheng, Elisa Riccietti, and Rémi Gribonval. Efficient identification of butterfly sparse matrix factorizations. *SIAM Journal on Mathematics of Data Science*, 5(1):22-49, 2023.
NVIDIA. Matrix multiplication background user's guide, 2023b. https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html [Accessed: April 2024].

NVIDIA. CUDA C++ programming guide, 2024. https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html [Accessed: April 2024].

Quemener, E. and Corvellec, M. SIDUS—the solution for extreme deduplication of an operating system. *Linux Journal*, 2013(235):3, 2013.

Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. Language models are unsupervised multitask learners. *OpenAI blog*, 1(8):9, 2019.

Spconv Contributors. Spconv: Spatially sparse convolution library. https://github.com/traveller59/spconv, 2022.

Tang, H., Liu, Z., Li, X., Lin, Y., and Han, S. Torchsparse: Efficient point cloud inference engine. In *Conference on Machine Learning and Systems*, 2022.

Tang, H., Yang, S., Liu, Z., Hong, K., Yu, Z., Li, X., Dai, G., Wang, Y., and Han, S. Torchsparse++: Efficient point cloud engine. In *CVPR Workshops*, 2023.

Vahid, K. A., Prabhu, A., Farhadi, A., and Rastegari, M. Butterfly transform: An efficient FFT based neural architecture design. In *CVPR*, 2020.

Van Loan, C. F. The ubiquitous kronecker product. *Journal of computational and applied mathematics*, 123(1-2):85-100, 2000.

Wang, P. Scaled dot-product attention implementation, 2024a. https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html [Accessed: April 2024].

Wang, P. Simple ViT implementation, 2024b. https://github.com/lucidrains/vit-pytorch/blob/main/vit_pytorch/simple_vit.py [Accessed: April 2024].

Yan, Y., Mao, Y., and Li, B. SECOND: Sparsely embedded convolutional detection. *Sensors*, 18(10):3337, 2018.

Ye, Z., Lai, R., Shao, J., Chen, T., and Ceze, L. SparseTIR: Composable abstractions for sparse compilation in deep learning. In *ACM International Conference on Architectural Support for Programming Languages and Operating Systems*, 2023.

Zhai, X., Kolesnikov, A., Houlsby, N., and Beyer, L. Scaling vision transformers. In *CVPR*, 2022.

Zheng, L., Riccietti, E., and Gribonval, R. Efficient identification of butterfly sparse matrix factorizations. *SIAM Journal on Mathematics of Data Science*, 5(1):22–49, 2023.
# Appendices

## A Related works

We now review the numerical results we found in the literature about time efficiency of existing algorithms for sparse butterfly matrix multiplication.

It is reported in *Dao et al.* [2022b] that some butterfly networks were twice faster to train than their dense counterparts for image classification and language modeling, without additional information.

In *Fu et al.* [2023] is reported an acceleration of $X \mapsto W^{-1}(K \odot WX)$ where $K$ is some dense weight matrix, $\odot$ is the element-wise multiplication, and $W$ is the DFT matrix (which admits a butterfly factorization), as soon as the dimensions of $W$ are at least equal to 4096.

These results do not provide *extensive* information on the efficiency of sparse butterfly matrix multiplication, motivating the benchmark in this paper.

## B Experiments

### B.1 Details on the experiments

The pytorch package version is 2.2 and pytorch-cuda is 12.1.

**Matrix sizes.** In all our experiments with matrices, we set the batch size to $K = 128 \times 196 = 25088$, a very standard choice for ViTs, as this quantity corresponds to the standard number of tokens per sequence (192) multiplied by the standard number of sequences in a batch of inputs (128). When dealing with a batch of images in neural networks, we choose the standard choice of batch size $K = 128$.

**Matrix entries.** The coordinates of any butterfly factor $B \in \mathbb{R}^{abd \times acd}$ with sparsity pattern $(a, b, c, d)$ are drawn i.i.d. uniformly in $[-\frac{1}{\sqrt{c}}, \frac{1}{\sqrt{c}}]$, as for the initialization chosen for training in *Dao et al.* [2022b], while the coordinates of the inputs $X$ are drawn i.i.d. according to a standard normal distribution $N(0, 1)$.

**Benchmarking time execution.** All the experiments measuring the time execution of the implementations (Tables 2, 3, 4 and 7, Figures 5 to 6 and 9 to 17) are done on a single NVIDIA A100-PCIE-40GB GPU on an Intel(R) Xeon(R) Silver 4215R CPU @ 3.20GHz with 377G of memory. The full benchmark took approximately 3 days in an isolated environment, ensuring that no other processes were running concurrently.

Measurements are done using the PyTorch tool `torch.utils.benchmark.Timer`. The medians are computed on at least 10 measurements of 10 runs. In 94.2% of the cases, we have an interquartile range (IQR) that is at least 100 times smaller than the median (resp. 98% for 50 times smaller, and 99.7% for 10 times smaller).

**Benchmarking energy consumption.** Measurements of the energy consumption (Figure 7) is done on a single NVIDIA Tesla V100-PCIE-16GB GPU on an Intel(R) Xeon(R) Silver 4215R CPU @ 3.20GHz with 754G of memory. The full benchmark took approximately 1.5 days in an isolated environment. Measurements are made using the pyJoules software toolkit. The medians are computed on at 10 measurements of at least 16 runs. In 96% of the cases, the IQR is at least 10 times smaller than the median, and in all the cases, it is 5 times smaller.

**Patterns benchmarked for time measurements (Section 5).** The considered patterns are generated by the Python code written in Figure 8. In all the cases, we only consider patterns $(a, b, c, d)$ with $b = c$ or $b = 4c$ or $c = 4d$ to have an input size $N$ and an output size $M$ such that $N = M$ or $N = 4M$ or $M = 4N$. This choice is motivated by the fact that fully-connected layers in ViTs satisfy have input and output sizes satisfying these constraints.

The first "for" loop in Figure 8 generates a wide range of patterns $(a, b, c, d)$ with $a = 1$, as this represents the simplest scenario. Indeed, the case $a > 1$ simply corresponds to the case $a = 1$ but repeated $a$ times in parallel.
# Appendices

**Lay summary.** Today's AI models are costly to run because they spend most of their time and energy multiplying huge tables of numbers. One way to speed this up is to use tables with lots of zeros, since multiplying by zero does nothing and can be skipped. If we know where the zeros are in advance, we can organize the work to skip them efficiently—but how we skip them matters.

In this work, we focus on a promising zero pattern that has been widely studied. We show that, even with this pattern, the fastest current Graphical Processing Unit (GPU) methods waste up to half their time just shuffling numbers around before doing the real calculations. While this shuffling is meant to help (to reorganize the table in big chunks that can be skipped), our work reveals that it is unnecessary, and we propose a faster workaround.

By keeping the table data in place and changing how the GPU handles the work, we skip the zeros without extra shuffling. Our method runs up to 40% faster and uses 15% less energy across 600 test cases. We provide open-source code for both GPUs and CPUs, plus a one-line rule of thumb to tell when it helps. Plugged into large AI models like Transformers, it cuts total running time by up to 22%, paving the way for faster, cheaper AI.

# A. Existing GPU Implementations of Kronecker-Sparse Matrix Multiplication

This section describes the existing GPU implementations for KS matrix multiplication on PyTorch. All source code is available in the forward pass of the `KSLinear` module we release at github.com/PascalCarrivain/ksmm.

## A.1. Specialized Baselines

**BMM (batched GEMM).** The implementation released with the “Monarch” layers of Dao et al. (2022b)<fnref n="14" /> permutes the weight matrix once offline and stores $\tilde{K} \in \mathbb{R}^{ad \times b \times c}$ as a dense `float/half` tensor of shape $(ad, b, c)$. Given a batch of inputs $X \in \mathbb{R}^{B \times a \times d}$, the forward pass executes Algorithm 1 with fast tensor reshapes for the input and output permutations, and with the fast routine `torch.bmm` for the multiplication with $\tilde{K}$:

1) reshape $X$ so that multiplying by $Q^T$ is a view;

2) perform the batched GEMM `torch.bmm`;

3) reshape the result so that multiplying by $P^T$ is a view.

A concise PyTorch version follows (layout `batch-size-first`; swap first/last dimensions for `batch-size-last`), where we denote by `K_bmm` the 4D PyTorch tensor storing $\tilde{K}$ and by `T_bsf` the `batch-size-first` PyTorch tensor corresponding to a matrix $T$:

```python
def kronecker_bmm(X_bsf, K_bmm):
    batch_size = X_bsf.shape[0]
    X_perm = (
        X_bsf.view(batch_size, a, c, d)
        .transpose(-1, -2)
        .reshape(batch_size, a * d, c)
        .contiguous()
        .transpose(0, 1)
    )
    Y_perm = torch.empty(batch_size, a * d, b).transpose(0, 1)
    Y_perm = torch.bmm(X_perm, K_bmm transpose(-1, -2))
    Y_bsf = (
        Y_perm transpose(0, 1)
        .reshape(batch_size, a, d, b)
        .transpose(-1, -2)
        .reshape(batch_size, a * b * d)
    )
    return
```

**BSR (block-sparse GEMM).** The same algebraic steps apply, but $\tilde{K}$ is stored once for all in **Block-compressed Sparse Row** format as a tensor `K_bsr`, and Step (2) calls `torch.sparse.mm`:

```python
def kronecker_bsr(X_bsf, K_bsr):
```

<footnote n="14">¹⁴The original release supported only $\pi = (a, b, c, d)$ with $a=1$ or $d=1$; we extended it to the general case.</footnote>
The second "for" loop in Figure 8 generates patterns with $a > 1$ offering fewer choices for $d$ to keep the benchmark concise in terms of execution time. This loop also imposes additional conditions on $b$ and $c$ (line 28 of the code) that we now explain. Many graphs are plotted based on the ratio $(b+c)/bc$, as introduced in Equation (2). Because of that, our goal was to include as many distinct ratios $(b+c)/bc$ as possible while keeping the benchmark brief. We excluded certain $(b, c)$ values because they resulted in a ratio that was very close to one already in the benchmark and were more computationally intensive.

```python
import itertools

batch_size = 25_088
size_limit = 2_147_483_647

a_list = [1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, 96, 128]
b_list = [48, 64, 96, 128, 192, 256, 384, 512, 768, 1024]
c_list = [48, 64, 96, 128, 192, 256, 384, 512, 768, 1024]
d_list1 = [1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, 96, 128]
d_list2 = [4, 16, 64]

def get_patterns_benchmark():
    patterns_list = []

    def add_pattern(a, b, c, d):
        if batch_size * a * c * d <= size_limit and \
           batch_size * a * b * d <= size_limit and \
           a * b * c * d <= size_limit:
            patterns_list.append((a, b, c, d))

    for b, c, d in itertools product(b_list, c_list, d_list1):
        a = 1
        if (b == c or b == 4 * c or c == 4 * b):
            add_pattern(a, b, c, d)

    for a, b, c, d in itertools product(a_list, b_list, c_list,
                                     d_list2):
        if a != 1 and \
           (b, c) not in [(1024, 256), (256, 1024), (128, 512), (512,
                                 128), (64, 256), (256, 64)] and \
           (b == c or b == 4 * c or c == 4 * b):
            add_pattern(a, b, c, d)

    return patterns_list
```

<caption>Figure 8: Python code to generate the patterns benchmarked for the execution time in the numerical experiments of Section 5.</caption>

**Patterns benchmarked for energy measurements (Section 5).** For the energy measurements, the goal is to have diverse sparsity patterns $(a, b, c, d)$ corresponding to many different ratios $d(b+c)/bc$ to observe the trend in Figure 7, while keeping the benchmark as short as possible. We chose to consider the cartesian product of

```python
a_list = [1, 4, 16, 32, 64]
b_list = [48, 64, 96, 128, 192, 256, 384, 512, 768, 1024]
c_list = [48, 64, 96, 128, 192, 256, 384, 512, 768, 1024]
d_list = [1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64]
```

by skipping as in Figure 8 all the patterns with

```python
(b,c) in [(1024, 256), (256, 1024), (128, 512), (512, 128),
          (64, 256), (256, 64)]
```

and also all the patterns such that
```python
batch_size = X_bsf.shape[0]
X_perm = (
    X_bsf.view(batch_size, a, c, d)
    .transpose(-1, -2)
    .reshape(batch_size, a * c * d)
)
Y_perm = torch.nn-functional.linear(
    X_perm, K_bsr
)
Y_bsf = (
    Y_perm.view(batch_size, a, d, b)
    .transpose(-1, -2)
    .reshape(batch_size, a * b * d)
)
return Y_bsf
```

**EINSUM (tensor contraction).** Following the public prototype in the Monarch repository (Dao et al., 2022b), the *abcd* nonzeros of *K* are packed in $K_{\text{einsum}} \in \mathbb{R}^{a \times b \times c \times d}$. After reshaping the input $X \in \mathbb{R}^{B \times a \times c \times d}$ to $X_{\text{einsum}} \in \mathbb{R}^{B \times a \times c \times d}$ and the output $Y \in \mathbb{R}^{B \times abd}$ to $Y_{\text{einsum}} \in \mathbb{R}^{B \times a \times b \times d}$, the contraction $(Y_{\text{einsum}})_{:,a,b,d} = \sum_c (X_{\text{einsum}})_{:,a,c,d} (K_{\text{einsum}})_{a,b,c,d}$ is issued via `einops.einsum`:

```python
def kronecker_einsum(X_bsf, K_einsum):
    X_perm = einops.rearrange(X_bsf, "... (a c d) -> ... a c d", a=a, c=c, d=d)
    Y_perm = einops.einsum(X_perm, K_einsum, "... a c d, a b c d -> ... a b d")
    Y_bsf = einops.rearrange(Y_perm, "... a b d-> ... (a b d)")
    return Y_bsf
```

The second line of this code does at the same time all the matrix multiplications $Y[:, row] \leftarrow X[:, col]K^{\top}[col, row]$ for all the pairs $(row, col)$ in Algorithm 2.

**Other KS implementations** The “Monarch” batched-GEMM code already covered above (BMM) is the maintained reference for Butterfly/Monarch layers of Dao et al. (2019; 2022b). We could not compile the original Butterfly Transform kernels (Dao et al., 2019; Vahid et al., 2020) on recent CUDA; an in-house re-implementation turned out much slower than other baselines and is therefore omitted from the benchmark.

## A.2. Generic Baselines

**DENSE.** All *MN* entries of *K* are stored (even zero entries) and the forward pass calls `torch.nn-functional.linear`. With *batch-size-first* layout this is the PyTorch default; in *batch-size-last* we select `torch.matmul` after an internal benchmark of alternatives.

**SPARSE.** Only the *abcd* non-zeros are kept (CSR format) and the same `linear/matmul` entry point is used. This baseline uses the sparsity but not its specific *Kronecker* structure.

## A.3. Prior Empirical Reports on the Efficiency of Existing KS Matrix Multiplication Algorithms

Dao et al. (2022b) reported a ×2 training speedup on image classification and language-modeling tasks by replacing dense layers with products of two KS factors.

Fu et al. (2023) measured a speedup exceeding ×2 at inference for FFT-based layers $X \mapsto W^{-1}(K \odot WX)$, where *K* is dense, ⊙ is element-wise multiplication, and *W* is the DFT matrix in dimension ≥ 4096. These gains are relevant to KS matrices since the fast implementation of the DFT (the FFT) is equivalent to a sequential multiplication with Kronecker-sparse matrices (see Figure 1).

Our study complements the above by isolating the most elementary operation in these different pipelines: the multiplication by a *single* KS factor $X \mapsto KX$. We benchmark more than 600 distinct KS patterns spanning six orders of magnitude in matrix shapes, comparing existing GPU baselines on PyTorch.
b != c and b != 4 * c and c != 4 * b

for the same reasons as the ones explained above in the case of the benchmark on the time execution.

**Details on boxplots.** In all boxplots (Figures 5 to 7 and 9 to 17), the orange line corresponds to the median, the boxes to the first and third quartile and the whiskers to the 5th and the 95th percentile. Outliers are not represented on the graph.

## B.2 Estimating the time for memory rewrites in the bmm implementation (Section 3)

**Protocol.** Given a pattern $\pi = (a, b, c, d)$ and an input $\mathbf{X} \in \mathbb{R}^{K \times acd}$ for some batch size $K$, we first measure the time $\Delta t$ to compute $\mathbf{Y} := \mathbf{X}\mathbf{B}^{\top}$ using the bmm implementation. Then, we measure the time $\Delta \tilde{t}$ to perform only the multiplication operations $\mathbf{Y}[\cdot, \text{row}] = \mathbf{X}[\cdot, \text{col}]\mathbf{B}^{\top}[\text{col}, \text{row}]$ in the bmm implementation (line 2 of Algorithm 2). Therefore, the estimated relative time to perform the memory rewrites of lines 1 and 3 of Algorithm 2 is simply $\frac{\Delta t - \Delta \tilde{t}}{\Delta t}$.

**Results.** Figure 5, which is replicated in the left part of Figure 9, shows that the relative time spent doing memory rewrites in bmm increases with the ratio $(b+c)/(bc)$, in the *batch-size-first* memory layout. Figure 9 shows that this is similar for both *batch-size-first* and *batch-size-last*.

<figure><img src="image_10.png" /><figcaption>Figure 9: Estimated relative time spent on memory rewrites in bmm for the multiplication with $B \in \Sigma^\pi$, for several $\pi = (a, b, c, d)$. We regroup patterns by their value of $(b+c)/(bc)$, and plot a boxplot to summarize the corresponding measurements.</figcaption></figure>

## B.3 Details on min time(<u>kernel</u>, <u>bmm</u>, <u>bsr</u>, <u>einsum</u>) vs. min time(<u>dense</u>, <u>sparse</u>) (Section 5)

Figure 10 shows that the speed-up factor of implementations specialized to the butterfly sparsity (<u>kernel</u>, <u>bmm</u>, <u>bsr</u>, <u>einsum</u>) over the generic <u>dense</u> and <u>sparse</u> implementations increases with the matrix size $M \times N$. We recall that $M = acd$ and $N = abd$ for a butterfly factor with pattern $\pi = (a, b, c, d)$.

## B.4 Details on time(<u>bmm</u>) vs. min time(<u>bsr</u>, <u>einsum</u>) (Section 5)

Figure 11 shows that for a sufficient large matrix size $M \times N$, we always have $\text{time}(\text{bmm}) < \text{min time}(\text{bsr}, \text{einsum})$, i.e., the bmm implementation is the most efficient among all baseline implementations (<u>bmm</u>, <u>einsum</u>, <u>bsr</u>).
# B. Connections to Other Sparsity Schemes

## B.1. Kronecker-Sparsity in the Wider Sparsity Landscape

*Which sparsity patterns studied in prior work are instances of Kronecker-structured designs?*

Table 4 shows that several sparsity patterns previously proposed for neural network compression can be expressed as special cases of KS patterns, illustrating the broad applicability of the framework studied in this paper.

<table><caption>Table 4. KS patterns used in prior work to replace dense layers in neural networks. For a product $K_1...K_L$ we list $(a_\ell, b_\ell, c_\ell, d_\ell)_{1\le\ell\le L}$.</caption><thead><tr><th></th><th>MATRIX SIZE</th><th>KRONECKER PATTERNS</th></tr></thead><tbody><tr><td>DENSE</td><td>M × N</td><td>(1, M, N, 1)</td></tr><tr><td>LOW-RANK</td><td>M × N</td><td>(1, M, r, 1), (1, r, N, 1)</td></tr><tr><td>SQUARE DYADIC (DAO ET AL., 2019; VAHID ET AL., 2020)</td><td>2<sup>L</sup> × 2<sup>L</sup></td><td>(2<sup>L-1</sup>, 2, 2, 2<sup>L-ℓ</sup>)<sub>ℓ=1</sub><sup>L</sup></td></tr><tr><td>KALEIDOSCOPE (DAO ET AL., 2022B)</td><td>2<sup>L</sup> × 2<sup>L</sup></td><td>(2<sup>L-1</sup>, 2, 2, 2<sup>L-ℓ</sup>)<sub>ℓ=1</sub><sup>L</sup> ∪ (2<sup>L-ℓ</sup>, 2, 2, 2<sup>L-1</sup>)<sub>ℓ=1</sub><sup>L</sup></td></tr><tr><td>BLOCK BUTTERFLY (DAO ET AL., 2022A)</td><td>2<sup>L</sup>t × 2<sup>L</sup>t</td><td>(2<sup>L-1</sup>, 2t, 2t, 2<sup>L-ℓ</sup>)<sub>ℓ=1</sub><sup>L</sup></td></tr><tr><td>MONARCH (DAO ET AL., 2022B; FU ET AL., 2023)</td><td>M × N</td><td>(1, M/p, min(M, N)/p, p), (p, min(M, N)/p, N/p, 1)</td></tr><tr><td>DEFORMABLE BUTTERFLY (LIN ET AL., 2021)</td><td>a<sub>1</sub>b<sub>1</sub>d<sub>1</sub> × a<sub>L</sub>c<sub>L</sub>d<sub>L</sub></td><td>(a<sub>ℓ</sub>, b<sub>ℓ</sub>, c<sub>ℓ</sub>, d<sub>ℓ</sub>)<sub>ℓ=1</sub><sup>L</sup> WITH a<sub>ℓ</sub>c<sub>ℓ</sub>d<sub>ℓ</sub> = a<sub>ℓ+1</sub>b<sub>ℓ+1</sub>d<sub>ℓ+1</sub>.</td></tr></tbody></table>

*How does the Kronecker-sparsity pattern compare to other forms (unstructured or structured) of sparsity and when is each preferable?*

Structured patterns-Kronecker, block, channel-wise, ...-tend to win on latency and energy as the prior knowledge of the sparsity pattern enables the design of more efficient algorithms for inference on GPU. Beyond implementation, structured sparsity gives stronger theoretical guarantees. For example, the associated function space is closed and a minimizer of the loss is guaranteed to exist, unlike in the unstructured case where solutions may diverge toward non-attaining infima (Theorem 4.2 in (Le et al., 2023)). That said, more specific performance comparisons ultimately depend on the particular task, hardware, model, and sparsity structure considered.

## B.2. Dense→KS Conversion (Out of Scope in this Paper)

*How can one convert a dense layer to a Kronecker-sparse one and what is its computational cost?*

Given a dense matrix W (e.g., a weight matrix obtained from training), one can learn factors $K_1...K_L$ such that $W \approx K_1...K_L$ (in Frobenius norm) using a quasi-optimal butterfly factorization algorithm (with provable error guarantees) that has a roughly $O(MN)$ time complexity for the factorization of a matrix of size $M \times N$ (Le et al., 2024; Zheng et al., 2023; Le et al., 2022). This factorization algorithm is implemented in the lazylinop package (see Appendix D.1 for details about this package). In particular, it is shown that this factorization algorithm is more accurate and efficient than alternating least squares (Lin et al., 2021) or gradient descent (Dao et al., 2019). In our work, we focus on pure inference mode where the KS matrix K is given and fixed. The cost of a potential upstream conversion is therefore ignored in our study.

## B.3. Analogy with Sparse 3D Convolutions

An interesting analogy can be drawn between the positioning of our new implementation relative to existing GPU-based KS matrix multiplication algorithms, and how recent works on sparse 3D convolutions (Tang et al., 2023; Spconv Contributors, 2022; Yan et al., 2018) relate to TorchSparse (Tang et al., 2022).

TorchSparse (Tang et al., 2022) performs sparse 3D convolutions in three stages: GATHER-MATMUL-SCATTER. The works (Tang et al., 2023; Spconv Contributors, 2022; Yan et al., 2018) improved performance by overlapping memory transfers with computation. Similarly, for the sparse problem we tackle, our new dataflow strategy implementing Algorithm 2 leads to a comparable improvement: input/output permutations and matrix multiplication can now be fused into a single kernel (Figure 3, right), enhancing efficiency.

That said, the analogy seems to stop there. The underlying structures of the two problems are quite different, and it does not seem that there is a straightforward way to adapt the tiling strategy from Tang et al. (2023) to the nested Kronecker layout we study.
<figure><img src="image_11.png" /><figcaption>Figure 10: Speed-up factor of min time(<u>kernel</u>, <u>bmm</u>, <u>bsr</u>, <u>einsum</u>) compared to min time(<u>dense</u>, <u>sparse</u>) as a function of the matrix size <i>M</i> × <i>N</i>.</figcaption></figure>

<figure><img src="image_12.png" /><figcaption>Figure 11: Speed-up factor of time(<u>bmm</u>) compared to min time(<u>einsum</u>, <u>bsr</u>) as a function of the matrix size <i>M</i> × <i>N</i>.</figcaption></figure>

## B.5 Details on the impact of the memory layout (Section 5)

Figure 12 shows the impact of the memory layout on the execution time of each implementation.

<figure><img src="image_13.png" /><figcaption>Figure 12: Boxplots of the ratio time of batch-size-first / time of batch-size-last</figcaption></figure>

Table 4 shows the percentage of patterns for which the *kernel* implementation improves over all baseline implementations, either in the *batch-size-first* or the *batch-size-last* memory layout. When restricting all implementations to the *batch-size-first* layout, the *kernel* still improves on 20% of the tested patterns despite non-contiguous memory accesses (Section 4).

<table><caption>Table 4: Percentage out of 600 patterns (a, b, c, d) where algo1 is *faster* than the algo2 (denoted by time(algo1) < time(algo2)), and the median acceleration factor in such cases (that is, the median ratio time of algo2 / time of algo1).</caption><thead><tr><th></th><th>time(<u>kernel</u>) &lt; min time(<u>bmm</u>, <u>einsum</u>, <u>bsr</u>, <u>dense</u>, <u>sparse</u>)</th></tr></thead><tbody><tr><td><u>Batch-size-first</u></td><td>20.0% (&times;1.28)</td></tr><tr><td><u>Batch-size-last</u></td><td>88.1% (&times;1.39)</td></tr></tbody></table>
## B.4. Relation to Sparse-Tensor Compilers

State-of-the-art sparse-tensor compilers, e.g., SparseTIR (Ye et al., 2023) or Fractal (Guan et al., 2024), are primarily designed for conventional sparse formats such as CSR or BSR, combined with regular block tilings. KS matrices, by contrast, exhibit a nested sparsity pattern: a block-diagonal layout with $b \times c$ sub-blocks, each refined internally by a $d \times d$ identity structure. Representing such a hierarchy would require support for an outer block-sparse format with an inner identity pattern—capabilities that current compilers do not offer natively.

Furthermore, while our Algorithm 2 uses standard tiling strategies when possible, it sometimes favors non-contiguous tile shapes along certain axes to reduce memory transfers and improve efficiency. This form of layout flexibility is not yet supported by existing sparse compilers, to the best of our knowledge.

Until compiler infrastructure evolves to support nested sparsity and more flexible tiling schemes, hand-optimized kernels seem to remain an effective solution for running KS workloads.

## C. Perfect-Shuffle Permutations Used in Algorithm 1

In practice, existing implementations of Algorithm 1 select the permutation matrices

$$ \mathbf{P} := (\mathbf{I}_a \otimes \mathbf{P}_{b,d}), \quad \mathbf{Q} := \mathbf{I}_a \otimes \mathbf{P}_{c,d}, \qquad (5) $$

where $\mathbf{P}_{p,q}$ is the $(p, q)$ perfect-shuffle permutation matrix. This appendix reviews perfect-shuffle permutations and shows why the choice in (5) produces the permuted matrix

$$ \tilde{\mathbf{K}} := \mathbf{P}^{\top} \mathbf{K} \mathbf{Q}^{\top} $$

that is block-diagonal with dense sub-blocks—specifically, ad diagonal dense blocks of size $b \times c$. Because each block can be multiplied with highly optimized dense routines, Step (2) of Algorithm 1 becomes fast, motivating the use of these particular permutations.

### C.1. Perfect-Shuffle Permutations: Definition

The following definition and properties follow the presentation in (Van Loan, 2000).

**Definition C.1** ($(p, q)$ perfect-shuffle permutation matrix (Van Loan, 2000)). Let $p, q \in \mathbb{N}$ and set $r = pq$. The $(p, q)$ perfect-shuffle permutation matrix $\mathbf{P}_{p,q}$ is the $r \times r$ permutation obtained by reshaping a length-$r$ vector into a $p \times q$ matrix in row-major order and then reading it column-wise. Equivalently,

$$ \mathbf{P}_{p,q} := \begin{pmatrix} \mathbf{I}_r[R_0,:] \\ \mathbf{I}_r[R_1,:] \\ \vdots \\ \mathbf{I}_r[R_{q-1},:] \end{pmatrix} \qquad (6) $$

where $R_i := \{i+qj | j \in [0, p-1]\}$ for $i \in [0, q-1]$ and $\mathbf{I}_r[R_i,:] \in \{0, 1\}^{p \times r}$ is the restriction of the identity matrix $\mathbf{I}_r$ to the rows indexed in $R_i$.

Because the sets $R_0, \dots, R_{q-1}$ partition $\{0, \dots, r-1\}$, $\mathbf{P}_{p,q}$ is a (row) permutation of $\mathbf{I}_r$, and in particular is itself a permutation matrix.

### C.2. Block-Diagonalizing a KS Matrix with Perfect-Shuffle Permutations

We now show that $\tilde{\mathbf{K}}$ as defined in Algorithm 1

$$ \tilde{\mathbf{K}} := \mathbf{P}^{\top} \mathbf{K} \mathbf{Q}^{\top} $$

is block-diagonal when $\mathbf{P}, \mathbf{Q}$ are chosen as perfect-shuffle permutations as in (5) and $\mathbf{K}$ is a KS matrix with pattern $\pi = (a, b, c, d)$. Consider the KS pattern $\tilde{\pi} = (ad, b, c, 1)$, whose support $\mathbf{S}_{\tilde{\pi}} = \mathbf{I}_{ad} \otimes \mathbf{1}_{b \times c}$ corresponds to ad dense diagonal blocks of size $b \times c$. It suffices to establish

$$ \mathbf{S}_{\pi} = \underbrace{(\mathbf{I}_a \otimes \mathbf{P}_{b,d})}_{:=\mathbf{P}} \mathbf{S}_{\tilde{\pi}} \underbrace{(\mathbf{I}_a \otimes \mathbf{P}_{c,d})^{\top}}_{:=\mathbf{Q}} = \mathbf{P} \mathbf{S}_{\tilde{\pi}} \mathbf{Q}. \qquad (7) $$
## B.6 Time spent in linear layers in vision transformers

This section gives a numerical lower bound estimate on the time spent in fully-connected layers in a Vision Transformer (ViT).

**Results.** Table 5 shows that, for different ViTs, the fraction of computation time solely dedicated to linear layers in feed-forward network modules varies between 31% and 53% in *half-precision*, and 46% and 61% in *float-precision*. This proportion increases with the size of the architecture. This shows that a non-negligible amount of ViTs inference is dedicated to fully-connected layers. Note that the time for the fully-connected linear layers in the multi-head attention module is not included in our measurements, so our estimate is *only a lower bound* on the time effectively devoted to all fully-connected layers in transformer architectures.

<table><caption>Table 5: Median execution times (ms) of the forward pass in a ViT, and the forward pass in an MLP containing only all the linear layers involved in the feed-forward network modules of the ViT. The latter is reported with its ratio over the first. FP16 is *half-precision*, FP32 is *float-precision*.</caption><thead><tr><th rowspan="2">ARCHITECTURE</th><th colspan="2">FP16 (S)</th><th colspan="2">FP32 (S)</th></tr><tr><th>COMPLETE</th><th>LINEAR IN FFNS</th><th>COMPLETE</th><th>LINEAR IN FFNS</th></tr></thead><tbody><tr><td>ViT-S/16</td><td>0.014</td><td>0.0046 (31%)</td><td>0.090</td><td>0.04 (46%)</td></tr><tr><td>ViT-B/16</td><td>0.036</td><td>0.015 (42%)</td><td>0.30</td><td>0.16 (54%)</td></tr><tr><td>ViT-L/16</td><td>0.11</td><td>0.050 (46%)</td><td>1.0</td><td>0.58 (58%)</td></tr><tr><td>ViT-H/14</td><td>0.31</td><td>0.16 (53%)</td><td>2.6</td><td>1.6 (61%)</td></tr></tbody></table>

**Details on the estimation.** The transformer architecture is composed of a sequence of transformer blocks, where each block contains a multi-head attention module and a feed-forward network module. The feed-forward network module is an MLP with one hidden layer of neurons, involving two fully-connected linear layers. Table 5 reports the time to perform sequentially all the fully-connected linear layers (without biases) appearing in feed-forward network modules of the considered ViT. This is compared to the total forward time of the transformer network. This is expected to yield a lower bound since we did not measure the time spent in fully-connected linear layers in the multi-head attention module.

**Experimental settings.** The architecture ViT-S/16 corresponds to the one in Zhai et al. [2022], while the architecture ViT-B/16, ViT-L/16 and ViT-H/14 correspond to those in Dosovitskiy et al. [2020]. Input images are of size 224 × 224. In *float-precision*, the PyTorch implementation of ViT architecture are taken from Wang [2024b]. In *half-precision*, the considered implementation of the transformer architecture uses FlashAttention Dao et al. [2022c] to compute the scaled dot product attention, like in Wang [2024a]. The MLP containing only the linear layers of the feed-forward modules in the transformer architecture is implemented using `torch.nn.Sequenceal` and `torch.nn.Linear`. Experiments are done on a single A100-40GB GPU on AMD EPYC 7742 64-Core Processor. Measurements are done using the PyTorch tool `torch.utils.benchmark.Timer` for benchmarking. The image batch size is set at 128.

## B.7 Details on the acceleration of the inference of a ViT (Section 6)

**Chosen butterfly matrices.** The weight matrices are replaced by butterfly matrices associated to the following butterfly architectures of depth 2 (Definition 2.1): (1, 192, 48, 2), (2, 48, 192, 1) for the size $N \times N$, (1, 768, 192, 2), (6, 64, 64, 1) for the size $4N \times N$, (1, 768, 192, 2), (6, 64, 64, 1) for the size $4N \times N$.

**Additional results.** Table 6 provides additional results to Table 3 on linear submodules of a ViT-S/16.
Because permutation matrices are orthogonal, $P^T S_\pi Q^T = S_{\pi}$. For any compatible matrices A, B, C we have

$$ \text{supp}(ABC) \subseteq \text{supp}(\text{supp}(A) \text{supp}(B) \text{supp}(C)), $$

where $\text{supp}(\cdot)$ is the binary support mask. Since $\text{supp}(P) = P$ and $\text{supp}(Q) = Q$,

$$ \text{supp}(P^T K Q^T) \subseteq \text{supp}(S_{\pi}) = I_{ad} \otimes I_{b \times c}, $$

which is the support of a block-diagonal matrix with ad dense $b \times c$ blocks, proving the claim.

*Proof of (7).* A key identity from (Van Loan, 2000) states that for all positive integers $b, c, d$,

$$ P_{b,d}^{\top} (I_{b \times c} \otimes I_d) P_{c,d} = I_d \otimes I_{b \times c}. $$

Because $S_\pi = I_a \otimes I_{b \times c} \otimes I_d$, we obtain

$$ S_\pi = I_a \otimes (P_{b,d} (I_d \otimes I_{b \times c}) P_{c,d}^\top). $$

Using $(AB) \otimes (CD) = (A \otimes C)(B \otimes D)$ for any matrices A, B, C, D of compatible sizes, we factor:

$$ \begin{aligned} S_\pi &= I_a \otimes (P_{b,d} (I_d \otimes I_{b \times c}) P_{c,d}^\top) \\ &= (I_a \otimes P_{b,d}) (I_a \otimes (I_d \otimes I_{b \times c}) P_{c,d}^\top) \\ &= (I_a \otimes P_{b,d}) (I_a \otimes I_d \otimes I_{b \times c}) (I_a \otimes P_{c,d}^\top) \\ &= (I_a \otimes P_{b,d}) (I_{ad} \otimes I_{b \times c}) (I_a \otimes P_{c,d}^\top). \end{aligned} $$

which is exactly (7).

**Algorithmic Implications.** Because $\tilde{K}$ is block-diagonal, Step (2) of Algorithm 1 can leverage optimized dense kernels on each sub-block. Both BMM and BSR follow this strategy, as summarized in Table 1. The cost, however, is two full-tensor permutations (Steps 1 and 3), which we measure to be a non-negligible share of runtime. Section 4 introduces a mathematically equivalent approach that eliminates these permutations altogether.

# D. Software Integration and Implementation Details of the Fused KS Kernel

This appendix complements Section 4 with practical information about the released software, the lazylinop integration, and the design choices that guided our CUDA kernel.

## D.1. Integration with lazylinop

We provide two complementary entry points for working with Kronecker-sparse matrix multiplications (KSMM), each designed for a specific audience and purpose.

* The GitHub repository ksmm is targeted at the machine learning community. It includes all CUDA/OpenCL source code, benchmarks, and a plug-and-play PyTorch module, KSLinear, that can replace standard Linear layers in neural networks.

* Separately, we are also integrated our fast matrix factorization routines into the lazylinop ecosystem, a package more oriented toward signal processing and fast linear transforms. These two entry points are intended to remain independent moving forward.

lazylinop<fnref n="15" /> is a Python library that supports delayed execution: it builds an operator expression graph lazily and only evaluates it when a final result is requested (delayed evaluation paradigm). One of its key features with respect to

<footnote n="15">¹⁵https://faustgrp.gitlabpages.inria.fr/lazylinop/index.html</footnote>
<table><caption>Table 6: Acceleration of submodules of a ViT-S/16 using butterfly matrices.</caption><thead><tr><th></th><th>time(bmm)<br>time(fully-connected)</th><th>time kernel<br>time(fully-connected)</th></tr></thead><tbody><tr><td>Linear N × N</td><td>0.82</td><td><strong>0.50</strong></td></tr><tr><td>Linear N × N + bias</td><td>0.97</td><td><strong>0.66</strong></td></tr><tr><td>Linear 4N × N</td><td>0.80</td><td><strong>0.78</strong></td></tr><tr><td>Linear 4N × N + bias</td><td>0.93</td><td><strong>0.90</strong></td></tr><tr><td>Linear N × 4N</td><td>0.91</td><td><strong>0.58</strong></td></tr><tr><td>Linear N × 4N + bias</td><td>0.94</td><td><strong>0.61</strong></td></tr></tbody></table>

## B.8 Additional results in half-precision

For the sake of completeness we perform the benchmark described in Section 5 in *half-precision*. The equivalent of Table 2, Figure 6, Figures 9 to 12 in *half-precision* are Table 7, Figure 13, Figures 14 to 17, respectively. Note that just as Figure 6, the Figure 13 only considers sparsity patterns for which `min time kernel, bmm, bsr, einsum) < min time dense, sparse). This corresponds to 87% of the tested patterns in *half-precision*, cf. Table 7.

<table><caption>Table 7: Percentage out of 600 patterns (a, b, c, d) where algo1 is faster than the algo2 in *half-precision* (denoted by `time(algo1) < time(algo2)`), and the median acceleration factor in such cases (that is, the median ratio $\frac{\text{time of algo2}}{\text{time of algo1}}$). For each implementation, we take the minimum time between the *batch-size-first* and the *batch-size-last* memory layout. Experiments are carried in *half-precision*.</caption><thead><tr><th>min time<br>(kernel)<br>bmm<br>einsum<br>bsr</th><th>&lt; min time<br>(dense<br>sparse)</th><th>time(bmm) &lt; min time<br>(einsum<br>bsr<br>dense<br>sparse)</th><th>time kernel<br>&lt; min time<br>(bmm<br>einsum<br>bsr<br>dense<br>sparse)</th></tr></thead><tbody><tr><td>86.95% (×8.45)</td><td></td><td>83.22% (×1.83)</td><td>36.69% (×1.46)</td></tr></tbody></table>

<figure><img src="image_14.png" /><figcaption>Figure 13: Speedup factor of kernel compared to min(bmm, einsum, bsr) in *half-precision*. For each implementation, we take the minimum time between the *batch-size-first* and the *batch-size-last* memory layout. We regroup the (a, b, c, d) patterns by their value of $(b+c)/(bc)$, and use a boxplot to summarize the corresponding measurements. Experiments are carried in *half-precision*.</figcaption></figure>
Kronecker-sparsity is the ability to approximate a dense matrix W into a fast transform, i.e., a product of structured KS factors:

$$W \approx K_1 \cdots K_L$$

This results in fast-transform approximations that might be more efficient in time and energy, especially on parallel hardware.

CUDA and OpenCL backends enable efficient evaluation of these factor chains $K_1 \cdots K_L X$ in a *batch-size-last* layout (batch dimension last). The Butterfly module of lazylinop supports execution on both CPUs and NVIDIA/AMD GPUs. The documentation is available at https://faustgrp.gitlabpages.inria.fr/lazylinop/index.html and the low-level source code is available at https://gitlab.inria.fr/faustgrp/lazylinop.

The current lazylinop package offers the following features for Butterfly and Kronecker-sparse operators:

1) **Butterfly factorization (products of KS matrices) (Le et al., 2024).** Given a dense matrix W, lazylinop approximates it as $W \approx K_1 \cdots K_L$ using a quasi-optimal factorization algorithm that achieves the theoretical error bounds of Le et al., all in $O(MN)$ time for an $M \times N$ matrix.

2) **Compatibility with PyTorch tensors (in progress).** PyTorch tensor inputs and outputs will soon be supported directly within lazylinop, streamlining integration into PyTorch-based workflows.

## D.2. Why an Output-Stationary Dataflow?

Like CUTLASS dense kernels (NVIDIA, 2023a) and recent sparse 3-D-convolution kernels (Tang et al., 2023; Spconv Contributors, 2022; Yan et al., 2018), our fused KS kernel is output-stationary. Two considerations motivate this choice:

1) *Weight-stationary seems to offer little reuse.* In a KS pattern $(a, b, c, d)$, when the parameter $a$ is large, the matrix is partitioned into $a$ submatrices that act on disjoint regions (Figure 2). Thus, weights are not reused across multiple input/output regions, limiting the benefits of keeping them stationary.

2) *Input-stationary seems harder to parallelize.* Due to the non-contiguous memory accesses involved (Figure 3), read and write operations on inputs/outputs come at a higher cost, so it is preferable to keep one of them stationary to reduce these costs. Both costs are largely driven by the parameter $d$ of KS patterns $(a, b, c, d)$, as it determines the distance between consecutive data elements that should be loaded together when considering one of the dense subproblems (Section 4.1). Since their reuse costs are similar, we had to consider other factors. Input-stationarity poses parallelization challenges as different thread blocks cannot accumulate into the same output coefficient (no possible synchronization) (NVIDIA, 2023a). In contrast, output-stationarity avoids this issue, hence our choice.

## D.3. Kernel-Level Optimizations

Our implementation follows standard high-performance GPU practices:

* **Vectorization.** Wherever possible we use `float4` and `half2` loads/stores to coalesce memory traffic (NVIDIA, 2023b; 2024; 2023a; Boehm, 2022).

* **Double buffering.** Each thread block overlaps computation on the current tile with prefetching of the next tile into shared memory (Li et al., 2019; NVIDIA, 2023b).

* **Structured epilogue.** Partial results are first written to shared memory and then flushed to global memory in a contiguous pattern, avoiding scattered global writes (NVIDIA, 2023a).

Kernel launch parameters (block size, tile shape, vector width) are auto-tuned per pattern $\pi = (a, b, c, d)$ and per GPU architecture, mirroring the tuning strategy of CUTLASS.

# E. Experiments

## E.1. Details on the Experiments

The pytorch package version is 2.2 and pytorch-cuda is 12.1.
<figure><img src="image_15.png" /><figcaption>Figure 14: Estimated relative time spent on memory rewrites in bmm for the multiplication with $B \in \Sigma^\pi$, for several $\pi = (a, b, c, d)$. We regroup patterns by their value of $(b+c)/(bc)$, and plot a boxplot to summarize the corresponding measurements. Experiments are carried in *half-precision*.</figcaption></figure>

<figure><img src="image_16.png" /><figcaption>Figure 15: Speed-up factor of min time kernel, bmm, bsr, einsum compared to min time dense, sparse vs. the matrix size $M \times N$. Experiments are carried in *half-precision*.</figcaption></figure>

<figure><img src="image_17.png" /><figcaption>Figure 17: Boxplots of the ratio time of batch-size-first / time of batch-size-last in *half-precision*.</figcaption></figure>
**Matrix Sizes.** In all our experiments with matrices, we set the batch size to $B = 25088$ (196 tokens × 128 sequences), a standard choice for ViTs as it corresponds to the standard number of tokens per sequence (196) multiplied by the standard number of sequences in a batch of inputs (128). When dealing with a batch of images in neural networks, we choose the standard choice of batch size $B = 128$.

**Matrix Entries.** The nonzero entries of any Kronecker-sparse matrix $K \in \mathbb{R}^{abd \times acd}$ with sparsity pattern $(a, b, c, d)$ are drawn i.i.d. uniformly in $[-\frac{1}{\sqrt{c}}, \frac{1}{\sqrt{c}}]$, corresponding to the initialization used for training in Dao et al. (2022b). The entries of the inputs $X$ are drawn i.i.d. according to a standard normal distribution $N(0, 1)$.

**Benchmarking Time Execution.** All the experiments measuring time execution of a Kronecker-sparse matrix multiplication algorithm (Tables 2, 5, 6, 9 and 10, Figures 4, 6a, 9, 10 and 12 to 18) are performed on a NVIDIA A100-PCIE-40GB GPU associated with an Intel(R) Xeon(R) Silver 4215R CPU @ 3.20GHz with 377G of memory. The full benchmark took approximately 3 days in an isolated environment, ensuring that no other processes were running concurrently.

Measurements are done using the PyTorch tool `torch.utils.benchmark.Timer`. The medians are computed on at least 10 measurements of 10 runs. In 94.2% of the cases, we have an interquartile range (IQR) that is at least 100 times smaller than the median (resp. 98% for 50 times smaller, and 99.7% for 10 times smaller).

**Benchmarking Energy Consumption.** Measurements of the energy consumption (Figure 6b) is done on a NVIDIA Tesla V100-PCIE-16GB GPU associated with an Intel(R) Xeon(R) Silver 4215R CPU @ 3.20GHz with 754G of memory. The full benchmark took approximately 1.5 days in an isolated environment. Measurements are made using the pyJoules software toolkit. The medians are computed on 10 measurements of at least 16 runs. In 96% of the cases, the IQR is at least 10 times smaller than the median, and 5 times smaller in all the cases.

**Kronecker-Sparsity Patterns Benchmarked for Time Measurements (Section 5).** The considered patterns are generated by the Python code written in Figure 7. In all the cases, we only consider patterns $(a, b, c, d)$ with $b = c$ or $b = 4c$ or $c = 4b$ to have an input size $N$ and an output size $M$ such that $N = M$ or $N = 4M$ or $M = 4N$. We make this choice because fully-connected layers in Transformers usually satisfy these dimension constraints.

The first “for” loop in Figure 7 (line 21 to 24) generates a wide range of patterns $(a, b, c, d)$ with $a = 1$, as this represents the simplest scenario. Indeed, the case $a > 1$ simply corresponds to repeating $a$ times the case $a = 1$ in parallel.

The second “for” loop in Figure 7 (line 26 to 30) generates patterns with $a > 1$ offering fewer choices for $d$ to keep the benchmark concise in terms of execution time. This loop also imposes additional conditions on $b$ and $c$ (line 28 of the code) that we now explain. Many graphs are plotted based on the ratio $(b+c)/bc$, as introduced in Equation (4). Because of that, our goal was to include as many distinct ratios $(b+c)/bc$ as possible while keeping the benchmark brief. We excluded certain $(b, c)$ values because they resulted in a ratio that was very close to one already in the benchmark and were more computationally intensive.

**Patterns Benchmarked for Energy Measurements (Section 5).** For the energy measurements, the goal is to have diverse sparsity patterns $(a, b, c, d)$ corresponding to many different ratios $d(b+c)/bc$ to observe the trend in Figure 6b, while keeping the benchmark as short as possible. We chose to consider the cartesian product of

```
a_list = [1, 4, 16, 32, 64]
b_list = [48, 64, 96, 128, 192, 256, 384, 512, 768, 1024]
c_list = [48, 64, 96, 128, 192, 256, 384, 512, 768, 1024]
d_list = [1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64]
```

by skipping as in Figure 7 all the patterns with

```
(b,c) in [(1024, 256), (256, 1024), (128, 512), (512, 128), (64, 256), (256, 64)]
```

and also all the patterns such that

```
b != c and b != 4 * c and c != 4 * b
```

for the same reasons as explained above for time measurements.

**Details on boxplots.** In all boxplots (Figures 4, 6a to 6b, 10 and 12 to 18), the orange line corresponds to the median, the
# C Theoretical results

## C.1 Details on perfect shuffle permutations

The goal is to prove Equation (1), which we recall here for convenience:

$$S_{\pi} = \underbrace{(\mathbf{I}_a \otimes \mathbf{P}_{b,d})}_{:=\mathbf{P}} \underbrace{(\mathbf{I}_{ad} \otimes \mathbf{1}_{b \times c})}_{:=\mathbf{S}_{\pi}} \underbrace{(\mathbf{I}_a \otimes \mathbf{P}_{c,d})^{\top}}_{:=\mathbf{Q}} = \mathbf{P} \mathbf{S}_{\pi} \mathbf{Q},$$

where the matrix $\mathbf{P}_{p,q}$ is the so-called $(p, q)$ perfect shuffle permutation introduced below. To prove this formula, we will use the next lemma.

**Lemma C.1.** For any positive integers $b, c, d$:

$$\mathbf{P}_{b,d}^{\top}(\mathbf{1}_{b \times c} \otimes \mathbf{I}_d)\mathbf{P}_{c,d} = \mathbf{I}_d \otimes \mathbf{1}_{b \times c},$$

where $\mathbf{P}_{p,q}$ denotes the $(p, q)$ perfect shuffle of $r := pq$ [Van Loan, 2000], which is the permutation matrix of size $r \times r$ defined as:

$$\mathbf{P}_{p,q} := \begin{pmatrix} \mathbf{I}_r[R_0,:] \\ \mathbf{I}_r[R_1,:] \\ \vdots \\ \mathbf{I}_r[R_{q-1},:] \end{pmatrix}, \qquad (3)$$

where $R_i := \{i+qj | j \in [0, p-1]\}$ for $i \in [0, q-1]$.

*Proof of Lemma C.1.* This is a direct consequence of a more general result claiming that the Kroenecker product commutes up to some perfect shuffle permutation matrices [Van Loan, 2000, Section 1]. □

We now turn to the proof of Equation (1).

*Proof of Equation (1).* By definition, $S_{\pi} = \mathbf{I}_a \otimes \mathbf{1}_{b \times c} \otimes \mathbf{I}_d$ when $\pi = (a, b, c, d)$. By Lemma C.1,

$$S_{\pi} = \mathbf{I}_a \otimes \mathbf{1}_{b \times c} \otimes \mathbf{I}_d = \mathbf{I}_a \otimes (\mathbf{P}_{b,d}(\mathbf{I}_d \otimes \mathbf{1}_{b \times c})\mathbf{P}_{c,d}^{\top}).$$

By the equality $(AB) \otimes (CD) = (A \otimes C)(B \otimes D)$ for any matrices $A, B, C, D$ of compatible sizes, we get the result:

$$\begin{align*} S_{\pi} &= \mathbf{I}_a \otimes (\mathbf{P}_{b,d}(\mathbf{I}_d \otimes \mathbf{1}_{b \times c})\mathbf{P}_{c,d}^{\top}) \\ &= (\mathbf{I}_a \otimes \mathbf{P}_{b,d}) (\mathbf{I}_a \otimes ((\mathbf{I}_d \otimes \mathbf{1}_{b \times c})\mathbf{P}_{c,d}^{\top})) \\ &= (\mathbf{I}_a \otimes \mathbf{P}_{b,d})(\mathbf{I}_a \otimes \mathbf{I}_d \otimes \mathbf{1}_{b \times c})(\mathbf{I}_a \otimes \mathbf{P}_{c,d}^{\top}) \\ &= (\mathbf{I}_a \otimes \mathbf{P}_{b,d})(\mathbf{I}_{ad} \otimes \mathbf{1}_{b \times c})(\mathbf{I}_a \otimes \mathbf{P}_{c,d}^{\top}). \end{align*}$$

## C.2 Existing variants of butterfly factorization

Table 8 summarizes how Definition 2.1 [Lin et al., 2021, Le, 2023] captures all the variants of butterfly factorizations that have been empirically tested for deep neural networks in the literature Dao et al. [2019, 2022a,b], Vahid et al. [2020], Lin et al. [2021], Fu et al. [2023]. Note that the butterfly chain that we consider from Dao et al. [2022b] is the one used in their actual implementation rather than in their paper. A detailed explanation about this unification can be found in Chapter 6 of Le [2023], except for the case of block butterfly Dao et al. [2022a]. We now deal with the latter: we explain why block butterfly matrices are indeed covered by Definition 2.1.

**Block butterfly matrices are covered by Definition 2.1.**

Block butterfly factors [Dao et al., 2022a] are a direct generalization of square dyadic butterfly factors [Dao et al., 2019]. Indeed, replacing each entry of a square dyadic butterfly factor by a block matrix of size $t \times t$ gives a block butterfly factor (and in particular, the square dyadic butterfly factors are
```python
import itertools

batch_size = 25_088
size_limit = 2_147_483_647

a_list = [1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, 96, 128]
b_list = [48, 64, 96, 128, 192, 256, 384, 512, 768, 1024]
c_list = [48, 64, 96, 128, 192, 256, 384, 512, 768, 1024]
d_list1 = [1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, 96, 128]
d_list2 = [4, 16, 64]

def get_patterns_benchmark():
    patterns_list = []

    def add_pattern(a, b, c, d):
        if batch_size * a * c * d <= size_limit and \
           batch_size * a * b * d <= size_limit and \
           a * b * c * d <= size_limit:
            patterns_list.append((a, b, c, d))

    for b, c, d in itertools product(b_list, c_list, d_list1):
        a = 1
        if (b == c or b == 4 * c or c == 4 * b):
            add_pattern(a, b, c, d)

    for a, b, c, d in itertools product(a_list, b_list, c_list, d_list2):
        if a != 1 and \
           (b, c) not in [(1024, 256), (256, 1024), (128, 512), (512, 128), (64, 256),
                          (256, 64)] and \
           (b == c or b == 4 * c or c == 4 * b):
            add_pattern(a, b, c, d)

    return patterns_list
```

<caption>Figure 7. Python code to generate the patterns benchmarked for the execution time in the numerical experiments of Section 5.</caption>

boxes to the first and third quartile and the whiskers to the 5th and the 95th percentile. Outliers are not represented on the graph.

## E.2. Performance Across Sparsity Levels (Section 5)

Figure 8 illustrates the distribution of sparsity levels across the 600 KS patterns used in our benchmark. A clear skew toward highly sparse patterns is observed: the median sparsity is 97.9%, and 75% of all patterns have a sparsity of at least 91.7%. As a result, any aggregated statistic computed over the full set of 600 patterns—such as the medians reported in Table 2—tends to be dominated by these very sparse cases.

To assess whether our findings hold across the entire sparsity spectrum, we complement the overall results in Table 2 with a stratified analysis by sparsity level. Table 5 reports the performance of KERNEL compared to three existing KS-aware baselines—BMM, EINSUM, and BSR—within six sparsity bins. These bins were defined by equally spaced quantiles in log-sparsity scale, to reflect the long-tailed distribution of sparsity levels.

Table 5 confirms that our fused KERNEL consistently outperforms the baselines in each of the different sparsity regimes. Furthermore, within each bin, we compute the correlation between the speedup of KERNEL and the heuristic $h(b, c) = \frac{b+c}{bc}$. In all six bins, we observe strong positive correlations (ranging from 0.76 to 0.91 in log-log scale), indicating that $h(b, c)$ remains a strong predictor of performance independently of sparsity. This supports the conclusion that our kernel’s speedup is driven primarily by the structural pattern captured by $h(b, c)$, rather than by the density of nonzero entries.

We also provide a visual counterpart in Figure 9, which extends the analysis of Figure 6a by showing the variation in speedup as a function of $h(b, c)$ at fixed sparsity levels, rather than aggregated bins. This further illustrates the dominant role of the heuristic in explaining performance gains.
<table><caption>Table 8: Definition 2.1 unifies the different factorizations tested for deep neural networks in the literature, as documented in Lin et al. [2021], Le [2023]. The fact that it also captures block butterfly matrices is new (Appendix C.2).</caption><thead><tr><th></th><th>MATRIX SIZE</th><th>BUTTERFLY ARCHITECTURE</th></tr></thead><tbody><tr><td>DENSE</td><td>M × N</td><td>(1, M, N, 1)</td></tr><tr><td>LOW-RANK</td><td>M × N</td><td>(1, M, r, 1), (1, r, N, 1)</td></tr><tr><td>SQUARE DYADIC [DAO ET AL., 2019, VAHID ET AL., 2020]</td><td>N × N WITH N = 2<sup>L</sup></td><td>(2<sup>L-1</sup>, 2, 2, 2<sup>L-ℓ</sup>)<sub>ℓ=1</sub><sup>L</sup></td></tr><tr><td>KALEIDOSCOPE [DAO ET AL., 2022B]</td><td>N × N WITH N = 2<sup>L</sup></td><td>CONCATENATE (2<sup>L-1</sup>, 2, 2, 2<sup>L-ℓ</sup>)<sub>ℓ=1</sub><sup>L</sup> AND (2<sup>L-ℓ</sup>, 2, 2, 2<sup>L-1</sup>)<sub>ℓ=1</sub><sup>L</sup></td></tr><tr><td>BLOCK BUTTERFLY [DAO ET AL., 2022A]</td><td>N × N WITH N = 2<sup>L</sup></td><td>(2<sup>L-1</sup>, 2t, 2t, 2<sup>L-ℓ</sup>)<sub>ℓ=1</sub><sup>L</sup></td></tr><tr><td>MONARCH [DAO ET AL., 2022B, FU ET AL., 2023]</td><td>M × N</td><td>(1, M/p, min(M, N)/p, p), (p, min(M, N)/p, N/p, 1)</td></tr><tr><td>DEFORMABLE BUTTERFLY [LIN ET AL., 2021]</td><td>INTRODUCED THE GENERAL FRAMEWORK OF DEFINITION 2.1.</td><td></td></tr></tbody></table>

particular cases of block butterfly factors with block size 1 × 1.) The support constraints of block butterfly factors can be expressed by the following binary matrix [Dao et al., 2022a]:

$$ \tilde{\mathbf{S}}_{\ell} := \mathbf{I}_{2^{\ell-1}} \otimes \mathbf{I}_{2 \times 2} \otimes \mathbf{I}_{2^{L-\ell}} \otimes \mathbf{I}_{\ell \times t}, \quad \ell = 1, \dots, L, \qquad (4) $$

where we have an additional parameter $t$ in comparison to the tuples of butterfly factors (Definition 2.1) to control the block size.

We argue that using a chain of block butterfly factors (4) is equivalent to using a chain of butterfly factors in the sense of Definition 2.1. In order to prove that, we first introduce the set of matrices that can be factorized according to the two alternative definitions.

**Definition C.2 (Set of block butterfly matrices).** Consider the chain of $L$ block butterfly factors of block size $t$, define $\mathcal{B}_b^{L,t} := \{\prod_{\ell=1,...,L} \mathbf{X}_\ell \mid \text{supp}(\mathbf{X}_\ell) \subseteq \text{supp}(\tilde{\mathbf{S}}_\ell)\}$ ($\tilde{\mathbf{S}}_\ell$ is defined as in (4)) the set of matrix admitting an exact factorization into block butterfly factors (shortened to block butterfly matrices).

**Definition C.3 (Set of butterfly matrices associated to a butterfly chain)** $\Theta := (\mathbf{t}_\ell)_{\ell=1}^L$. Consider a butterfly chain $\Theta := (\mathbf{t}_\ell)_{\ell=1}^L$, we define $\mathcal{B}^\Theta := \{\prod_{\ell=1,...,L} \mathbf{X}_\ell \mid \text{supp}(\mathbf{X}_\ell) \subseteq \text{supp}(\mathbf{S}_{\mathbf{t}_\ell})\}$ the set of matrix admitting an exact factorization into butterfly factors (shortened to butterfly matrices).

The equivalence between block butterfly factors and butterfly chain is shown in the following lemma:

**Lemma C.4.** Consider $\Theta := (\mathbf{t}_\ell)_{\ell=1}^L$ where $\mathbf{t}_\ell := (2^{L-1}, 2t, 2t, 2^{L-\ell})$, $\ell = 1, \dots, L$, $\Theta$ is equivalent to the chain of $L$ block butterfly factors of block size $t$ in the sense that there exists two permutation matrices $(\mathbf{P}, \mathbf{Q})$ of size $2^L t \times 2^L t$ such that:

$$ \mathcal{B}_b^{L,t} = \mathbf{P}\mathcal{B}^\Theta\mathbf{Q} := \{\mathbf{P}\mathbf{B}\mathbf{Q} \mid \mathbf{B} \in \mathcal{B}^\Theta\}. $$

In words, the two sets $\mathcal{B}_b^{L,t}$ and $\mathcal{B}^\Theta$ are *expressively equivalent* up to permutation of rows and columns.

*Proof of Lemma C.4.* Consider the permutation matrices:

$$ \mathbf{P}_{\ell} := \mathbf{I}_{2^{\ell-1}} \otimes \mathbf{T}_{\ell} \qquad \mathbf{Q}_{\ell} := \mathbf{I}_{2^{\ell}} \otimes \mathbf{T}_{\ell+1}^{\top}, \quad \ell = 1, \dots, L, \qquad (5) $$

where $\mathbf{T}_\ell$ (whose size is $2^{L-\ell+1}t \times 2^{L-\ell+1}t$) is the permutation matrix corresponding to the permutation:

$$ (1, t+1, \dots, (2^{L-\ell+1}-1)t+1, 2, t+2, \dots, (2^{L-\ell+1}-1)t+2, \dots, t, 2t, \dots, 2^{L-\ell+1}t). $$

The proof relies on the following claim:

$$ \mathbf{P}_{\ell} \tilde{\mathbf{S}}_{\ell} \mathbf{Q}_{\ell} = \mathbf{S}_{\mathbf{t}_{\ell}}, \quad \forall \ell = 1, \dots, L, \qquad (6) $$

which means that applying $\mathbf{P}, \mathbf{Q}$ to the left and right of the $\ell$th block butterfly factors turn it into a $\mathbf{t}_\ell$-butterfly factors. Before proving (6), we explain how we can finish the proof of Lemma C.4 based on (6). Firstly, to prove the inclusion $\mathcal{B}_b^{L,t} \subseteq \mathbf{P}_1^\top \mathcal{B}^\Theta \mathbf{Q}_\ell^\top$, observe that:

$$ \mathbf{Q}_{\ell-1} \mathbf{P}_{\ell} = (\mathbf{I}_{2^{\ell}} \otimes \mathbf{T}_{\ell}^{\top})(\mathbf{I}_{2^{\ell-1}} \otimes \mathbf{T}_{\ell}) = \mathbf{I}_{2^{\ell-1}} \otimes (\mathbf{T}_{\ell}^{\top} \mathbf{T}_{\ell}) = \mathbf{I}_{2^{\ell-1}} \otimes \mathbf{I}_{2^{L-\ell+1}t} = \mathbf{I}_{2^L t}, \qquad (7) $$
<figure><img src="image_8.png" /><figcaption>Figure 8. Boxplot and histogram of the sparsity levels covered by the 600 KS patterns of our benchmark in Section 5.</figcaption></figure>

<table><caption>Table 5. Win rate and median speed-up of KERNEL vs min{BMM, EINSUM, BSR} across sparsity bins, in FP32. The correlation is between log speedup) and log(h) within each bin, where h(b, c) = $\frac{b+c}{bc}$.</caption><thead><tr><th>Sparsity bin</th><th>Win rate</th><th>Median × faster</th><th>Number of patterns</th><th>Corr(log-log)</th></tr></thead><tbody><tr><td>99.48–99.93%</td><td>100.0%</td><td>1.84</td><td>102</td><td>0.91</td></tr><tr><td>98.96–99.48%</td><td>93.1%</td><td>1.41</td><td>101</td><td>0.82</td></tr><tr><td>97.92–98.96%</td><td>79.6%</td><td>1.38</td><td>108</td><td>0.83</td></tr><tr><td>95.83–97.92%</td><td>81.5%</td><td>1.41</td><td>92</td><td>0.80</td></tr><tr><td>87.49–95.83%</td><td>89.5%</td><td>1.30</td><td>114</td><td>0.80</td></tr><tr><td>49.99–87.49%</td><td>86.0%</td><td>1.23</td><td>86</td><td>0.76</td></tr></tbody></table>

### E.3. Measuring the Cost of the Permutation Steps in BMM (Section 3)

**Protocol.** The BMM baseline follows Algorithm 1, which for a pattern (a, b, c, d) performs two global permutations (lines 1 and 3) plus one matrix multiplication with the permuted factor $\tilde{K}$. Because $\tilde{K}$ itself has pattern (ad, b, c, 1), we isolate the permutation overhead by timing:

1. $\Delta t$: the full BMM kernel on the original pattern (a, b, c, d);

2. $\Delta \tilde{t}$: the same kernel restricted to the multiplication step only, i.e., BMM executed on (ad, b, c, 1).

The share of runtime devoted to the two permutations is

$$ \frac{\Delta t - \Delta \tilde{t}}{\Delta t} $$

**Results.** Figure 10 shows the distribution of memory transfer overheads at a fixed batch size of 25088. The patterns included in this plot are of the form (1, b, c, d) with b, c ∈ {48, 64, 96, 128, 192, 256, 384, 512, 768, 1024} such that b = c, b = 4c, or c = 4b, and d ∈ {1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, 96, 128}.

In the *batch-size-first* layout (Figure 10a, reproduced from Figure 4), the overhead of BMM due to the two global permutations in Algorithm 1 increases with the ratio $h(b, c) = \frac{b+c}{bc}$, reaching up to 50%.

The same trend is observed in the *batch-size-last* layout (Figure 10b).

### E.4. Details on min time(KERNEL, BMM, BSR, EINSUM) vs. min time(DENSE, SPARSE) (Section 5.4)

Figure 12 shows that the speedup factor of implementations specialized to the Kronecker-sparsity (KERNEL, BMM, BSR, EINSUM) over the generic DENSE and SPARSE implementations increases with the matrix size $M \times N$. Figure 11 describes the distribution in sizes (the product $MN$) over the benchmark: most of the values of $MN$ are in between $4.2 \times 10^6$ (first quartile) and $6.0 \times 10^8$ (third quartile). We recall that the output and input dimensions are $M = abd$ and $N = acd$,
where we use the identity $(A \otimes B)(C \otimes D) = (AC) \otimes (BD)$ (applicable only when the matrix multiplication AC and BD is well-defined) and $PP^\top = I$ for any permutation matrix. Thus, consider $(\tilde{X}_\ell)_{\ell=1,...,L}$ a sequence of L block butterfly factors of block size t, we have:

$$ \begin{aligned} \tilde{X}_1 \tilde{X}_2 \dots \tilde{X}_L &\stackrel{(7)}{=} (P_1^\top P_1) \tilde{X}_1 (Q_1 P_2) \tilde{X}_2 \dots (Q_{L-1} P_L) \tilde{X}_L (Q_L Q_L^\top) \\ &\stackrel{(6)}{=} P_1^\top \underbrace{(P_1 \tilde{X}_1 Q_1)}_{t_1 \text{ butterfly factor}} \dots \underbrace{(P_2 \tilde{X}_2 Q_2)}_{t_2 \text{ butterfly factor}} \dots \underbrace{(P_L \tilde{X}_L Q_L)}_{t_L \text{ butterfly factor}} Q_L^\top \in P_1^\top B^\Theta Q_L^\top. \end{aligned} $$

Therefore, $B_b^{L,t} \subseteq P_1^\top B^\Theta Q_\ell^\top$. Secondly, to see the other inclusion, from (6), we also have:

$$ S_{\ell}^{b} = P_{\ell}^{\top} S_{t_{\ell}} Q_{\ell}^{\top}, \quad \forall \ell = 1, \dots, L. \qquad (8) $$

Similar to the proof of $B_b^{L,t} \subseteq P_1^\top B^\Theta Q_L^\top$, we can consider $(X_\ell)_{\ell=1,...,L}$ where $X_\ell$ is a $t_\ell$-butterfly factor. We have:

$$ \begin{aligned} P_1^\top X_1 X_2 \dots X_L Q_\ell^\top &\stackrel{(7)}{=} P_1^\top X_1 (P_2 Q_1)^\top X_2 \dots X_{L-1} (P_L Q_{L-1})^\top X_L^b Q_L^\top \\ &\stackrel{(8)}{=} \underbrace{(P_1^\top X_1 Q_1^\top)}_{1\text{st block butterfly factor}} \dots \underbrace{(P_L^\top X_L Q_L^\top)}_{L\text{th block butterfly factor}} \in B_b^{L,t}. \end{aligned} $$

Thus, $P_1^\top B^\Theta Q_\ell^\top \subseteq B_b^{L,t}$. We can conclude that $P_1^\top B^\Theta Q_\ell^\top = B_b^{L,t}$ and the pair of permutation matrices $(P_1^\top, Q_L^\top)$ satisfies Lemma C.4.

Finally, it remains to prove (6). Using the identity: $(A_1 \otimes \dots \otimes A_n)(B_1 \otimes \dots \otimes B_n)(C_1 \otimes \dots \otimes C_n) = (A_1 B_1 C_1) \otimes \dots \otimes (A_n B_n C_n)$ (when $A_\ell B_\ell C_\ell$, $\ell = 1, \dots, n$ makes sense), we have:

$$ \begin{aligned} P_{\ell} S_{\ell}^{b} Q_{\ell} &= (I_{2^{\ell-1}} \otimes T_{\ell})(I_{2^{\ell-1}} \otimes I_{2 \times 2} \otimes I_{2^{L-\ell}} \otimes I_{1 \times t})(I_{2^{\ell}} \otimes T_{\ell+1}^{\top}) \\ &= I_{2^{\ell-1}} \otimes [T_{\ell}(I_{2 \times 2} \otimes I_{2^{L-\ell}} \otimes I_{1 \times t})(I_{2} \otimes T_{\ell+1}^{\top})], \forall \ell = 1, \dots, L. \end{aligned} $$

Thus, it is sufficient to prove that:

$$ T_{\ell}(I_{2 \times 2} \otimes I_{2^{L-\ell}} \otimes I_{1 \times t})(I_{2} \otimes T_{\ell+1}^{\top}) = I_{2t \times 2t} \otimes I_{2^{L-\ell}}, \forall \ell = 1, \dots, L. $$

To show the equation above, we will use an ad-hoc argument. Note that the binary matrix $I_{2 \times 2} \otimes I_{2^{L-\ell}} \otimes I_{1 \times t}$ has exactly $2^{L-\ell}$ different columns (resp. rows). If we partition them into $2^{L-\ell}$ groups $\{1, \dots, 2^{L-\ell}\}$ and label each of them by the partition they belong to, then the columns and rows will have the form:

$$ \underbrace{(1, \dots, 1)}_{t \text{ times}}, \underbrace{2, \dots, 2}_{t \text{ times}}, \dots, \underbrace{2^{L-\ell}, \dots, 2^{L-\ell}}_{t \text{ times}}, \underbrace{1, \dots, 1}_{t \text{ times}}, \underbrace{2, \dots, 2}_{t \text{ times}}, \dots, \underbrace{2^{L-\ell}, \dots, 2^{L-\ell}}_{t \text{ times}}. $$

This procedure of labeling can be visualized as in Figure 18 for a simple case where $L-\ell = 1$ and $t = 3$. To turn $I_{2 \times 2} \otimes I_{2^{L-\ell}} \otimes I_{1 \times t}$ into $I_{2t \times 2t} \otimes I_{2^{L-\ell}}$, the permutation matrices on the left and right (in this case, $T_\ell$ and $I_2 \otimes T_{\ell+1}^\top$) need to permute the columns and rows such that the label becomes:

$$ \underbrace{(1, 2, \dots, 2^{L-\ell})}_{1\text{st}}, \underbrace{1, 2, \dots, 2^{L-\ell}}_{2\text{nd}}, \dots, \underbrace{1, 2, \dots, 2^{L-\ell}}_{2\text{th}}. $$

It can be seen directly that $T_\ell$ permutes the label of the rows correctly. For $I_2 \otimes T_{\ell+1}^\top$, this permutation matrix permutes the columns in the first half and the second half separately (due to the Kronecker product with $I_2$) and each half is permuted by the permutation described in (5). Thanks to this interpretation, it can be seen that $I_2 \otimes T_{\ell+1}^\top$ also permutes the label of the columns correctly as well. This concludes the proof. $\square$

# D Implementations

## D.1 Details on baseline GPU implementations

To keep it short, we only give the code in the case of the batch-size-first memory layout (except for dense and sparse where the codes are small). The case of batch-size-last can simply be obtained by inverting the first and last positions in all tensor reshavings.

**einsum implementation.** This implementation uses tensor contractions with the high-performance einops library. The abcd nonzero entries of the butterfly factor B (Figure 2) are stored in a PyTorch 4D-tensor B_einsum of shape $(a, b, c, d)$. The implementation uses Einstein notations.
<figure><img src="image_9.png" /><figcaption>Figure 9. Speedup factor of KERNEL compared to baselines (BMM, BSR, EINSUM) vs. the heuristic $h(b, c) = \frac{b+c}{bc}$, when fixing the sparsity level at a given value, in FP32. Each subplot corresponds to different sparsity bins involved in Table 5.</figcaption></figure>
<figure><img src="image_18.png" /><figcaption>Figure 18: Illustration for the proof of (6). Green and white squares represent the 1s and 0s of the binary matrices respectively.</figcaption></figure>

```python
def butterfly_einsum(X_bsf, B_einsum):
    X_perm = einops.rearrange(X_bsf, "... (a c d) -> ... a c d", a=a,
                             c=c, d=d)
    Y_perm = einops.einsum(X_perm, B_einsum, "... a c d, a b c d ->
                             ... a b d")
    Y_bsf = einops.rearrange(Y_perm, "... a b d-> ... (a b d)")
    return Y_bsf
```

The second line of this code does at the same time all the matrix multiplications $Y[:, row] \leftarrow X[:, col]B^{\top}[col, row]$ for all the pairs $(row, col)$ in Algorithm 1.

**bsr implementation.** This is an implementation of Algorithm 2 using the high-performance Block compressed Sparse Row (BSR) PyTorch library. The matrix $\tilde{B}$ is stored as a tensor $B_{bsr}$ stored in the BSR format.

```python
def butterfly_bsr(X_bsf, B_bsr):
    batch_size = X_bsf.shape[0]
    X_perm = (
        X_bsf.view(batch_size, a, c, d)
        .transpose(-1, -2)
        .reshape(batch_size, a * c * d)
    )
    Y_perm = torch.nn(functional.linear(
        X_perm, B_bsr
    )
    Y_bsf = (
        Y_perm.view(batch_size, a, d, b)
        .transpose(-1, -2)
        .reshape(batch_size, a * b * d)
    )
    return Y_bsf
```

**bmm implementation.** This is an implementation of Algorithm 2 using the high-performance Block compressed Sparse Row (BSR) PyTorch library. The matrix $\tilde{B}$ is stored as a tensor $B_{bmm}$ stored in the BSR format. This implementation using `torch.bmm`, which is based on high-performance batched matrix multiplication NVIDIA routines. The non-zero entries of $\tilde{B}$ are stored in a four-dimensional PyTorch tensor $B_{bmm}$ of shape $(a * d, b, c)$.

```python
def butterfly_bmm(X_bsf, B_bmm):
    batch_size = X_bsf.shape[0]
    X_perm = (
```
<figure><img src="image_10.png" /><figcaption>Figure 10. Estimated share of runtime devoted to the permutation steps in the BMM baseline for many patterns $\pi = (a, b, c, d)$, in FP32. Patterns are grouped by the value of $h(b, c) = \frac{b+c}{bc}$ and the distribution within each group is shown as a boxplot.</figcaption></figure>

respectively, for a Kronecker-sparse matrix with pattern $\pi = (a, b, c, d)$.

<figure><img src="image_11.png" /><figcaption>Figure 11. Distribution of matrix sizes in the benchmark.</figcaption></figure>

### E.5. Details on time(BMM) vs. min time(BSR, EINSUM) (Section 5.4)

Figure 13 shows that for a sufficient large matrix size $M \times N$, we always have $time(BMM) < min time(BSR, EINSUM)$, i.e., the BMM implementation is the most efficient among all baseline implementations (BMM, EINSUM, BSR).
```python
X_bsf.view(batch_size, a, c, d)
 transpose(-1, -2)
 reshape(batch_size, a * d, c)
 contiguous()
 transpose(0, 1)
)

Y_perm = torch.empty(batch_size, a * d, b, device=x device, dtype=
    x . type). transpose(0, 1)
Y_perm = torch.bmm(X_perm, B_bmm. transpose(-1, -2))
Y_bsf = (
    Y_perm. transpose(0, 1)
    . reshape(batch_size, a, d, b)
    . transpose(-1, -2)
    . reshape(batch_size, a * b * d)
)
return
```

**dense implementation.** This ignores the sparsity of the butterfly factor $\mathbf{B}$, that is stored as a dense matrix in a 2d-tensor $\mathbf{B}_{\text{dense}}$.

*batch-size-first:* `torch.nn-functional(linear(X_bsf, B_dense))`

*batch-size-last:* `torch.matmul(B_dense, X_bsl)`

The implementation in *batch-size-first* is the default PyTorch implementation of a forward pass of a linear layer. For *batch-size-last*, we had to choose an implementation since PyTorch uses *batch-size-first* by default. We made our choice based on a small benchmark of different alternatives.

**sparse implementation.** This exploits the sparsity of the butterfly factor $\mathbf{B}$ but not its structure (recall that the support are not arbitrary, they are structured since they must be expressed as Kronecker products, see Definition 2.1).

*batch-size-first:* `torch.nn-functional(linear(X_bsf, B_csr))`

*batch-size-last:* `torch.matmul(B_csr, X_bsl)`

## D.2 Details on the kernel implementation

**Classical optimizations that we build upon.** The proposed implementation use *vectorization* as soon as an operation can be vectorized. Concretely, the float4 and half2 vector types are used to mutualize read/write operations [NVIDIA, 2023b,a, 2024, Boehm, 2022]. An *epilogue* [NVIDIA, 2023a] is also implemented to avoid writing in global memory in a disorganized way. Indeed, after having accumulated the output in registers, each thread has specific rows and columns of the output to write to global memory, and may finish its computation before the others. To avoid that, the *epilogue* starts to write in the shared memory, in a disorganized way, and then organize the writing from shared to global memory. Another implemented optimization is *double buffering* NVIDIA [2023b,a], Boehm [2022], Li et al. [2019]: a thread block is always both computing the output of a tile, and loading the next tile from global to shared memory. This allows us to hide some latency that arises when loading from the global memory.

Note that as with any CUDA kernel, the constants (such as the number of threads) need to be tailored to each specific case of use —here, each butterfly sparsity pattern $\pi = (a, b, c, d)$— and to each GPU.
<figure><img src="image_12.png" /><figcaption>Figure 12. Speedup factor of min time(KERNEL, BMM, BSR, EINSUM) compared to min time(DENSE, SPARSE) as a function of the matrix size $M \times N$.</figcaption></figure>

<figure><img src="image_13.png" /><figcaption>Figure 13. Speedup factor of time(BMM) compared to min time(EINSUM, BSR) as a function of the matrix size $M \times N$.</figcaption></figure>

### E.6. Details on the Impact of the Memory Layout (Section 5.4)

Figure 14a shows how the memory layout (*batch-size-first* vs. *batch-size-last*) impacts the runtime of each implementation. In particular, it shows that * KERNEL* achieves a median speedup of approximately ×2 when switching from *batch-size-first* to *batch-size-last*.

<figure><img src="image_14.png" /><figcaption>(a) *float-precision*.</figcaption></figure>

<figure><img src="image_15.png" /><figcaption>(b) *half-precision*.</figcaption></figure>

<caption>Figure 14. Boxplots of the ratio *time of batch-size-first* / *time of batch-size-last*, in FP32.</caption>
Table 6 reports the percentage of KS patterns for which *KERNEL* outperforms all baseline implementations, in either the *batch-size-first* or *batch-size-last* memory layout. Even when all implementations are restricted to the *batch-size-first* layout, *KCHOOL* still improves over all baselines on 20% of the tested patterns, despite relying on some non-contiguous memory accesses (Section 4.2).

Table 6 also confirms that among existing baselines (excluding the new *KCHOOL*), *BMM* is the fastest across both layouts, outperforming *EINSUM*, *BSR*, *DENSE*, and *SPARSE*.

<table><caption>Table 6. Percentage out of 600 patterns (a, b, c, d) where *algol* is faster than the *algol2* (denoted by time(algol) < time(algol2)), and the median acceleration factor in such cases (that is, the median ratio time of *algol2* / time of *algol* for a given memory layout), in FP32.</caption><thead><tr><th rowspan="2">Comparison</th><th colspan="2">Batch-size-first</th><th colspan="2">Batch-size-last</th></tr><tr><th>Win rate</th><th>Median × faster</th><th>Win rate</th><th>Median × faster</th></tr></thead><tbody><tr><td>min{KCHOOL, BMM, EINSUM, BSR} &lt; min{DENSE, SPARSE}</td><td>97.61%</td><td>×18.50</td><td>98.09%</td><td>×6.53</td></tr><tr><td>BMM &lt; min{EINSUM, BSR, DENSE, SPARSE}</td><td>94.1%</td><td>×1.44</td><td>89.47%</td><td>×1.66</td></tr><tr><td>KCHOOL &lt; min{BMM, EINSUM, BSR, DENSE, SPARSE}</td><td>20.41%</td><td>×1.26</td><td>85.49%</td><td>×1.39</td></tr></tbody></table>

## E.7. Time Spent in Linear Layers in Vision Transformers (Section 6)

This section provides a numerical lower bound estimate of the time spent in fully-connected layers within Vision Transformers (ViTs).

**Results.** Table 7 shows the proportion of computation time dedicated solely to the linear layers in the feed-forward network (FFN) modules across various ViT architectures. This proportion ranges from 31% to 53% in *half-precision* and from 46% to 61% in *float-precision*. The fraction increases with model size, indicating that a significant part of ViT inference is spent in fully-connected layers. Note that our measurement excludes the linear layers in the multi-head attention modules, so these values represent a *lower bound* on the actual time spent in *all* linear layers of the Transformer.

<table><caption>Table 7. Median execution times (ms) for the forward pass in various ViTs, and for an MLP composed solely of the linear layers from the feed-forward network modules of the ViTs. The table also reports the proportion of the latter over the former. FP16 corresponds to *half-precision*, and FP32 to *float-precision*.</caption><thead><tr><th rowspan="2">ARCHITECTURE</th><th colspan="2">FP16 (s)</th><th colspan="2">FP32 (s)</th></tr><tr><th>COMPLETE</th><th>LINEAR IN FFNS</th><th>COMPLETE</th><th>LINEAR IN FFNS</th></tr></thead><tbody><tr><td>ViT-S/16</td><td>0.014</td><td>0.0046 (31%)</td><td>0.090</td><td>0.04 (46%)</td></tr><tr><td>ViT-B/16</td><td>0.036</td><td>0.015 (42%)</td><td>0.30</td><td>0.16 (54%)</td></tr><tr><td>ViT-L/16</td><td>0.11</td><td>0.050 (46%)</td><td>1.0</td><td>0.58 (58%)</td></tr><tr><td>ViT-H/14</td><td>0.31</td><td>0.16 (53%)</td><td>2.6</td><td>1.6 (61%)</td></tr></tbody></table>

**Details on the Estimation.** A Transformer architecture consists of a sequence of blocks, each containing a multi-head attention module and a feed-forward network module. The feed-forward network (FFN) is a two-layer MLP (without biases) with an intermediate non-linear activation. Table 7 reports the cumulative execution time of all linear layers in the FFNS, computed sequentially. This is then compared with the total forward pass time of the full ViT model. Since we do not include linear layers from the attention modules, the reported ratios represent only a lower bound on the total time spent in linear operations.

**Experimental Settings.** The ViT-S/16 architecture follows (Zhai et al., 2022), while ViT-B/16, ViT-L/16, and ViT-H/14 follow (Dosovitskiy et al., 2020). Input images are of size 224 × 224. For *float-precision*, we use the PyTorch ViT implementations of Wang (2024b). For *half-precision*, the Transformer uses FlashAttention (Dao et al., 2022c) to compute scaled dot-product attention, as in (Wang, 2024a). The FFN-only MLPs are implemented using `torch.nn.Sequenceal` and `torch.nn.Linear`. All experiments were run on a single NVIDIA A100-40GB GPU paired with an AMD EPYC 7742 64-Core Processor. Execution times were measured using `torch.utils.benchmark.Timer`, with a batch size of 128.
## E.8. Details on End-to-End Transformers Inference Acceleration (Section 6)

**Why Do We Focus More on Transformers Rather Than, e.g., Convolutional Neural Networks?** In Section 6, we argued that injecting KS in linear layers is straightforward—making architectures like Transformers a natural fit.

By contrast, applying KS to convolutional neural networks is more challenging. There are at least two ways to inject Kronecker sparsity in convolutional layers, none of which we found convincing enough to be considered in this paper.

The first approach—replacing the convolution kernel with a Kronecker-sparse matrix—is limited by the small size of typical convolution kernels (e.g., 7×7), leaving little room for meaningful gains.

The second approach—recasting convolutions as matrix multiplications with butterfly-structured weights, as in (Lin et al., 2021)—introduces significant overhead in the form of input/output folding and unfolding operations. In our view, this makes the approach impractical for efficient inference.

For these reasons, we focus on fully connected layers, where KS integration is both simple and impactful.

**Selected Kronecker-Sparse Patterns for ViT-S/16.** In our ViT-S/16 experiments, we replace the dense weight matrices in all fully-connected layers—namely, the projection matrices in the self-attention blocks and the linear layers in the feed-forward network blocks—with products of two Kronecker-sparse matrices $K_1K_2$, where each $K_i$ follows a specific sparsity pattern $\pi_i$ as defined in Definition 2.1. The patterns used are:

* $(\pi_1, \pi_2) = ((1, 192, 48, 2), (2, 48, 192, 1))$ for square matrices of size $N \times N$ (projections in self-attention blocks),

* (1, 768, 192, 2), (6, 64, 64, 1) for $4N \times N$ matrices (UP projection matrices),

* (6, 64, 256, 1), (1, 128, 128, 3) for $N \times 4N$ matrices (DOWN projection matrices).

These patterns were selected based on two criteria: (i) the total number of nonzero entries across both KS factors is at least 75% of the original dense matrix size ($a_1b_1c_1d_1 + a_2b_2c_2d_2 \ge 0.75MN$), ensuring sufficient layer expressivity to raise hope that it could yield competitive task accuracy compared to the dense implementation, if it were fine-tuned or trained from scratch on a given task; (ii) the patterns yield a high value of the heuristic $h(b, c) = \frac{b+c}{bc}$, which favors the performance of our fused KERNEL implementation.

**Additional Results for ViT-S/16.** Table 8 presents the latency reduction achieved by replacing standard Linear layers with Kronecker-sparse layers using our fused KERNEL implementation, as compared to the BMM baseline. These results, measured on an Nvidia A100 40GB GPU, confirm that the gains observed at the layer level carry over to larger functional blocks in the model.

<table><caption>Table 8. Acceleration of submodules of a ViT-S/16 using Kronecker-sparse matrices, in FP32.</caption><thead><tr><th></th><th>time(BMM)<br>time(fully-connected)</th><th>time(KERNEL)<br>time(fully-connected)</th></tr></thead><tbody><tr><td>Linear N × N</td><td>0.82</td><td><strong>0.50</strong></td></tr><tr><td>Linear N × N + bias</td><td>0.97</td><td><strong>0.66</strong></td></tr><tr><td>Linear 4N × N</td><td>0.80</td><td><strong>0.78</strong></td></tr><tr><td>Linear 4N × N + bias</td><td>0.93</td><td><strong>0.90</strong></td></tr><tr><td>Linear N × 4N</td><td>0.91</td><td><strong>0.58</strong></td></tr><tr><td>Linear N × 4N + bias</td><td>0.94</td><td><strong>0.61</strong></td></tr><tr><td>FFN (2 × Linear+GELU+LN)</td><td>0.91</td><td><strong>0.77</strong></td></tr><tr><td>MHA (QKV + proj)</td><td>0.87</td><td><strong>0.79</strong></td></tr><tr><td>Transformer block</td><td>0.90</td><td><strong>0.78</strong></td></tr><tr><td><strong>ViT-S/16 full</strong></td><td>0.89</td><td><strong>0.78</strong></td></tr></tbody></table>

**Selected Kronecker-Sparse Patterns for GPT-2 Medium.** For GPT-2 Medium, we inject Kronecker sparsity into the down-projection linear layer within the feed-forward block. The layer in question has dimensions $1024 \times 4096$ (hidden size $d = 1024$). The applied KS patterns are:

$$ \pi_1 = (1, 64, 256, 16), \quad \pi_2 = (64, 64, 64, 1), $$
yielding the factorization $W = K_1K_2$. These patterns are selected following the same rationale as for ViT-S/16: they are expressive enough (nonzero density above 75%) and yield high values of the heuristic $h(b, c)$, which predicts good performance with *KERNEL*.

**Inference Setup.** Inference is conducted in PyTorch using the *batch-size-first* memory layout (batch-size-first), which is the framework’s default. All operations aside from the KS layers use standard PyTorch kernels. For ViT-S/16, we use an Nvidia A100 GPU with a batch of 128 image token sequences, each containing 196 tokens. For GPT-2 Medium, we use an Nvidia RTX6000 GPU under similar batch settings.

## E.9. Benchmark Results in Half-Precision

We now reproduce the benchmark from Section 5 in *half-precision*. The *half-precision* equivalents of our key *float-precision* results—Tables 2 and 6, Figure 6a, Figure 10, Figure 12, Figure 13, and Figure 14a—are respectively shown in Tables 9 and 10, Figure 18, Figure 17, Figure 15, Figure 16, and Figure 14b.

As in the FP32 benchmark, Figure 18 reports results only for patterns where at least one KS-aware method (KERNEL, BMM, EINSUM, BSR) is faster than the generic baselines (DENSE, SPARSE). This filters out roughly 14.3% of the patterns and retains 85.7% of the benchmarked grid in *half-precision* (see Table 9).

Overall, KERNEL shows fewer speedups compared to the FP32 setting, raising two possible interpretations: either existing baselines have been better optimized for half-precision and benefit more from tensor cores, or there remains untapped room for optimization in our kernel for *half-precision*.

1) Figure 15 and Table 10 confirm that, as in FP32, the implementations tailored to KS matrices (KERNEL, BMM, EINSUM, BSR) still largely outperform the generic baselines (DENSE, SPARSE).

2) Figure 16 and Table 10 show that among the existing KS-aware implementations (BMM, EINSUM, BSR), BMM remains the fastest when the matrix size is sufficiently large, regardless of the memory layout.

3) Figure 17 highlights that in *half-precision*, the relative cost of the permutation steps in BMM grows even larger than in *float-precision*—ranging from 40% to 80%, compared to 10% to 50% in *float-precision*—further motivating efforts to eliminate these costly memory operations.

4) Figure 18 shows that the new KERNEL can still yield up to ×2 speedups, but only for KS patterns where a large portion of the runtime is spent on permutations—those with high values of the heuristic $h(b, c) = \frac{b+c}{bc} \ge 0.02$. This contrasts with the FP32 case, where improvements already appear from $h(b, c) \ge 0.004$ (see Figure 6a). Consequently, KERNEL improves over existing methods on 38% of the patterns in *half-precision* (Table 9), compared to 85% in *float-precision* (Table 2).

5) Table 10 shows that the current KERNEL remains competitive in *batch-size-last* under *half-precision*, but falls behind in *batch-size-first*. This shows that memory layout plays a key role in making Kronecker sparsity efficient on GPUs, and suggests further study on how layout affects other neural network operations.

These findings suggest that while KERNEL remains competitive in *half-precision*, its performance gains are currently more limited. This highlights potential for further optimization, particularly to better leverage half-precision hardware features such as tensor cores.
<figure><img src="image_16.png" /><figcaption>**Figure 15.** Speedup factor of min time(KERNEL, BMM, BSR, EINSUM) compared to min time(DENSE, SPARSE) vs. the matrix size $M \times N$. Experiments are carried in *half-precision*.</figcaption></figure>

<figure><img src="image_17.png" /><figcaption>**Figure 16.** Speedup factor of time(BMM) compared to min time(EINSUM, BSR) vs. the matrix size $M \times N$. Experiments are carried in *half-precision*.</figcaption></figure>

<table><caption>**Table 9.** Percentage out of 600 patterns (a, b, c, d) where *algol* is faster than the *algo2* in *half-precision* (denoted by time(algol) < time(algo2)), and the median acceleration factor in such cases (that is, the median ratio time of algo2 / time of algol). For each implementation, we take the minimum time between the *batch-size-first* and the *batch-size-last* memory layout. Experiments are carried in *half-precision*.</caption><thead><tr><th>min time<br>(<strong>KERNEL</strong><br><strong>BMM</strong><br><strong>EINSUM</strong><br><strong>BSR</strong>) &lt; min time<br>(<strong>DENSE</strong><br><strong>SPARSE</strong>)</th><th>time(BMM) &lt; min time<br>(<strong>EINSUM</strong><br><strong>BSR</strong><br><strong>DENSE</strong><br><strong>SPARSE</strong>)</th><th>time(KERNEL) &lt; min time<br>(<strong>BMM</strong><br><strong>EINSUM</strong><br><strong>BSR</strong><br><strong>DENSE</strong><br><strong>SPARSE</strong>)</th></tr></thead><tbody><tr><td>85.7% (×8.14)</td><td>80.7% (×1.52)</td><td>37.8% (×1.51)</td></tr></tbody></table>

<table><caption>**Table 10.** Across 600 patterns, this table reports how often each algorithm is the fastest, along with the corresponding median speedup, under each memory layout (*batch-size-first* and *batch-size-last*) in FP16.</caption><thead><tr><th rowspan="2">Comparison</th><th colspan="2">Batch-size-first</th><th colspan="2">Batch-size-last</th></tr><tr><th>Win rate</th><th>Median × faster</th><th>Win rate</th><th>Median × faster</th></tr></thead><tbody><tr><td>min{<strong>KERNEL</strong>, <strong>BMM</strong>, <strong>EINSUM</strong>, <strong>BSR</strong>} &lt; min{<strong>DENSE</strong>, <strong>SPARSE</strong>}</td><td>83.18 %</td><td>×6.31</td><td>85.34 %</td><td>×8.14</td></tr><tr><td><strong>BMM</strong> &lt; min{<strong>EINSUM</strong>, <strong>BSR</strong>, <strong>DENSE</strong>, <strong>SPARSE</strong>}</td><td>82.25 %</td><td>×1.54</td><td>79.48 %</td><td>×2.65</td></tr><tr><td><strong>KERNEL</strong> &lt; min{<strong>BMM</strong>, <strong>EINSUM</strong>, <strong>BSR</strong>, <strong>DENSE</strong>, <strong>SPARSE</strong>}</td><td>0.0 %</td><td>-</td><td>41.78 %</td><td>×1.56</td></tr></tbody></table>
<figure><img src="image_18.png" /><figcaption>Figure 17. Estimated relative time spent on memory rewrites in BMM for the multiplication with $K \in K_\pi$, for several $\pi = (a, b, c, d)$. We regroup patterns by their value of $(b+c)/(bc)$, and plot a boxplot to summarize the corresponding measurements. Experiments are carried in *half-precision*.</figcaption></figure>

<figure><img src="image_19.png" /><figcaption>Figure 18. Speedup factor of KERNEL compared to min(BMM, EINSUM, BSR) in *half-precision*. For each implementation, we take the minimum time between the *batch-size-first* and the *batch-size-last* memory layout. We regroup the $(a, b, c, d)$ patterns by their value of $(b+c)/(bc)$, and use a boxplot to summarize the corresponding measurements. Experiments are carried in *half-precision*.</figcaption></figure>