Fast Inference with Kronecker-Sparse Matrices
Authors: Antoine Gonon
Léon Zheng
Pascal Carrivain
TUNG QUOC LE
Venue: ICML 2025 poster (2025)
Keywords: GPU matrix multiplication
Kronecker-sparse
Primary area: deep_learning->algorithms
TL;DR: Improves GPU matrix multiplication algorithms for Kronecker-sparse matrices

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 x1.4 and lowers energy consumption by 15%. A simple heuristic based on KS pattern parameters predicts when our method outperforms existing ones. We release all code at [github.com/PascalCarrivain/ksmm](https://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.

Fast Inference with Kronecker-Sparse Matrices
Antoine Gonon Léon Zheng
134
Pascal Carrivain Quoe-Tung Le
Abstract
2022a:b). Among the many structures explored, the Butler- ftyMonarch 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 Kro- necker product of identities and all-ones blocks. We call them Kronecker-sparse (KS) matrices (Definition 2.1), and denote by (a, c, d) the pattern of KS matrix with support L @ T1xe LIg, where the Kronecker product, I, is the n identity matrix, and Laxe is matrix of size b x full of ones. Factorizing dense matrix W into KS matrices K ... KL can reduce inference cost, as the se- quential matrix-vector multiplications x + K₁ .... K2 can be more efficient than computing the dense multiplica- tion x + Wx. The canonical example the fast Fourier transform (FFT): the N N discrete Fourier transform (DFT) matrix can be factored into L. (N) KS factors of size N x N, with at most ((N) nonzero entries factor (Dao et al., 2019). This reduces the overall cost of comput- ing the DFT from O(N2) (via naive multiplication with the original matrix W) to O(N log, .N) (through sequential multiplication with the KS factors Ky); see Figure 1.
Kronecker-sparss (KS) matrices-whort supports are Kronecker products of identity and all-ones uldeks-underpin the structure of Butterfly and Monarch matrices and offer the promise of more efficient models. However, existing GPU ker- nels for KS matrix multiplication suffer from high data movement costs, with up to 50 % of time spent on memory-bound tensor permuta- tions. We propose fused, ouput-stationary GPU kernel that eliminates these overheads, re- ducing global memory traffic threefold. Across 600 KS patterns, our kernel achieves in FP32 a median speedup of x 1.4 and lowers energy con- sumption by 15 %. A simple heuristic based on KS pattern parameters predicts when our method outperforms existing ones. We release all code' at wttts//o//aaaaaaaiiiiii
in- KSLinear layer.
cluding a and demonstrate in FP32 end-to-end latency re- ductions of up to 22 % in ViT-S/16 and 16 % in GPT-2 medium.
1. Introduction
W
Matrix multiplications dominate the time and energy bud- gets of modem-day networks involving large linear layers. For instance, single forward of a Vision Transformer (ViT- H/14) spends ~0 % of its runtime in linear layers (Ap- pendix E.7). Using structured sparse factors in place of dense weights-eite either learned from scratch or introduced after training-off a compelling way to reduce mem- ory usage and inference FLOPs (Dao et al., 2019; 2020:
K1
K
K3
K,
Figure 1. FFT fast Kronecker-sparse transform. The dense DFT matrix W can he factored as W = K ... K2 (up to a column permutation), where each KS factor K, has sup- port 1 @ 12xz @ L This factorization underlies the O(N log N) efficiency of the FFT algorithm.
"Equal contribution ENS Lyon, CNRS, Inria, Université Claude Bernard Lyon 1, LIP, UMR 5668, 69342, Lyon cedex 07, France 'Institute of Mathematics, EPFL, Lausanne, Switzer- land vvlleo.i,, Paris, France 'Huawei Lagrange Mathematics and Computing Research Center, Paris, France 'Toulosee School of Economics, Toulouse, France. Correspondence to: Antoine Gonon <antoine.gonon @ eptl.ch>.
Scope of This Work: Focus on Inference. This work fo- cuses on efficiently computing Y XKT, where e RM is a fixed Kronecker-sparse matrix and X R is dense matrix representing a batch of inputs. We assume K can be fully preprocessed offline, but inputs and outputs must follow fixed, dense Mayout-PyTore PyTorch's 's default in our case case-which cannot be altered. This reflects standard inference scenarios, where models are already trained and deployed, and data arrives a prescribed format. Optimiz- ing this phase is essential, as inference accounts for over % of large-scale MI. costs (HPCwire, 2019; Barr, 2019).
Proceedings of the 4nnd International Conference on Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
Also integrated into the signal-processing oriented package lazy1 nop (Appendix D.1).

---

Fast Inference with Kronecker-Sparse Matrices
Cost of Memory Operations. The fastest current GPU implementations for Kronecker-sparse multiplica- tion-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 in- put to ensure that specific elements are stored contiguously in memory; (ii) invoke a perhhperformacce multiplication routine, such general matrix multiplication (GEMM); and (iii) apply an inverse permutation to restore the output to its original layout.
the time they spend on memory permutations. Section 4 presents the fused kernel and its analysis. Section bench- marks 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 restrictions the values of those entries. The support is defined by a Kronecker product of three simple mnatrices-identity all-ones, and identity-which forms the common structural core of the Butterfly, Monarch, Kalei- doscope, and related layers (Dao et al., 2019; 2022a;b; Lin et al., 2021; Fu et al., 2023).
However, our profiling (in Section 3 below) reveals that these two permutations- before and after the multiplica- tion-along with the associated memory reads and writes, can account for up to 50 % of the total runtime (Figure 4). Notably, find that the overhead from these data move- ment operations increases with the ratio h(b, c) be simple expression that we motivate analytically and validate empirically across 600 KS patterns in this paper.
Definition 2.1 (KS pattern and matrix). A KS pattern is an integer 4-tuple π = (a, b, c, d). Its corresponding support is S I, < 1axe s La (see Figure 2). matrix K E Raiddaed is π ---ronncker--parse if supp(K) S supp(S,). where supp(-) is the set of matrix indices corresponding to nonzero entries. We write K E K
Contributions. We show that the data movement costs caused by the two permutations are not unavoidable. By in- troducing a new ouput-stationnry tiling scheme for KS ma- trix multiplication, we fuse the three GPU srinnlsspermuute. GEMM, and inverse permute into a single kernel, elimi- nating two global memory passes. Our main contributions are as follows:
book dingonal matrix with blockes
sub blockbookk
x. diagonal matrix
(i) Time & Energy Benchmark. We introduce the first large-scale public benchmark of KS matrix multiplica- tion on GPU, spanning 6 orders of magnitude in matrix sizes, several sparsity levels, and supporting both single (FP32) and half (FP16) floating- -pintt formats. Our re- sults confirm that the proportion of the runtime spent at memory rewritings grows with h(b, c)
Figure 2. A h K--Krneckee-sppe sparse matrix with π (a, b, d) is block-diagonal matrix with a blocks, where each block itself a block matrix composed by b x diagonal matrices of size d x 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.
bic
bc
iii) Single-Kernel Implementation. We propose an outpu-- stationary tiling that reduces the cost of memory opera- tions. We release publicly our kernel implementation in CUDA and OpenCL, plus a drop-in KSLineax layer for PyTorch. Our kernel achieves a median speedup of x1.4, and median energy reduction of 15 % in FP32.
The mask S has shape abd x acd and abed nonzeros, so the density (of nonzero entries) is 1/(ad). Increasing either a or d makes the factor sparser. See Appendix B for:
(iii) Design Heuristic. For a KS pattern (a, b, c, d), we em- pirically show that the ratio h(b, c) bic predicts la- tency, while d h(b, c) predicts energy, giving practi- tioners a one-line line rule for pattern selection.
discussion on where KS fits in the wider sparsity land- scape, with in particular Table 4 showing how popular structured layers in neural networks map to concrete KS patterns;
(iv) End-to-End Impact. We demonstrate that injecting KS sparsity in linear layers of Transformers such as ViT-
links between KS matrix multiplication, sparse 3D convo- lutions, and sparse-tensor compilers:
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.
a brief overview of techniques with provable guarantees for approximating dense matrix W by KS product K Note that this approximation problem lies outside the scope of this paper-we assume fixed KS matrix K is given and focus on efficient multiplication with dense input batches X; see Section 1.
Outline. Section formalizes KS matrices and reviews existing GPU algorithms on PyTorch. Section 3 quantifies
2

---

Fast Inference with Kronecker-Sparse Matrices
2.1. Previous PyTorch GPU Implementations for KS
cost: it requires two full-tensor permutations in Steps (1) and (3), which we find to consume non-negligible portion of the runtime (Section 3). To address this, we introduce new mathematically equivalent algorithm that avoids these permutations entirlly (Section 4).
The fastest publicly available PyTorch implementations op- timized 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 K, pre-permuted, block-diagonal version of K; (3) permute the output back to its original layout.
2.2. Reminder on Memory-Layout
PyTorch uses row-major storage for tensors. Placing the batch dimension first (bulch-size-fissf) ensures that each sample is stored contiguously in memory; placing it last (batch-size-las)) instead makes the features contiguous. The historical default most machine learning pipelines (in- cluding PyTorrch) 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-late performs better for KS matrices. To support both conventions, we provide a Py KSLinear layer compatible with both layouts.
Algorithm Permutation-baat based KS matmul
Require: KS pattern π (a, b, C, d), inputs X E T Bxaed, permutation matrices P, Q (as in Appendix C.I)
pre-permuted KS matrix K := =P
Ensure: outputs Y : -K E Bxabd
1: X XQ T (permute columns)
2: Y X
matrix multiply]
3: Y < YPT
{permute columns back)
Any pair of permutation matrices P and Q in Algo- rithm 1 yields the correct KS matrix multiplication re- sult. Indeed, any permutation matrices and Q satisfy PPT Q I, so this ensures that the multiplica- tion of batch of inputs X Bxedd with KS factor K Rabdxacd is equal to:
3. Cost of the Permutations in Baseline Implementations
As we show in Section 5, BMM is the fastest among exist- ing implementations of KS matrix multiplication. It im- plements Algorithm 1 directly, with two explicit permuta- tions-lines 1 and 3-that move data between global mem- ory and registers before and after the multiplication with the block-diagone matrix K. Figure (left) illustrates the data flow in this baseline.
Y = XQ PP = PT.
Existing implementations use a specific pair of permuta- tions (P, Q) to achieve high efficiency. These are chosen so that the permuted matrix K == PTPP KQ becomes block- diagonal with dense sub-blocks (details Appendix C.1), enabling fast parallel multiplication in Step (2) using stan- dard routines each dense sub-block.
In this section, we empirically demonstrate that these two permutations can account for up to 50 % of the total runtime, motivating the design of new tiling strategy that avoids this overhead (Section 4).
The main differences between existing implementations of Algorithm 1 lie in how they perform Step (2): ei- ther 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 direct tensor contraction baseline tailored to KS struc- ture EINSUM EINSUM-inspired inspired by Dao et al. (2022b). Implemen- tation details are provided in Appendix A.
3.1. Why Data Transfers Matter
GPU memory management plays critical role in perfor- mance optimization. Memory on GPU is organized hier- archically: 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 mem- ory. When a kernel executed, each GPU thread typi- cally loads data from global memory into registers, per- forms register-level computations, and writes results back to global memory (Figure 3). Because memory access often becomes performance bottleneck, minimizing data move- ment between global memory, shared memory, and registers is essential for an efficient implementation (NVIDIA, 2024, Section 5.3).
Table 1. Two PyTorch realizations of Algorithm 1.
BMM
BSR
K storage
(ad, b, c) tensor (abd, aod) in BSR torch. reshape
Permute Q (line 1) Multiply step (line 2)
torch. lomm linear
Permute P (line 3)
torch. reshape
We name the two memory layouts batch-size-firs and batch- sice-last, by analogy with PyTorch's 's channels-lan optimization. which moves the channels dimension to the end for convolutions.
The convenience of reducing the problem to a matrix mul- tiplication with the block-diagonal matrix K comes at a
3

---

Fast Inference with Kronecker-Sparse Matrices
bmm implementation
kernel implementatiio
Global memory
Xuf (topleil ac11
Xeeen phapedas, 9, shasetal,,
Kmmm
Y_bsf Shopoi8 abst
X bsT
K. Kernel shapeled
Y_bsf stapulB. abst
shapolad
srapett
Kronecker adapted reading
Kronecker adapted writing
Shared memory
Xad
16:00
IRV3
18.00
Registers
Matrix multiplication
Permutation
multiplication
Permutation
Figure 3. Data flow in BMM (three global passes) (Dao et al., 2022h) versus our fused KERNEL. (one pass).
3.2. Cost of Data Transfers in Baseline Implementations
4. Novel Tiling for KS Multiplication with Reduced Memory Transfers
We focus BMM will find out it is the fastest existing baseline. To isolate the cost of the memory rewriting op- erations corresponding to the permutations in Algorithm (lines and 3), we compare the runtime of BMM with a modified version in which these permutations are simply re- moved. Figure (y-axis) shows that across 600 KS patterns, permutations consume up 45 % of the time", and the share grows with h6b, e) b+c, we provide more explanation about this ratio in Section 4.
Baseline implementations call three separate GPU ker- nels
every call round-trips through global memory (Figure 3, left).
A naive approach prevent such back and forths would be to directly try merge these three different kernels into 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/snnchre nization across thread blocks). Thus, if a thread assigned tiles (submatrices) on which to apply three operations SULUTEEMMUTPLP-PPEEEA
Rei ative time in memory
its workload cannot de- pend the results obtained by 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.
0.5
rewritings bmm
0.3
0.2
To this end, we introduce in Section 4.1 a novel, mathemat- ically equivalent reformulation of the permutation-ban algorithm (Algorithm 1), formalized Algorithm 2. This reformulation leads to new tiling strategy that enables us to implement the multiplication in single GPU kernel, as described in Section 4.2. We then perform an analytical comparison of memory operations between our implemen- tation and existing baselines (Section 4.3), and derive a heuristic that predicts when our kernel is expected to re- duce the cost of the permutation- memory operations of Algorithm 1. We validate this prediction empirically in Section 5.
1-1
0.1
0.002
0.003:
0.004
0.003
0.006
7,000
0.008
2,000
0.1
0.02
0.03
0.00
(bc)
Figure Share runtime spent only on permutations BMM
(FP32, BSF layout). Each box aggregates patterns with the same
h(b, c)
bic
be
Take-Away: Reducing the cost of memory transfers is a promising direction for improving the performance of Kronecker-sparss matrix multiplication. This the main objective of our implementation, described in the fol- lowing section.
4.1. New Tiling Strategy
We introduce new tiling strategy for KS matrix multipli- cation, to reduce input/output (I/O) transfers between the different memory levels. Tiling consists of partitioning the
Regardless of the memory layout convention, batch-size-fisst or batch-size-lat as shown Appendix F.3.

---

Fast Inference with Kronecker-Sparse Matrices
Initial sparse problem
Dense matrix multiplications in parallel
KT [col1,
[col2, row2]
K T E R acd xabd
X F col1]
X | co12]
X E R
Bxacd
Y E R Bx
[ row
[ rowz)
Figure 5. Illustration of Algorithm 2 for sparsity pattern π = (2,3,2,3) and batch-size B = 8. The rows and columns (row), col1) are associated with () = (0, 1) the "for" loop Algorithm 2, whereas (row2, col2) associated with (i.j) (1,
matrices into smaller submatrices called tiles, which are processed in parallel, with the final result obtained by ac- cumulating intermediate results from each tile (NVIDIA, 2024; 2023a). Our strategy arises from mathematically equivalent reformulation of Algorithm 1, formalized as Al- gorithm 2, which we now explain.
where M[; denotes matrix M restricted to columns J, and M[T, J] denotes the submatrix of M with rows / and columns J. By grouping rows that share the same support, the full KS matrix multiplication reduces to set of smaller, independent dense multiplications one per row,, that can be processed in parallel. This is illustrated in Figure 5. Hence, Algorithm 2 is mathematically equivalent to Algo- rithm 1.
Algorithm 2 New mathematically equivalent tiling of Algo- rithm 1 (no global memory permutations), see Figure 5.
Require: pattern π (a, b, c, d). inputs X = Bxadd
4.2. New Kernel Implementation
KS factor K E KGr
We now introduce our fused implementation of Algorithm 2, denoted KERNEL in the remainder. compact pseudocode sketch is given in Algorithm 3.
Ensure: output Y XK E IR Bxadd
1: Y <- 0 (in global memory)
2: for i = 0 to a j= to d in parallel do
3: row - M 0 b}
4: col f- (i +j+
5: Y[ + +oroow+X[ss coll K [col, row]
KERNEL versus the permutation- based baseline. While both Algorithms 1 and 2 ultimately compute the same set of dense sub- bub-products- one per tile row,,- -hey differ in how they reach them, and hence in their natural GPU implementations:
6: end for
Why Algorithm is Mathematically Equivalent to Al- gorithm I? Given KS pattern π (a, c, d),, the corresponding support S, == Lg @ 1axe & La has shape abd x acd. We partition its row indices [0, - 1] {O...., abd 1} = U U row,y where:
Algorithm 1. This baseline performs two explicit permu- tations of X and Y map the entries of X|: colig] and Y row.e] columns are spaced by d- contigu- ous regions in the permuted tensors X and 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 global memory.
rowijj = { M + kd: : =
(1)
Due to the structure of S., all b rows in each rowi, share the same support, which we denote coli,j
coli,j = +j+
A Algorithm 2. This variant avoids global permutations altogether. It computes each dense sub-product directly, even when the corresponding rows and columns in row, and coli,; j.j not contiguous in global memory. Our fused implementation KERNEL follows this approach: each thread block directly loads X|: coli,] and K [coli,,, row;,] from global to shared memory, performs the multiplication, and writes the result into Y[: row,,), thus bypassing the two extra permutation passes.
(2)
The family {colis) (s.j) forms partition of the column indices [0, acd 1]. This means the matrix multiplication Y = -XKT, when restricted to the rows in r row,,j, can be expressed as:
|1 rows,s]
= (KKT T row,Jl)
= -/s colij) K' [coli,, row,yl
5

---

Fast Inference with Kronecker-Sparse Matrices
Thread-block tiling. We assign to each thread block single dense sub-product { col,,] K KT [coli,; row.yl-
Algorithm 3 Sketch of the fused output-stationary kernel (one tile (zowe,j. col..)) assigned to each thread block).
We make threads within a block:
Require: pattern π = (a, b, C, d), inputs X E Bxacd
KS factor K E 10
ii) cooperatively load the two tiles from global memory
Ensure: output Y == XKT BR
to shared memory,
1: identify (i.j) tile assigned to current thread block
compute zow, col in Algorithm 2
iii) move their assigned subtiles to registers,
3: global - shared: load X[: col], K [colc, row]
(iii) perform multiply-and-ccuumll operations,
4: for stride 0 to c- 1 by size-subtile do
5: identify subtile assigned to current thread
(iv) store intermediate results to shared memory, and
6: shared-rregister: load subtiles of {(; col] and
K [[row,co.] col]
(v) write cooperatively the final output tile back to global memory.
7: Multiply and accumulate into registers
8: Prefetch next subtile in parallel (double buffering)
registers-vgloall store Y|: row]
Standard CUDA optimizations such as double buffering, vectorized memory access, and warp-level parallelism are applied (see Appendix D.3 for implementation details).
4.3. How Much Traffic Does the Fuse Save?
Read/write efficiency. Replacing two permutation ker- nels with direct memory access leads to different tradeoffs depending on the layout:
Against Specialized Baselines. Recall #nz = abed is the number of nonzeros of 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 ker- nels (PERMUTE- SEA-IE--EEMM--EEEMUTE involves 2BN global memory operations (read X E BBN once, write its permuted version X once). The GEMM involves BN + BM global memory operations (read once the permuted inputs X, write once the permuted outputs Y). The last PERMUTE involves 2BM global memory opera- tions of the input and output tensors (read Y once, write Y once). The total number of global reads and writes on the input and output tensors is 3B( M). Our kernel moves each entry exactly once, saving the extra 2 M) global reads/writes of the two permutations. The ratio of wasted global memory traffic to useful multiplies is therefore twice
KS matrix tiles. As in BMM, the factor K is preprocessed and stored so that every tile K [cole,,, row.,] is already contiguous in global memory. This setup cost is paid once and amortized across inference runs with different inputs X, we exclude it from our analysis (Section 1).
The first PERMUTE
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;,)] and Yi: rows,j is stored contiguously in global memory, enabling efficient coalesced access within columns. In this setting, KERNEL. benefits from both eliminating permula- tions and maintaining memory locality.
B(N + M)
b+c be
In contrast, under the batch- major layout (batch-size-firsi in row-major), entries within the same column of X(: colig] or row.,] are not contiguous in global memory, as the indices in coli,j and row,,; 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.
h(bb,c),
(4)
B #nz
which yields the heuristic that correlated with the time spent on permutations in Figure 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 Py PoTocc 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.
In summary, KERNEL eliminates two global-memory round trips compared permutation- 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 butch-size-firs 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 outper- forms baselines in several cases, as observed in Section 5.
Take-Away: The fused, output-stationan 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.
6

---

Fast Inference with Kronecker-Sparse Matrices
Table 2. Across the 600 patterss, how often is an algorithm faster and by how much (median speedup in parentheses); in FP32.
Comparison
Win rate
Median x faster
min m [KERNEL, BMM, EINSUM, BSR] < (DENSE SPARSE]
98.1
x6.5
BMM < DiiNHINSUM,, BSR, DENSE, SPARSE]
90.0 %
x x1.36
KERNEL < min (BMM, EINSUM, BSR, DENSE, SPARSE]
85.5 %
x x1.39
5. Benchmark: 600 KS patterns, Time & Energy
5.1. Time: KERNEL Wins on 85 % of Patterns in FP32
Table 2 highlights three key findings: (i) KS- KS-aware imple- mentations (KERNEI,, BMM, FINSUM, BSR) more than x6 faster than generic baselines that ignore the KS struc- ture (DENSE, SPARSE); (ii) among existing methods, BMM is the fastest baseline (excluding our new KERNEL from the comparison); (iii) our fused, output-stationary KERNEL surpasses all baselines on 85 % of the benchmarked grid in FP32, achieving a median speedup of x1.4.
We now quantify the benefit of KERNEL over the fastest public baselines-BMM, BSR, and FINSUM-on single Nvidia A100 (40 GB) for time and on a single Nvidia V100 (32 GB) for energy".
Benchmarked KS Patterns. We explore patterns of the form π = (a,b, c, d). sweeping over the space a 8 x 3 x a, with a= 21, 4, 6,8, 12, 16, 24, 32, 48, 64, 96, 128) and B (48, 64, 96, 128, 192, 256, 384, 512, 768, 1024). We constrain the shapes by enforcing b = c, b = 4c, or 4b, which reflect common configurations in Transformer architectures." This results in over 600 distinct KS matrix configurations, with sizes ranging from 102 x 102 up to 131072 x 100722-sannnngg six orders of magnitude and broadly covering the hidden dimensions found in vision and language Transformers. The input batch size is fixed to B = 128x x 196 = 25 088, corresponding to batch of 128 sequences with context length of 196 tokens, which is a typical value encountered in Transformer inference." All results below take, for every implementation, the most efficient (in time or energy) of fatch-size-firss and batch- size-last layouts; the full split is in Appendix E.6.
Because the benchmark grid includes many highly sparse KS patterns, we also compare performance at fixed spar- sity levels in Appendix E.2. This analysis confirms that KERNEL provides consistent speedups across all sparsity regimes-non just in the extremely sparse or dense cases.
5.2. What Drives the Speedup? Mostly the Heuristic h(b, c) be
Stc
To understand what drives the speedup achieved by KER- NEL, we perform multiple linear regression in log-log scale using two predictors: the density 1 and the heuristic h(b, c) = bic Before fitting the regression, we verified that the two predictors- density and the heuristic l(b,, cre are only weakly correlated across the 600 KS patterns (Pear- son correlation: 0.28). This low correlation ensures that the regression is not affected by multicollineaarty, and both predictors can be interpreted as contributing distinct infor- mation to the model.
ad
be
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 2 29..7% Additional details on the sparsity distribution are provided in Appendix E.2.
Letting speedup == yields:
min time BMM. .RINIU))M) time (KERNEL)
the regression
*pyJoules is not yet compatible with Nvidia A A100 GPUs so we benchmarked the energy consumption on V100 GPUs. Note that observed the same relative time efficiency ranking the two GPUs.
log(speedup) = 1.69 - 0.031 log + 0.325 log(h).
This model achieves an adjusted R2 of 0.697 and passes the Jarque-Bera normality test on residuals (p 0.12). Both predictors statistically significant (p < 0.001).
'Recall that for a KS pattern (a, b, df. the input and output dimensions of the matrix are given by N acd and MM = abd. c. we have N MM. yielding square matrix suitable for weight matrices in the self-attention block. When b 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. 'The largest LLaMA model (40SB) includes FFN layers with matrices up to 53 248 x 53 248.
When
Notably, h(b, c) is approximately x x10 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 li(b, c).
'e.g.. ViT with a patch size of 16 x 16 for image resolution of 2244* x 224
'We excluded fully dense patterns from this regression, i.c., those with density = 1.
'We chose density over sparsity here, as gave better regres- sion results.
10We were able to pass the normality test in log-log scale, but not inear-linear or log-linear settings.
7

---

Fast Inference with Kronecker-Sparse Matrices
energy(kernll/mm) energy einsum, dense)
1.8
3
min time time(bmm, (bmm, einsum, bsr)
divided by time(kernel)
1.6
1.4
1.2
1.0
0.8
0.9
0000
0.00
0000
1044
0.00
0.008
0.00
808
0.02
0.00
0.00
0.6
0
00
0.4
10
10
10
100
bc
d(b c) (bc)
(a) Time speedup factor.
(b) Energy reduction factor.
Figure 6. Comparison of kernel terms of (a) times speed factor and (b) energy consumption, relative to min (BMM, EINSUM, BSR), on KS patterns FP32, each implementation, we minimum time/energy spent between the setch-size-ffree and batch-size-las memory layouts. We regroup patterns by their value of h1(b, c) bite for(a) and d x h(b, c) for (b).
For each sparsity bin in Appendix E.2, we observe strong and consistent correlation in log-log scale (between 0.76 and 0.91), confirming that h(b, is reliable predictor of performance regardless of sparsity; see Appendix E.2 for details.
(BSF) to batch-size-last, thanks to fully coalesced writes as we saw Section 4.2; full numbers are in Appendix E.6.
Impact of Size. Speedups increase with matrix size M x N = abd x acd (Figure 12); our largest test case (131k x 131k) already exceeds Llama 3-405B's B's hidden layers.
In conclusion, maximizing hbb, c) is a simple and effective design principle for selecting KS patterns that benefit from KERNEL'S ' speedup over prior baselines (see Figure 6a). contrast, density has only a weak and inconsistent influence once h is controlled for.
Portability. All results above use CUDA. For broader hard- ware support, we also provide an OpenCL version, enabling benchmarks on non-NVIDIA devices like AMD GPUs and CPUs.
5.5. Take-Away
5.3. Energy: 15 Median Saving in FP32, Predicted by d c))
Exploiting the KS structure is essential: x6 vs. dense/CSR generic baselines. Removing the two permutations gives in foat-precision additional x1.4 in time and 15 % energy reduction. The simple rule "maximise lib, c)" predicts where the fused kernel shines; sparsity alone does not.
Because each column gathered by permutation is spaced by d, the distance travelled by memory requests grows with d, and dxh(b, c) is 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 existing baselines on 72 % of the patterns, achieving median saving of 15 % in FP32.
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- gains, we inject KS sparsity into Vision Transformer (Dosovitskiy al., 2020) and GPT-style language model (Radford et al., 2019), and measure inference time.
5.4. Additional Remarks
BSF vs. BSL. Switching baselines to batch-size-last (BSL) does not change the runtime of "I However, KER- NEL gains stable x2 when switching from butch-size-first
"Not that BMM remains the best basclines compared to EIN- SUM, both BSF BSL layouts (Appendix E.6 for details).
Target Architecture. Fully connected (FC) layers are natural candidates for KS basba acceleration. In our
8

---

Fast Inference with Kronecker-Sparse Matrices
experiments, we simply replace the standard PyTorch Linear module with our publicly available near module."2 12 Transformer architectures such as VIT and GPT offer suitable testbed, as they contain numerous FC lay- ers ers-accounting for 30-60% of total inference time de- pending on model size (Appendix E.7). 13 We evaluate our method on V (21M parameters) and GPT-2 Medium (345M), which both fit comfortably on a single GPU.
7. Conclusion and Discussion
We revisited GPU multiplication with Kronecker-sparse (KS) matrices, identified that permutation overheads dom- inate existing kernels, and proposed an output-stationary tiling strategy that fuses both permutations with the GEMM into single CUDA/OpenCL, kernel. Across 600 KS pat- terns, our kernel is x 1.44 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.
Setup. Following the protocol of Dao al. (2022b), we replace specific weight matrices of linear layers in ViT and GPT with a product of two KS factors, using patterns de- tailed in the appendices. All other components of the Trans- former forward pass use the default Py PyTorch kernels in the default butch-size-firs layout. Complete setup details are provided in Appendix E.8.
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 kemnll.
Results. Table shows that the fused kemnel reduces end-to- end latency by 22 % on ViT-S/16 in FP32, corresponding to x x1.28 speedup over the original dense implementation. For comparison, the best public baseline achieves an 11 % reduction (x 1.12 speedup). On GPT-2 Medium, the fused kernel yields a 16 % latency reduction (x1.19 speedup), while the best baseline achieves 12 % (x 1.14 speedup). Ad- ditional experimental results are provided in Appendix E.8.
Simple Design Rule. A one-line heuristic h(b, c) D+c helps identify patterns well-suited for this kernel.
be
We release 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 inves- tigations: is it our kernel that be further improved, or the other implementations that have been particularly optimized in these regimes?
Table 3. Latency (fully-connected) Lower is better.
ratio
versus the dense baseline
Sparsity level
time(BMM) (ime(DENSE)
dime (KENEL) (DENSB)
VIT-S/16
57.1 % 24.8 %
0.89
0.78
GPT-2 Medium
0.88
0.84
iii
Batch-size-last layout outperforms batch-size-firg in our KERNEL (Figure 14), but PyTorch defaults to BSF. Using BSL in BSF pipelines requires costly transpositions that may offset the gains. Our KSLi near r layer supports both: BSF for drop-in use, and BSL for higher perfor- mance when compatible. More broadly, the benefits of BSL raise an open question: how does batch dimension placement affect neural network efficiency in general?
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 batch-size-last can yield up to x2 speedup within KS layers. Extending batch-size-last support to other core components of neural network inference- such as activation functions, softmax, and LayerNorm-coull 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 engi- neering challenge.
(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 accu- racy when training/finetuning with KSL nea r layers.
(iv) Cache behavior was not analyzed in depth; future work may reveal strengths or limitations of our tiling strategy and suggest further improvements.
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.
Twae-Away:
(v) Speedups depend batch size, KS pattern, and few kernel hyperparameters (e.g., thread count, tile shape). Tuning is sptup-specific and tedious, so we provide pre- sets for the 600 patterns we benchmarked and leave general tuning strategies to future work.
12 A'aallable at
See Appendix E.8 for why applying KS sparsity convolu- tional networks poses additional challenges.
(vi) Beyond NVIDIA hardware, the OpenCL port opens the door to studies on AMD GPUs and other platforms.
9

---

Fast Inference with Kronecker-Sparse Matrices
Acknowledgments
Dao, T., Chen, B., Sohoni, N. S., Desai, A. D., Poli, M., Gro- gan, J., Liu, A., Rao, A., Rudra, A., and Ré, C. Monarch: Expressive structured matrices for efficient and accurate training. In ICML, 2022b.
This work was supported in part by the AssegroAssai BUR-I9-CHIA-OUR by the NuSCAP ANR-20- ANR-20-CE48- projects of the French Agence Nationale de la Recherche and by the SHARP ANR project NAR-22-PEIA--001 in the context of the France 2030 program.
Dao, T., Fu, D., Ermon, S., Rudra, A., and Ré, C. Flashat- tention: Fast and memory-efficient exact attention with io-awareness. In NeurIPS, 2022c.
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-P133-0004, and Air Force Office of Scientific Research, Air Force Material Command, USAF, under grant numbers FA8655-22-1-700
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.
The authors thank the Blaise Pascal Center for the computa- tional 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.
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.
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.
HPCwire. AWS Upgrades its GPU-Backed AI Inference Platform. https: -ttpgradddd--ts------bceeeeeee
com 2019
1ttp:///awsupprraaeee d-inferrenc-pplee foatfor// March 2019. Accessed: 2025].
References
Bar, J. Amazon EC2 Update- Infl Instances with AWS Inferentia Chips for High Performance Cost-Effective Inferencing, 2019. URL https:://wws.aaaz... hww.lloo///aassmmmmccoeeeeeeeepp
Le, Q.-.,, 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 sparse relu network training problem always admit an optimum? In NeurlPS, 2023..
oo-
nferencing/. Accessed:
Accessse [April 2025].
Boehm, S. How to optimize CUDA matmul kernel for cuBLAS-like performance: A worklog, 2022. http s: s://siboehn com comm/artiiaa /tt//aarccllse///uap//// -MM [Accessed: April 2024].
Le, Q.-.,, Zheng, L., Riceietti, E., and Gribonval, R. But- terfly factorization with error guarantees. arXiv preprint arXiv:4411.05506 2024.
Li, X., Liang, Y., Yan, S., Jia, L., and Li, Y. A coordi- nated tiling and batching framework for efficient GEMM on GPUs. In Symposium on Principles and Practice of Parallel Programming, 2019.
Dao, T., Gu, A., Fichhom, M., Rudra, A., and Ré, C. Lear- ing fast algorithms for linear transforms using butterfly factorizations. In /CML, 2019.
Dao, T., Sohoni, N. S., Gu, A., Eichhorn, M., Blonder, A., Leszezynski, M., Rudra, A., and Ré, C. Kaleidoscope: An efficient, learnable representation for all structured linear maps. In ICLR, 2020.
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.
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.
com /Om//VIIIIA//// ass / /blob maamaannmmee mewia http///cceeeeee.t.../// emm md [Accessed: April 2024].
10

---

Fast Inference with Kronecker-Sparse Matrices
NVIDIA. Matrix multiplication background user's guide, 2023b. https://docs.nvv com deepp l hrapl////perfoommmmccedeeeeerrre
Zhai, X., Kolesnikov, A., Houlsby, N., and Beyer, L. Scaling vision transformers. In CVPR, 2022.
mat rie-multipilcatioo//ndee...... html [Accessed: April 2024].
Zheng, L., Riccietti, E., and Gribonval, R. Efficient identi-
fication of butterfly sparse matrix factorizations, SIAM Journal on Mathematics of Data Science, 5(1):22-49, 2023.
NVIDIA. CUDA C+ programming guide, 2024, https: //docs avidia nwwiiiiiii..... com/ cuda/cdauiee cudaa c program ming-quide/ /index html [Accessed: April 2024].
Quemener, E. and Corvellec, M. SIDUS solution for extreme deduplication of an operating system. Linux Journal, 2013(235):3,000)
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Spconv Contributors. Spconv: Spatially sparse convolution library, /ttp:://githhn comm/travelleer59// poonv, 2022.
Tang, H., Liu, Z., Li, X., Lin, Y., and Han, S. Torchsparse: Efficient point cloud inference engine. In Conference 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. CVPR Workshops, 2023.
Vahid, K. A., Prabhu, A., Farhadi, A., and Rastegari, M. Butterfly transform: An efficient FFT based neural archi- tecture design. In CVPR, 2020.
Van Loan, F. The ubiquitous kronecker product. Journal of computational and applied mathematics, 123(1-2):85- 100, 2000.
Wang, Scaled dot-product attention implementation, 2024a. amnila//prrffmmanee///-rrtroomee
www.iic..iiiiii.ccc.o.m.. ctt.iddiaa.c///eeppp//eeeee deeple
tix-mltiiplictiionn April 2024).
html [Accessed:
Wang, P. Simple VIT implementation, 2024b. https:
//github. com/ http:///cciaaaaaaaaale----eeoeee
/ htalbb/main/vtt__yytrrhchh wtt.sstttttttttoooooeeeeeeseeeee
[Accessed: April 2024].
Yan, Y., Mao, Y., and Li, B. SECOND: Sparsely embedded convolutional detection. Sensors. 18(10):3337, 2018.
Ye, 7., Lai, R., Shao, J., Chen, T., and Ceze, L. SparseTIR: Composable abstractions for sparse compilation in deep learning. ACM International Conference on Architec- tural Support for Programming Languages and Operating Systems, 2023.

---

Fast Inference with Kronecker-Sparsc Matrices
Appendices
Lay summary. Today's s Al models 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-buu 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 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 helps. Plugged into large Al models like Transformers, it cuts total running time by up to 22%, paving the way faster, cheaper AI.
A. Existing GPU Implementations of Kronecker-Sparse Matrix Multiplication
This section describes the existing GPU implementations for KS matrix multiplication on PyTTrrc. All source code is available in the forward pass of the KSLinear module we release at gittub..om/PssalCaarvvinnns//s
A.1. Specialized Baselines
BMM (batched GEMM). The implementation released with the "Monarch" layers of Dao et al. (2022b) permutes the weight matrix once offline and stores K E gadxbxe as dense float/half tensor of shape (ad, 6, e). Given a batch of inputs X Xxod the forward pass executes Algorithm with fast tensor reshapes for the input and output permutations, and with the fast routine torch. bmm for the multiplication with K:
1) reshape X that multiplying by Q T is view;
2) perform the batched GEMM torch. lom;
3) reshape the result that multiplying by PT is view.
A concise PyTorch version follows (layout batch-size-first; swap first/last dimensions for batch-size-last) where we denote by bmm the 4D PyTorch tensor storing K and by b.b the batch-size-firs/ PyTocch tensor corresponding to a matrix T.
def kronecker_tme bmm (X k_bs,, K K_bmm)
2.
batch_size shapee
perrmm =
X (bat.cccwbbbtccc..oc size, a, c, d)
transpose (-1, -2)
6
reshape (batch_size, a d, c)
c)ntigucus() transpose 1)
perm terch, empty (batch size, * d, b) trranspose (0, 1) Y porm ferch.bmm perm,,,, "_bnm.trannpose (-1, )
11
bof =
Y __perm transpose (0, 1)
reshape (batch_sizee a,
15
transpose -2)
16
reshape seatch__siees a * b d)
return
BSR (block-sparse GEMM). The same algebraic steps apply, but K is stored once for all in Block- cookkccopeesseed Sparse Row format as a tensor Kssr,, and Step (2) calls torch. sparse mm:
def kronecker_bar (x bsf, K_bsr);
The original release supported only л = (a, b, c, d) with a=1 or d=1; we extended to general case.
12

---

Fast Inference with Kronecker-Sparse Matrices
batch_size shape
3
X_perm
bosf.view wiwttbhchhsssssee a, C, d)
(-1,, -2)
reshappe seatch_ssze,
d)
porm terch. narch.nn.ffin functioon lincar
perrm bsr
11
bsf
12
Leerm. (batch sez, a, d, b)
(-1, -2]
reshape (batch_size, size, d)
return bsf
EINSUM (tensor contraction). Following the public prototype in the Monarch repository (Dao et al., 2022b), the abed
nonzeros of K are packed in Keinum E Raxbxexd After reshaping the input X E RR Rxand to Xeinsum E R Axaxexd and the
output Y e R Babdd to Yeissum E RBxxxxxd the contraction (eiissu):...ba = Ee (X Neinsum S,,,e, d (Kinsum)a J, b, is issued via einopss. . einsum:
def kronecker _eteerreeiinummmeeeee (X_bsf, K_einsum)
X_perm e nops. reareerrre rearrange (X_bsf, " .... (a c d) -> ... a c d", ama, c=c, d-d)
3.
Y perm e einops. pepss.einse einsum (X perm, X eiinsum, "... a d, a b d -> b d")
Y bsf = eineps. rearrange (Y perm, "... b d-> (a b d)
return bsf
The second line of this code does at the same time all the matrix multiplications ( row] + coljK [col, row] for all the pairs (row, col) in Algorithm 2.
Other KS implementations The "Monarch" batched- GEEM code already covered above (BMM) is the maintained reference for Butterfly/Monarch layers of Dao al. (2019; 2022b). We could not compile the original Butterfly Transform kernels (Dao et al., 2019; Vahid et al., 2020) on recent CUDA; 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 stored (even zero entries) and the forward pass calls torch. nn. f fanctional... linear. With butch-size-first layout this is the PyTorch default; in butch-size-last we select torch. matmul after an internal benchmark of alternatives.
SPARSE. Only the abed non-zeros are kept (CSR format) and the same linear/matmul entry point is used. This baseline uses the sparsity but not specific Kronecker structure.
A.3. Prior Empirical Reports on the Efficiency of Existing KS Matrix Multiplication Algorithms
Dao et al. (2022b) reported a x2 training speedup on image classification and language-modellng tasks by replacing dense layers with products of two KS factors.
Fu et (2023) measured a speedup exceeding inference for FFT-based layers X + W- c 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 sequential multiplication with Kronecker-sparse matrices (see Figure 1).
Our study complements the above by isolating the most elementary operation these different pipelines: the multiplication by a single KS factor X + KX. We benchmark more than 600 distinct KS patterns spanning six orders of magnitude in matrix shapes, comparing existing GPU baselines on Py Torch.
13

---

Fast Inference with Kronecker-Sparse Matrices
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-structuke designs?
Table 4 shows that several sparsity patterns previously proposed for neural network compression can be expressed special cases of KS patterns, illustrating the broad applicability of the framework studied in this paper.
Table KS patterns used in prior work to replace dense layers neural networks. For a product K.... Ky we list (az, be, ce, deliseeL.
MATRIX SIZE
KRONECKER PATTERNS
DENSE
M N
N,
LOW-RANK
M N
1..
SQUARE DYADIC (DAo
2019; VAHID AL.,, 2020)
24 2
(2'-1,2,2 2,2,2
KALEIDOSCOPE (DAo ET AL., 2022o)
2 24
(2- 2, JEST (2 2,2,,
BLOCK AUTTERFLY Y (Da0 FT AI.,, 2022A)
2 24
(2 1,,, 24,2" -ih
MONARCH (Dao ET AL., 2022h; FU ET AL., 2023)
M x N
(1, Min(M. (g, M,,ii N/p N/p.1) (ac,be, cr difa WITH ocedde = aesabesidess
DEFORMABLE BUTTERFLY (LIN ET AL., 2021)
qidid x agerdd
How does the Kronecker-sparstt pattern compare to other forms (unstructured or structured) of sparsity and when is each preferable?
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 minimizer of the loss is guaranteed to exist, unlike in the unstructured case where solutions may diverge toward altaatiaiing 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.
Structured
B.2. Dense > KS Conversion (Out of Scope in this Paper)
How one convert dense layer to 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 .... K2 such that W = K₃ ... K2 (in Frobenius norm) using a quasi-optimal butterfly factorization algorithm (with provable error guarantees) that has a roughly O(MN N) time complexity for the factorization a matrix size M x N (Le et al., 2024; Zheng et al., 2023; 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 more accurate and efficient than alternating least squares (Lin et al., 2021) gradient descent (Dao et al., 2019). In our work, we focus pure inference mode where the KS matrix K is given and fixed. The cost of 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 al., 2022) performs sparse 3D convolutions three stages: (Tang et al., 2023; Spconv Contributors, 2022: Yan al., 2018) improved performance by overlapping memory transfers with computation. Similarly, for the sparse problem tackle, new dataflow strategy implementing Algorithm 2 leads to comparable improvement: input/output permutations and matrix multiplication can now be fused into single kemel (Figure 3, right), enhancing efficiency.
The works
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 straightforward way to adapt the tiling strategy from Tang et al. (2023) to the nested Kronecker layout we study.
14

---

Fast Inference with Kronecker-Sparse Matrices
B.4. Relation to Sparse-Tensor Compilers
Staee-of-the-art sparse-tensor compilers, e.g., SparseTIR (Ye al., 2023) Fractal (Guan al., 2024), are primarily designed for conventional sparse formats such CSR BSR, combined with regular block tilings. KS matrices, by contrast, cxhibit nested sparsity pattern: block-diagonal layout with X sub-blocks, each refined internally by d x d identity structure. Representing such a hierarchy would require support for outer block-sparse format with an inner identity patter--apaailities that current compilers do not offer natively.
Furthermore, while our Algorithm uses standard tiling strategies when possible, sometimes favors non-contiguoss 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, the best of our knowledge.
Until compiler infrastructure evolves to support nested sparsity and more flexible tiling schemes, hand-optimized kemelss seems to remain an effective solution for running KS workloads.
C. Perfect-Shuffte Permutations Used in Algorithm 1
In practice, existing implementations Algorithm select the permutation matrices
P == (Ia & Po.d),
Q I Ia @ Pa.ds
(5)
where Pp.q is the (p, q) perfect-shuffte permutation matrix. This appendix reviews perfect-shuffle permutations and shows why the choice in (5) produces the permuted matrix
k KQ
that is block-diagonal with dense sub-blocks bww.bbkk--spcciicclllll..... ad diagonal dense blocks of size bxc. Because each block can be multiplied with highly optimized dense routines, Step (2) of Algorithm 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.g) q) perfect-shuffle permutation matrix (Van Loan, 2000)). Let p, N and set r = pq. The (p, q) perfect-shuffle permutation matrix P p.g the r r permutation obtained by reshaping a length-r vector into a p x matrix in row-major order and then reading it column-wise. Equivalently,
I, Ro I, :
P Pp. =
(6)
I, [. 1.
where R; = +j Ij E [0,p- 1}} for E [0,g - 1] and I, [R, E {0, 1)pxr is the restriction the identity matrix I, to the rows indexed R.
Because the sets Ro..... Rq- partition r 1}, p.g is a (row) permutation of I,, and in particular is itself permutation matrix.
C.2. Dhool-iiagnnalizingg KS Matrix with Perfect-Shuffle Permutations
We now show that K as defined in Algorithm
K : :
block diagonaa when P, Q are chosen as perfect-shuffle permutations in (5) and K KS matrix with pattern (a, 6, c, d). Consider the KS pattern (ad, b, c, 1), whose support S Led & 1ixe corresponds ad dense diagonal blocks of size x It suffices to establish
Sn (I, 2 Pb.)) S [I. & P Poa)
PS Q
Q
15

---

Fast Inference with Kronecker-Sparse Matrices
Because permutation matrices are orthogonal, S S For any compatible matrices A, B, C we have
supp(ABC) [
supp(supp(( () supp B) Supp(C)),
where supp(-) is the binary support mask. Since supp(P) = P and supp Q) = =Q.
supp (PPPPP C supp(S#) Lad @ lbxes
which is the support of a block-diagonal matrix with ad dense b x blocks, proving the claim.
Proof of (7). A key identity from (Van Loan, 2000) states that for all positive integers b, c, d.
Po,d" (1oxe 2 Io) Pe.d La 1bxe-
Because 7 La & 18xe Id, we obtain
S = Ia @ Pb.d (Ia & l8xe) Ped)) ))
Using (AB) @ (CD)= = (A & (B @ D) for any matrices A, B, C, D of compatible sizes, we factor:
S = Ia @ (oo & loxe) Pa."
(I. s Pb.d) Ia 3 (Ia loxe)
(In Pb.d) IIa & La lxxe) Pea")
= (Ia @ Pb.d) Tad & (Ta & Ped") ).
which is exactly (7).
Algorithmic Implications. Because 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 and 3). which we measure to be a neg-neglig share of runtime. Section 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 plug-and-play PyTorch module, KSLinear. that can replace standard near layers in neural networks.
Separately, we are also integrated our fast matrix factorization routines into the lazyl i nop ecosystem, package more oriented toward signal processing and fast linear transforms. These two entry points are intended to remain independent moving forward.
lazylinog Python library that supports delayed execution: it builds an operator expression graph lazily and only evaluates it when final result is requested (delayed evaluation paradigm). One of its key features with respect
16

---

Fast Inference with Kronecker-Sparse Matrices
Kronecker-sparsity is the ability approximate dense matrix W into fast transform, i.e., product of structured KS factors:
W K K
This results in fast-transform approximations that might be more efficient time and energy, especially on parallel hardware.
CUDA and OpenCL backends enable efficient evaluation of these factor chains K , (batch dimension last). The Butterfly module of lazyl inop supports execution both CPUs and NVIDIA/AMD GPUs. The documentation is available at https://austtrrp. https///austgrrggiiaaaggssss wwwggggpgge..nnniiiie inria inriaa.f/lazylinoo/ideee html and the low-level source code is available at https://gitlab.nnie wtt.aaafffassssggp//aaaa
batch-size-last layout
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, lazyllliinn inop approximates it as W x K using quasi-optimal factorization algorithm that achieves the theoretical error bounds of Le et al., all in ((MN) time for an M x N matrix.
2) Compatibility with Py PyTorch tensors (in progress), Py Torch tensor inputs and outputs will soon be supported directly within lazylinng nop, 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) Neighh-sttationany seems to offer little reuse. In KS pattern (a, b, d), when the parameter 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 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 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 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; Bochm, 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 contiguous pattern, avoiding scattered global writes (NVIDIA, 2023a).
Kernel launch parameters (block size, tile shape, vector width) auto-tuned per pattern (a, 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 12.1.
17

---

Fast Inference with Kronecker-Sparse Matrices
Matrix Sizes. In all our experiments with matrices, set the batch size B 25088 (196 tokens 128 sequences), standard choice for ViTs it corresoonds 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 E RRoldx aod with sparsity pattern (a. b, c, d) are drawn i.idd uniformly in -- ve' Jel. corresponding to the initialization used for training in Dao et al. (2022b). The entries of the inputs X are drawn i..d. according to standard normal distribution N 1).
Benchmarking Time Execution. All the experiments measuring time execution of 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 A A100-PCIE-40GB GPU associated with an Intel(R) Xe) Silver 4215R CPU @ 3.20GHz with 377G of memory. The full benchmark took approximately days in an isolated environment, ensuring that no other processes were running concurrently.
Measurements are done using the Py Torch tool torch. ut 1s benchmark Timer. The medians are computed on 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-I6GB GPU associated with an Intel(R) Xeon(R) Silver 42159 CPU 3.20GHz with 754G of memory. The full benchmark took approximately 1.5 days 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 6 = 4e or = 46 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 wide range of patterns (a, b, c, d) with a 1, this represents simply corresponds to repeating a times the case 1 in parallel.
the simplest scenario. Indeed, the case
The second "for" loop in Figure 7 (line 26 to 30) generates patterns with > 1 offering fewer choices for d to keep the benchmark concise in terms of execution time. This loop also imposes additional conditions on 6 and (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, e) 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) /b to observe the trend in Figure 6b, while keeping the benchmark as short as possible. We chose to consider the cartesian product of
1 a_iis = [1, 4, 16, 32, 64]
b bist 148, 64, 96, 128, 192, 256, 384, 512, 768, 1024]
list 148, 64, 96, 128, 192, 256, 384, 512, 768, 1024] dist [1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64]
by skipping as in Figure 7 all the patterns with
(b, a) in [ (1024. , 256) , (256 1024) , (128 512) (512 128)
, (64
256)
(256
6
and also all the patterns such that
b != and b != 4 c and != 4 b
for the same reasons as explained above for time measurements.
Details on boxplots. In all boxplots (Figures 4, to 6b, 10 and 12 to 18), the orange line corresponds to the median, the
18

---

Fast Inference with Kronecker-Sparse Matrices
import itertools
batch_size 25_088 size_limi 2_147_483_647
6
a_list [1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, 96, 128] b_it [48, 64, 96, 128, 192, 256, 384, 512, 768, 1024]
c_list [48, 64, 96, 128, 192, 256, 384, 512, 768, 1024]
dlist1 [1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, 96, 128]
10
1istt [4, 16, 64]
11
def get_ patterrs_benchmane ()
patterna_list []
15
def add pat_erate b, c, d)
it batch size + c d <= size limit and \ batch_size
d size_limit and
aaba c patterns_list appppa ((a, b, c, d))
size_limit:
19
for b, c, d in peattools.rrouute (b list, c_list, d_listl): 20
=1 if cr or patternna (a, b, c, d)
or 4 b):
24
25
for a, b, c, d in itertools, product list, b list, c_list, list2) !- and
28
(b, c) not [11224, 256), (256, 1024), (128, 512), (512, (28), (64, 256),
(256, 64)] and
(b == c or ==
== 4 b):
add (a b, c,
31
32
return patterns_list
Figure 7. Python code generate the patterns benchmarked for the execution time in the numerical experiments of Section 5.
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. 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 result, any aggregated statistic computed over the full set of 600 patterns such as the medians reported in Table 2eness 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 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 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) bit In all six bins, 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 lb,,, c), rather than by the density of nonzero entries.
We also provide a visual counterpart in Figure 9, which extends the analysis 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 the heuristic explaining performance gains.
19

---

Fast Inference with Kronecker-Sparse Matrices
Boxplot
Histogram
160
140
120
100
Count
80
60
0.6
0.7
1.0
0.7 Sparsity
Sparsity 1-
Figure 8. Boxplot and histogram of the sparsity levels covered by the 600 KS patterns of benchmark in Scction 5.
Table 5. Win rate and median speed-up of KERNEL min m BMM, EINSUM, BSR) across sparsity bins, in FP32. The correlation is between
log(specdup) and log(h) within each bin, where h(b, c)
bc
Sparsity bin
Win rate
Median x faster
Number of patterns
Corr(log-log)
99.48-99 99.93%
100.0%
1.84
102
0.91
98. 98.96-99.48 488%
93.1%
1.41
101
0.82
97.92-98.900 9.%
79.6%
1.38
108
0.83
35.83-97.929%
81.5%
1.41
92
0,80
$7.49-95.833
89.5%
1.30
114
0.80
49. 19.99-87.4999
86.0%
1.23
86
0.76
E.3. Measuring the Cost of the Permutation Steps in BMM (Section 3)
Protocol. The BMM baseline follows Algorithm 1, which for pattern (a, b, c, d) performs two global permutations (lines 1 and 3) plus one matrix multiplication with the permuted factor K. Because K itself has pattern (ad, b, 1), we isolate the permutation overhead by timing:
1. At: the full BMM kernel on the original pattern (a, b, c, d);
2. A/: the same kernel restricted the multiplication step only, i.e., BMM executed (ad, b, b, 1).
The share of runtime devoted to the two permutations is
A/ AE
At
Results. Figure 10 shows the distribution of memory transfer overheads at fixed batch size of 25088. The patterns included
in this plot are of the form (1, b, c, d) with e {48, 64, 96, 128, 192, 256, 384, 512, 768, 1024) such that b c, b 4c. 4b, and 2,, 3, 4, 6, 8, 12, 16, 24,32, 32, 48, 64, 96, 128).
In the batch-size-firss layout (Figure 10a, reproduced from Figure 4), the overhead of BMM due to the two global permutations in Algorithm increases with the ratio h(b, c) be
bic reaching up to 50 %.
The same trend is observed in the batch-size-las 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 x 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 x 10a (first quartile) and 6.0 x 10 (third quartile). We recall that the output and input dimensions are M abd and acd,
20

---

Fast Inference with Kronecker-Sparse Matrices
Sparsity: (99.5%, 99.9%) (102 samples)
Sparsity: (99.0%, 99.5%) (101 samples)
min time (bmm, bsr, einsum) divided by time(kernel)
Sparsity
Sparsity
99.48% (33 samples)
98. 90% samples)
samples]
samples)
99. 74% (21 samples)
samples)
99 87% (6 samples)
99.90% (4 samples)
99.33% samples)
1
0.005
0.006
0.00
0.008
0.009
0.01
0.02
0.03
0.04
0.00
000
0000
0.00
0.007
0003
0001
0.02
0.00
0.04
0.
b C
b
Sparsity: (97.9%, 99.0%) samples)
Sparsity: (95.8%, 97.9%) samples)
min time(bmm, bsr, einsum) divided by time(kernel)
Sparsity
Sparsity samples)
samples) 98 44% (54 samples)
96. 88% (54 samples)
2
1
0.002
0,09
0.00
0.099
0.006
3,007
0.008
000
0.03
0.04
0.002
0,05
0.00
0.005
0.006
0.007
0.008
900
0.03
0.04
b+c
b+c
bc
Dc
Sparsity: (87.5%, 95.8%) (114 samples)
Sparsity: (50.0%, 87.5%) (86 samples)
min time bsr, einsum) divided by time(kernel)
Sparsity
Sparsity
87.50% samples)
00% samplesi 66. 67% (22 samples)
91.67%% (38 samples)
93.75% samples)
samplesi
83. 33% (22 samplesi
2
0.9
0.002
0008
0.00
0.005
0,006
0,007
0.098
0.01
0.3
0.04
0.002
0009
0004
0.005
0.006
000
0.008
0.9
0.01
0.03
0.04
b+C
b
bc
bc
Figure Speedup factor KERNEL compared baselines (BMM, BSR, EINSUM) the heuristic h(b, c)
bic when fixing the sparsity
level at given value, in FP32, Each subplot corresponds to different sparsity bins involved in Table 5.
21

---

Fast Inference with Kronecker-Sparse Matrices
0.50
0.45
Relative time in memory rewritings bmm
Relative time memory rewritings bmm
0.40
0.4
0.35
0.30
0.3
0.25
0.20
0.2
0.15
8
0.10
0.003
0.0
0.005
0.006
0.007
000
0.01
0.02
0.0
0.0
0002
0003
0004
0.006
0.007
000
0009
0.01
0.02
0.0
0.04
b+C
bc
(a) Batch-size-Biee first (same Figure 4).
(b) Batch-size-last.
Figure 10. Estimated share of runtime devoted to the permutation steps in the BMM baseline for many patterns я = (a, b, c, d), in FP32.
Patterns are grouped by the value of l(b, c) htc and the distribution within each group is shown as boxplot.
respectively, for Kronecker-spiin sparee matrix with pattern (a, c.
Boxplot
Histogram
50
40
Count
30
O
20
10
105
107
10
10s
107
10a
x value
M N value
Figure 11. Distribution matrix sizes the henchmark.
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 x N, we always have time(BMM) < min time(BSR, EINSUM), i.c. the BMM implementation is the most efficient among all baseline implementations (BMM, EINSUM, BSR).
22

---

Fast Inference with Kronecker-Sparse Matrices
divided by min (ierme. bmm, einsum, bsr)
10-
8-
min time(dense, sparse)
min timeleinsung bsr) divided by time(bmm)
0.9-
0.8
0.9
0.8
0.7
10a
10a
10c
107
10a
10a
10a
10a
10"
10"
10a
10a
M x
M x
Figure 12. Speedup factor of min time(KERNEL, BMM, BSR, EIN- SUM) compared to min time(DENSE, SPARSE) as function of the matrix size M x N.
Figure 13. Specdup factor timc(BMM) compared to min time(FINSUM, BSR) function of the matrix size M x N.
E.6. Details on the Impact of the Memory Layout (Section 5.4)
Figure 14a shows how the memory layout (batch-size-firs) vs. butch-size-last) impacts the runtime of each implementation. In particular, it shows that KERNEL achieves a median speedup of approximately x2 when switching from batch-size-first to batch-size satch-size-lass size-last.
20
Time batch- size-first Time batch-sized size--lzzlstt laat
Time batch-size-fin wen--siz--fiiie fee-irrs// Time batch-sizer sich-sizee-last
10.
2.
0.3
0.6
0.8
0.7
0.4
0.4
0.3
kernel bmm einsum bsr dense sparse
kernel bmm einsum bsr dense sparse
(a) floaat-precision.
(b) half-precision.
Figure 14, Boxplots the ratio biteo-site--frst FP32.
23

---

Fast Inference with Kronecker-Sparse Matrices
Table 6 reports the percentage of KS patterns for which KERNEL outperforms all baseline implementations, in either the batch-size-fist batch-size-last memory layout. Even when all implementations restricted to the batch-size-first layout, KERNEL still improves over all baselines on 20 % of the tested patterns, despite relying some non-contiguous memory accesses (Section 4.2).
Table 6 also confirms that among existing baselines (excluding the new KERNEL.,, BMM is the fastest across both layouts, outperforming EINSUM, BSR, DENSE, and SPARSE.
Table Percentage out 600 patterss (a, b, c, d) where algol is faster than the algo2 (denoted by time(algo1) < time time(algo2)) and the median accoleration factor in such cases (that is, the median ratio
time of a1ge2 for a given memory layout), in FP32.
time of aigol
Batch-size-fistt
Batch-size-last
Comparison
Win rate
Median faster
Win rate
Median x faster
min KERNEL, BMM, EINSUM, BSR) < (DENSE, SPARSE]
97.61 %
x x18.50
98.09 %
x6.53
BMM < min(EINSUM, BSR, DENSE, SPARSE}
94.1 %
x X1.44
89.47 %
x x1.66
KERNEL min BMM, EINSUM, BSR, DENSE, SPARSE]
20.41 %
x1.26
85.49
x1.39
E.7. Time Spent in Linear Layers in Vision Transformers (Section 6)
This section provides 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 inffoat-pree ffoat-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 7. Median exccution times (ms) for the forward pass in various ViTs, and for an MLP composed solcly of the lincar layers from the feed-forward network modules of ViTs. The table also reports the proportion of the latter over former. FP16 corresponds to hulf-precissinn and FP32 to fllaat-precisionn
ARCHITECTRRE
FP16 (S)
FP32 (S)
COMPLETE
LINEAR IN FFNS
COMPLETE
LINEAR IN FFNS
VIT-S/16
0.014
0.0046 (31 %)
0.090
0,04 (46 %)
VIT-B/IC
0.036
0.015 (42 %)
0.30
0.16 (54 %)
VIT-I/16
0.11
0.050 (46 (46%)
1.0
0.58 (58 %)
VIT-H/14
0.31
0.16 (53%)
2.6
1.6 (61 %)
Details the Estimation. A Transformer architecture consists of sequence of blocks, each containing multi-head attention module and feed-forward network module. The feed-forward network (FFN) is two-layer MLP (without biases) with an intermediate non- 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 x 224. For float-precision implementations of Wang (2024b). For hulf-precision, the Transformer uses FlashAttention (Dao et al., 2022c) to compute scaled dot-produce attention, as in (Wang. 2024a). The FFN-only MLPs are implemented using torch. nn Sequential and torch. nn. Linear. All experiments were run on a single NVIDIA 4100-40GB GPU paired with an AMD EPYC 7742 64-Core Processor. Execution times were measured using t torch. utils s. benchmark. rk. Timer, with a batch size of 128.
the PyTorch ViT
24

---

Fast Inference with Kronecker-Sparse Matrices
E.8. Details End-to-End Transformers Inference Acceleration (Section 6)
Why Do We Focus More on Transformers Rather Than, e-g., Convolutional Neural Networks? In Section argued that injecting KS in linear layers is straightforward-maakee 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 convineing enough to be considered in this paper.
The first approach upproach-replacingg the convolution kernel with a chonecker-spareee matrix-is limited by the small size of typical convolution kernels (e.g.. 7x 7), leaving little room for meaningful gains.
The second approach-recasting convolutions as matrix multiplications with butterfly-strucuured weights, as in (Lin et al., 2021)-introdee 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. our ViT-S/16 experiments, we replace the dense weight matrices in all fully-connected namers 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 K2 where each K; follows a specific sparsity pattern и, defined in Definition 2.1. The patterns used are:
(T1, π2) ((1, 192, 48, 2), (2, 48,48, 192, for square matrices of size N x N (projections in self-attentio blocks),
(1, 768, 192, 2), (6, 64, 64, 1) for 4N N matrices (UP projection matrices),
(6, 64, 256, 1), (I, 128, 128,3) 3) for N 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 (0: bi d1d, + azb202d2 2 0.75MMN, 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 high value of the heuristic h(b, c) be of our fused KERNEL. implementation.
bic which favors the performance
Additional Results for VIT-S/16. Table 8 presents the latency reduction achieved by replacing standard Linear layers with Kronecker-spane sparke 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 Acceleration of submodules of ViT-S/16 using Kronecker-sparse matrices, FP32.
time (RMM)
fime(KFRNFL))
lime Fu ly-conesteed)
linel (f (ullyfooocnceet
Linear N x V
0.82
0.50
Linear N x N + bias
0.97
0.66
Linear 4N x N
0.80
0.78
Linear 4N X N bias
0.93
0.90
Linear N x 4N
0.91
0.58
Linear N X 4N bias
0.94
0.61
FFN (2 x Linear+GELU+LN)
0.91
0.77
MHA (QKV + proj)
0.87
0.79
Transformer block
0.90
0.78
VIT-S/16 full
0.89
0.78
Selected Kronecker-Spars, 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 x 4096 (hidden size d 1024). The applied KS patterns are:
म = (1, 64, 256, 16),
πº (64, 64, 64, 1),
25

---

Fast Inference with Kronecker-Sparse Matrices
vielding the factorization W K2 These patterns selected following the same rationale for ViT-S/16: they 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 Py Torch using the batch-size-fiss, memory layout (batch-size-fing 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 A A100 GPU with 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 Precissoo
We now reproduce the benchmark from Section 5 in half-precision. The half-precision equivalents of our key float-precision results -abless 2 and 6, Figure 6a, Figure 10, Figure 12, Figure 13, and Figure 14a are respectively shown in Tables 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-preciison (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- salf-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-prreiison the relative cost of the permutation steps in BMM grows even larger than in dout-prreissonn-ranggi from 40 % to 80 %, compared to 10 % to 50 % in flottpreeissonn-futttee motivating efforts to climinate these costly memory operations.
4) Figure 18 shows that the new KERNEL still yield up x2 speedups, but only for KS patterns where large portion of the runtime is spent on permutations cermutations-those with high values of the heuristic h(b, c) big with the FP32 case, where improvements already appear from hbb,c) 2 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).
> 0.02. This contrasts
bc
5) Table 10 shows that the current KERNEL remains competitive in batch-size-last under half-precision, but falls behind in siteh- 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 aalf-srecision hardware features such as tensor cores.
26

---

Fast Inference with Kronecker-Sparse Matrices
divided by min time time (kernell bmm, einsum, bsr)
30
20
min time timeldense, sparse)
10.
min time(einsum,, bsr) divided by time(bmm)
0.7
0.6
A
0.5
0.4
0.6
104
10a
10
107
10a
10"
10
10a
10a
10
x N
M x
Figure Speedup factor of min time(KERNEL,, BMM, BSR, SUM) compared to min time(DENSE, SPARSF) vs. the matrix size M N. Experiments are carried in half-precision.
Figure 16. Speedup factor of time(BMM) compared min time(FINSUM, BSR) vs. the matrix size M x N. Experiments are carried half-precision.
Table Percentage of 600 patterns (a, b, c, d) where go faster than the go2 in half-precision (denoted by time(a (11go)) of time(a g)2)), and the median acceleration factor in such cases (that is, the median ratio of algel ). For each implementation, we take the minimum time between the batch-size firs-first and the batch-size-lass memory layout. Experiments are carried half precision.
BMM
KERNEL
FINSUM BSR DENSE SPARSE
EINSUM BSR
min time
BMM
< min time
DENSE SPARSE
time(BMM) < min time
time KERNEL) < min time
EINSUM BSR
SPARSE
85.7 % (X8.14)
80.7 % (x1.52)
37.8 % (x1.5)
Table Across 600 patterns, table reports how often each algorithu the fastest, along with the corresponding median speedup. under each memory layout (butch-size--irs) and batch-size-last) in FP16.
Batch-size-firs: Median x faster
Batch-size-last
Comparison
Win rate
Win rate
Median x faster
min KERNEL, BMM, EINSUM, ]SR) < miE(DE SPARSE]
83.18 %
x6.31 x1.54
85.34
x8.14
BMM < min EINSUM, BSR, DENSE, SPARSE)
82.25 %
79.48
x2.65
KERNEL < min BMM, EINSUM, BSR, DENSE, SPARSE]
0.0 %
41.78 %
x x1.56
27

---

Fast Inference with Kronecker-Sparse Matrices
0.8
Relative time in memory
Relative time in memory rewritings in bmm
0.7
rewrikings in bmm
0.7
0.6
0.5
0.4
0.4
0.3
0.3
00002
0003
0.004
0.00
0006
2.000
0.01
0.02
0.0
0.04
0.002
005
00.00
0.00
001
0.03
D+C
bc
bc
(a) Batch-size-firss first
(b) Batch-size-last
Figure 17. Estimated relative time spent memory rewritings in BMM for the multiplication with K E Kx. for several π (a, b, c, d). We regroup patterns by their value + e)/(bc). and plot boxplot summarize corresponding measurements. Experiments are carried in half-preeisst precision.
min time (bmm, einsum, bsr)
divided by (iee(kernel)
0.3
0.8
0.7
0.5
0.4
0.3
0.002
3000
0.000
2004
0,005
0.006
0007
0.02
0.3
004
b+c
bc
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-fire and the batch-size-last memory layout. We regroup the (a, b, c, d) patterns by their value of (b + (bc), and use boxplot to summarize the corresponding measurements, Experiments are carried in half-precision.
28


Reviews:
[{"review_id": "R9CJSxQAQI", "reviewer": "Reviewer_z2VJ", "summary": "The paper \"Fast inference with Kronecker-sparse matrices\" focuses on optimizing matrix multiplication algorithms for Kronecker-sparse matrices, which are used in neural networks to reduce parameters while maintaining accuracy. The main contributions include:\n- Benchmarking and Optimization: The authors benchmark existing GPU algorithms for multiplying Kronecker-sparse matrices and identify that up to 50% of runtime is spent on memory rewriting operations. They propose a new tiling strategy implemented in a CUDA kernel to reduce these memory transfers.\n- New CUDA Kernel: The kernel achieves a median speedup of ×1.4 and reduces energy consumption by 15% compared to baseline implementations.\n- Broader Impact: The new kernel accelerates the inference of neural networks, such as Vision Transformers (ViTs), by replacing dense layers with Kronecker-sparse matrices. This results in significant speedups in fully-connected layers.\n\nUpdate after rebuttal:\n\nI would like to sincerely thank the authors for taking the time to address all the issues I raised. However, as I pointed out in the initial review, the design and evaluation are quite limited in the single platform and software stacks. Even the authors claimed that they would provide the example in OpenCL, but it is hard for reviewers to evaluate without actually seeing the related implementation and experiments. So, for the current version, it may not meet the acceptance threshold for ICML.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review.\n\n# Regarding your questions\n\n1. > Extending the kernel to other hardwares / Limited Hardware Scope \n\nWe translated the CUDA kernel to openCL, so the kernel can now be used on other hardwares such as AMD GPUs or CPUs. The CUDA and openCL codes have also been integrated to an open-source python package available online. The revision will mention that and provide links to the (non-anonymous) repositories.\n\n2. > Extending the heuristic beyond Kronecker-sparsity\n\nThe heuristic is specifically tailored to Kronecker-sparse matrices with patterns (a, b, c, d), in order to predict the improvement in terms of memory transfer brought by the proposed kernel, compared to PyTorch baselines such as bmm. So there is no general and straightforward way to adapt it to other forms of sparsity and sparse multiplication algorithms. Nevertheless, counting the number of memory versus algebraic operations in a given sparse matrix multiplication for a certain sparsity pattern, and putting that in relation to its efficiency on a given hardware, is a general principle that might help in broader scenarios beyond Kronecker-sparse factors.\n\n3. > Influence of batch-size-last versus batch-size-first on other neural network operations\n\nWe agree that this is an interesting open question arising from the paper's results. For pointwise operations (e.g., activation functions like ReLU), changing the batch position is not expected to have any impact (and we indeed observed that internally). To investigate the impact of the batch position on other operations, it requires to carefully rewrite highly optimised kernels to convert them to batch-size-last, which requires time and expertise. Such a study falls out of the scope of what we claim to study (Kronecker linear layers), and is left to future work.\n\n4. > Measuring speedups in other architectures, beyond Vision Transformers\n\nThere is indeed an opportunity to extend the observations made in this paper to other architectures such as LLMs or RNNs, as they contain large linear layers where introducing Kronecker-sparsity is expected to speedup inference. The case of CNNs 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 option—replacing the convolutional kernel K with a Kronecker-sparse matrix—is limited because of the small size of typical kernels (e.g., 7×7). The second option—recasting convolutions as matrix products with a butterfly-structured weight matrix as in [1]—requires in practice costly folding/unfolding operations on the inputs and outputs, making it impractical in our view.\n\n5. > Generalisation of the speedups to other hardwares\n\nThe claimed scope of the paper is to focus on NVIDIA GPUs as they are the most common hardware used in AI clusters. It is left open to explore the opportunities and challenges related to benchmarking and optimising the kernel on other hardwares. Since the paper comes with an openCL version of the kernel, people with specific hardware needs can now easily include it in their benchmark.\n\n6. > Kronecker-sparsity versus other forms of sparsity\n\nStructured sparsity is clearly to be favored over unstructured sparsity for time and energy performance. Indeed, knowing the structure of the support in advance helps a lot to design efficient algorithms. Structured sparsity has also better theoretical proven properties (e.g., the functional space is closed and there always exists a minimizer of the loss for structured sparse networks, while it is not the case for unstructured ones, which might cause the solution to diverge, see theorem 4.2 in [2]). More specific comparisons would depend on the task, hardware, model and sparsity structure at hand.\n\n# Regarding the other points you mentioned\n\n7. > Formal proof of the algorithm equivalence, background on permutation shuffles matrix\n\nThank you for the suggestions, the revision will include both. \n \n8. > Energy measurements on V100 while time measurement on A100\n\nThe pyJoules package used in the paper to measure energy consumption is unfortunately not yet compatible with A100 GPUs.\n\n9. > On the robustness of the benchmark for different matrix sizes\n\nWe believe that the problem of robustness that you mention does not arise here, because the benchmark already covers all relevant matrix sizes encountered in practice. Indeed, it covers matrix sizes from 102×102 to 131072×131072 (6 orders of magnitude), while sizes typically used in transformers range from 500×500 to 15000 x 15000, and go up to 53248 in the largest models like Llama 3-405B [3].\n\n# References \n\n[1] Deformable Butterfly: A Highly Structured and Sparse Linear Transform. Lin et al NIPS 2022\n\n[2] Does a sparse ReLU network training problem always admit an optimum? Le et al NIPS 2023.\n\n[3] The Llama 3 Herd of Models, 2024.", "claims_and_evidence": "The submission provides extensive evidence to support its claims, primarily through benchmarks and theoretical analyses. However, some aspects could be scrutinized for clarity and robustness:\n\nBenchmarking Methodology: The paper presents a comprehensive benchmarking framework that compares various implementations of Kronecker-sparse matrix multiplication. The evidence is convincing, as it includes multiple scenarios and configurations, such as different memory layouts and precision levels (float and half-precision). However, the choice of specific hardware (NVIDIA GPUs) might limit the generalizability to other architectures.\n\nHardware Dependency: The performance benefits might not generalize equally across different hardware platforms.\n\nApplicability to Other Architectures: The focus on ViTs leaves room for further research on how well the approach works with other neural network architectures.\n\nRobustness Across Different Input Sizes: While the benchmark covers a range of sparsity patterns, the performance with very large or very small matrices might require additional validation.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria in the paper on \"Fast inference with Kronecker-sparse matrices\" appear well-suited for the problem at hand, which is optimizing matrix multiplication algorithms for Kronecker-sparse matrices in neural networks.\n\nThe benchmarks are conducted on NVIDIA GPUs. Extending the evaluation to other hardware platforms, such as AMD GPUs or CPUs, could provide a more comprehensive understanding of the methods' applicability.", "theoretical_claims": "Formal Proof for Algorithm Equivalence: While the paper explains the equivalence between Algorithm 1 and Algorithm 2, a formal proof might be beneficial for readers seeking rigorous mathematical validation.\n\nRobustness of the Heuristic: The heuristic for efficient patterns is empirically validated but might benefit from further theoretical analysis to ensure its applicability across different scenarios or hardware platforms.", "experimental_designs_or_analyses": "Hardware Dependency: While the paper focuses on NVIDIA GPUs, extending the benchmark to other hardware platforms could provide broader insights into the applicability of the proposed methods.\n\nGeneralizability Across Architectures: The paper primarily focuses on Vision Transformers. Investigating how well the new kernel performs with other neural network architectures (e.g., CNNs or RNNs) could enhance its utility.\n\nStatistical Analysis: The benchmark results are presented using medians and interquartile ranges. While this provides a good overview, additional statistical analysis (e.g., hypothesis testing) might further validate the significance of the observed improvements.\n\nEnergy Consumption Measurements: The energy measurements are conducted on a different GPU (V100) than the time benchmarks (A100). While this is noted, ensuring consistency across all measurements could strengthen the conclusions regarding energy efficiency.", "relation_to_broader_scientific_literature": "The paper builds on the concept of butterfly matrices, which are structured matrices that can be expressed as products of sparse factors with specific sparsity patterns, often described by Kronecker products. Butterfly matrices have been used to accelerate linear transforms like the Discrete Fourier Transform (DFT) and the Hadamard Transform.\n\nPrior work has shown that replacing dense matrices with sparse or structured matrices can improve neural network efficiency. For example, Dao et al. demonstrated that using butterfly matrices can speed up neural network training.\n\nThe paper introduces a new tiling strategy for matrix multiplication with Kronecker-sparse matrices, implemented in a CUDA kernel. This approach reduces memory transfers between different levels of GPU memory, leading to improved time and energy efficiency.\n\nSimilar heuristics have been used in other contexts to optimize sparse matrix operations, but the specific application to Kronecker-sparse matrices is novel and contributes to the broader literature on efficient neural network design.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Limited Hardware Scope: The paper focuses primarily on NVIDIA GPUs, which might limit the generalizability of the results to other hardware platforms like AMD GPUs or CPUs.\n\nOriginality of Heuristic: While the heuristic for efficient Kronecker-sparsity patterns is useful, it is based on a relatively straightforward analysis of memory operations. Further theoretical justification or exploration of its applicability beyond the current context could enhance its originality.\n\nClarity in Some Technical Details: Some sections, such as the explanation of perfect shuffle permutations, might be challenging for readers without a strong background in linear algebra. Additional explanations or references could improve clarity for a broader audience.\n\nBroader Applicability: The paper primarily focuses on Vision Transformers. Exploring how the new kernel performs with other neural network architectures could further demonstrate its significance and versatility.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. How do you envision extending the new CUDA kernel to work efficiently on other hardware platforms, such as AMD GPUs or CPUs? Are there any specific challenges or opportunities you foresee in this process?\n\n2. The heuristic based on the ratio (b+c)/bc is useful for identifying efficient Kronecker-sparsity patterns. Could you elaborate on how this heuristic might be refined or extended to accommodate different types of sparse matrices or computational contexts?\n\n3. The paper highlights the impact of memory layout (batch-size-first vs. batch-size-last) on performance. How do you think this might influence the design of other neural network operations, and are there opportunities for further optimization in this area?\n\n4. While the paper demonstrates significant speedups in Vision Transformers, what potential exists for applying these techniques to other neural network architectures, such as CNNs or RNNs? Are there specific challenges or opportunities in these contexts?\n\n5. The new kernel not only improves time efficiency but also reduces energy consumption. Could you discuss how these energy savings might be further optimized or generalized across different hardware platforms?\n\n6. How do you see Kronecker-sparse matrices comparing to other sparsity techniques, such as unstructured sparsity or other forms of structured sparsity? Are there scenarios where one might be preferred over the others?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5DBa7c2oB8", "reviewer": "Reviewer_nHJu", "summary": "This paper aims to speedup DNN inference with kronecker-sparse matrices by optimizing GPU memory accesses via customizing the CUDA kernels. The paper has made three key contributions: (1) analyzing the time and energy efficiency of existing implementations for multiplying kronecker-sparse matrices; (2) proposing a new tiling strategy in a new CUDA kernel implementation which reduces expensive GPU memory accesses; and (3) introducing a heuristic model that describes the time and energy efficiency. The experimental results have shown the proposed methods achieve 1.4x median speedup and 15% energy reduction.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review.\n# Regarding your questions\n1. >How did we measure the time spent on memory rewritings\n\nWe compared with the execution time where we removed the permutations/memory rewritings part, i.e. lines 1 and 3 in algorithm 1 (details in appendix B.2).\n\n2. >Implications on other transformer sizes, e.g. larger transformers like GPT/Llama\n\nWhile we benchmark end-to-end latency on a transformer involving matrices of size 768 x 768, note that the benchmark covers sizes between 102 x 102 and 131072 x 131072, largely covering the wide range used in practice, even the ones in large transformers (e.g. 53248 x 53248 in Llama3-405B). Since the speedup of the kernel increases with the matrix size (section 4), our resource-limited evaluation on a small transformer actually corresponds to one of the most challenging setup to observe a speedup. The revision will make this explicit.\n\n3. > Cache hit ratio\n\nWe measured the cache hit as suggested by the reviewer and will add a discussion about it.\n# Regarding the other points you mentioned\n4. >How Kronecker-sparse matrices are used in practice (motivation)\n\nKronecker-sparse matrices are the building block of Butterfly matrices. The literature has proposed using them in a variety of way in neural networks (Table 1) and the benchmark considers patterns aligned with these different use cases. Therefore, the potential misalignment problem, where studied patterns would be different to the ones used in practice, does not arise here. The revision will make it clear.\n\n5. >Relation to similar issues in sparse 3d convolutions\n\nA relevant analogy with this literature is that the new Algorithm 2 enables the implementation of a tiling strategy that optimises the dataflow in an analog way to how recent works [2,3,4] built on top of TorchSparse [1]. TorchSparse has three kernel calls: gather, multiplication, scatter. The subsequent works optimised this by overlapping memory and compute operations. For the sparse problem we consider, the dataflow optimisation enabled by Algorithm 2 has the same flavour: the kernel can now coalesce the input/output permutations with the multiplication part (Figure 3).\n\n6. >Specify whether the kernel is input/output or weight-stationary\n\nThe kernel is output-stationary, similarly to the dense cutlass kernels [5] and sparse 3d convolutional kernels [2,3,4] that have a similar dataflow. The revision will include a code template and explain this design choice as follows.\n* *Weight-stationary is not attractive*. In our Kronecker 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.\n* *Input-stationary is not attractive*. Due to the non-consecutive memory accesses involved (Figure 5), 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 the pattern $(a,b,c,d)$, which determines the distance between consecutive data elements that should be loaded together when considering one of the dense subproblems (Figure 5). 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) [5]. In contrast, output-stationarity avoids this issue, hence our choice.\n\n7. >Relation to sparse tensor compilers\n\nThese compilers efficiently handle unstructured or simple block/tile sparsity, but Kronecker-sparse matrices feature a block-diagonal pattern with $b\\times c$ sub-blocks further refined by $d\\times d$ identity matrices. Capturing such nested sparsity would require combining an outer block-sparse format with an inner identity layout—something not supported off-the-shelf. Moreover, our Algorithm 2 leads to a tiling strategy with *non-contiguous* tiles along certain axes to reduce memory operations, a design choice that these compilers cannot accommodate.\n\n8. >On adding results to show how the sparsity level affects speedups\n\nThanks for the suggestion, the revision will include a graph showing how the speedup increases with the sparsity level (percentage of zero, ranging from 0 to 99% in the benchmark). The revision will also explain that the 12% of the 600+ patterns where the baselines still outperform the kernel correspond to cases with a large density of nonzero or a small value of the proposed heuristic, i.e., a small value of (b+c)/bc.\n# References\n[1] TorchSparse: Efficient Point Cloud Inference Engine. Tang et al 2022\n\n[2] Torchsparse++: Efficient point cloud engine. Tang et al CVPR 2023\n\n[3] SpConv, Yan 2022\n\n[4] SECOND: Sparsely Embedded Convolutional Detection. Yan et al 2018\n\n[5] github.com/NVIDIA/cutlass/blob/main/media/docs/efficient_gemm.md", "claims_and_evidence": "While the paper has shown evidences that the proposed methods outperform existing frameworks, its evaluation results are not enough to fully support the claims. Specifically, the paper has never conducted ablation studies on how the percentage of non-zero elements affect the performance of the proposed kernels. This is important because when sparse matrix is used in practice, different level of sparsity might be used for better accuracy.", "methods_and_evaluation_criteria": "The paper claims that they are able to support general inference of transformers. However, they only evaluate the vision transformers of small sizes, which could have very different trade-offs compared to transformer-based language models. Thus, it is obscure how the proposed method would perform in practice.", "theoretical_claims": "I do not find any problems in the theoretical analysis in the paper but my expertise could be limited.", "experimental_designs_or_analyses": "Please refer to the \"Methods and Evaluation Criteria\" section of the review.", "relation_to_broader_scientific_literature": "While the paper has pointed out an interesting way to tackle the memory inefficiency in existing implementation of kronecker-sparse matrix multiplication, the paper should have discussed its relationship to other types of sparse matrix multiplication problems. In particular, similar issues have been seen in similar problems such as 3D sparse convolutions [1], where the permutations also needs to be done on inputs and outputs. The authors should have discussed the trade-offs of different design choices (e.g. input/output/weight stationary), and the reason why they pick up a specific one.\n\n[1] Tang, Haotian, et al. \"Torchsparse++: Efficient point cloud engine.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.", "essential_references_not_discussed": "The paper should also discuss how their methods related to (sparse) tensor compilers, since these compilers can also be easily used to speedup sparse-kernel inference. Specifically, the paper should discuss why kronecker-sparse matrix multiplication is considered as a challenging problem and cannot just be considered as a special case for sparse tensor compilers:\n\n* Ye, Zihao, et al. \"Sparsetir: Composable abstractions for sparse compilation in deep learning.\" Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3. 2023.\n* Guan, Yue, et al. \"Fractal: Joint multi-level sparse pattern tuning of accuracy and performance for DNN pruning.\" Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3. 2024.", "other_strengths_and_weaknesses": "The motivation of the paper is also lacking. Specifically, the paper has never discussed how the kronecker-sparse matrix is used in practice. This is important because different sparsities levels could result in different trade-offs and design choices.", "other_comments_or_suggestions": "I feel the paper is in general well written. However, the section captions could be simplified as it makes more clear for the readers to understand the topic of each section.", "questions_for_authors": "* How do you measure the memory access time in general (see line 209, \"We find that the memory rewritings can take up to 45% of the total runtime\")\n* How the proposed method speedup the inference of other types of networks such as GPT or Llama models?\n* What is the cache hit ratio of reading/writing global memory of the proposed CUDA kernels?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8acflHg3mO", "reviewer": "Reviewer_nKiB", "summary": "This paper proposes a novel CUDA kernel designed to accelerate neural network inference using Kronecker-sparse matrices. These matrices, characterized by sparsity patterns derived from the Kronecker product, offer a structured alternative to traditional dense matrices in neural networks. By optimizing memory access and reducing redundant operations, the proposed kernel achieves a 1.4× speedup and a 15% reduction in energy consumption compared to existing approaches, demonstrating its effectiveness in enhancing computational efficiency.\n\n## update after rebuttal\nI have read author's response, but I still believe that the conversion cost could be high when the proposed method operates on activations—such as in self-attention operations (e.g., Q @ Kᵀ, Attention-Score @ V)—which was my main concern. Since my concern regarding this point have not been fully resolved.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review.\n\n# Regarding your questions\n\n1. > End-to-end latency results in ViT\n\nTable 4 already provides an end-to-end latency result, showing a 22% relative time gain on a vision transformer when using the kernel. If you actually meant to ask about the *absolute* measurements in seconds rather than the *relative* comparison, we would be happy to add them upon request.\n\n2. > Comparison to other “existing Kronecker-sparse matrix calculation techniques” / “structured matrix optimization kernels” such as Monarch Transform or Butterfly Transform\n\nThe implementation associated with the Monarch Transform [1,2] corresponds to the “bmm” baseline in the benchmark. The official code associated with the Butterfly Transform [3,4] is no longer maintained and we were unable to make it work. However, we tested a faithful reproduction internally but found it significantly slower than the other baselines, so we chose not to include it in the benchmark. The revision will make this explicit.\n\n# Regarding the other points you mentioned\n\n3. > Further strengthen the study by extending it to other hardwares\n\nWe agree that exploring the opportunities and challenges related to benchmarking and optimising the kernel on other hardwares is an exciting open avenue that is now raised by this work. To support further exploration, we will release an OpenCL version of the kernel, enabling users with specific requirements to test it on platforms such as AMD GPUs or CPUs.\n\n4. > CUDA kernel code not included\n\nThe final version will include a template of the kernel code and a link to the open-source (non-anonymous) repository.\n\n5. > Computational cost of converting a dense matrix to the butterfly structure\n\nAlthough this issue falls outside the claimed scope of the paper—which is to study Kronecker-sparse matrix multiplication on GPUs and to showcase its potential to accelerate the inference of models having Kronecker-sparse matrices (e.g., models trained from scratch with such matrices, or dense models replaced by Kronecker-sparse ones after training, with potential subsequent fine-tuning)—it is worth mentioning that approximating a given target matrix of size $m \\times n$ by a product of Kronecker-sparse factors can be done efficiently in roughly $\\mathcal{O}(mn)$ time [5].\n\n# References\n[1] Monarch: Expressive Structured Matrices for Efficient and Accurate Training, Dao et al, PMLR 2022.\n\n[2] Monarch mixer: A simple sub-quadratic GEMM-based architecture. Fu et al. NeurIPS, 2023.\n\n[3] Butterfly transform:An efficient FFT based neural architecture design. Vahid et al. CVPR, 2020.\n\n[4] Learning fast algorithms for linear transforms using butterfly factorizations. Dao et al. ICML, 2019\n\n[5] Butterfly factorization with error guarantees. Le et al., preprint, 2024.", "claims_and_evidence": "- The study evaluates existing GPU implementations, identifies inefficiencies in memory access, and proposes a new CUDA kernel optimized for Kronecker-sparse matrix multiplications.\n- The paper is well-supported by empirical benchmarks and theoretical analysis, with the 1.4× speedup and 15% energy reduction demonstrated through extensive experiments across various sparsity patterns.\n- However, the study could be further strengthened by evaluating its performance on different hardware platforms (e.g., AMD GPUs, CPUs, FPGAs) and offering deeper insights into automated sparsity pattern selection.", "methods_and_evaluation_criteria": "- The paper proposes a new tiling strategy for matrix multiplication with Kronecker-sparse matrices, which reduces memory transfer overhead in GPU computations.\n- The study carefully selects a range of Kronecker-sparsity patterns and tests them across different conditions, ensuring that results are not limited to specific cases.\n- The evaluation extends to practical scenarios like transformer inference acceleration, demonstrating the broader utility of the proposed method.", "theoretical_claims": "- The paper demonstrates that existing methods incur significant GPU memory access costs.\n\b- It quantifies the memory rearrangement cost as (b + c) / (bc) , arguing that a higher value of this ratio indicates the inefficiency of existing approaches.\n- However, the theoretical analysis presented in the paper does not clearly establish whether the proposed method guarantees consistent performance improvements across all Kronecker-sparse patterns.", "experimental_designs_or_analyses": "- The benchmarking of execution time and energy consumption on various GPU implementations is relevant and well-justified.\n- The comparison with existing PyTorch implementations (bmm, bsr, einsum) and generic dense/sparse approaches provides a meaningful baseline.\n- However, there is no direct comparison with other existing Kronecker-sparse matrix calculation techniques.", "relation_to_broader_scientific_literature": "- None", "essential_references_not_discussed": "- None", "other_strengths_and_weaknesses": "Strengths\n- The paper introduces a new tiling strategy for GPU matrix multiplication, reducing memory access overhead.\n- The proposed CUDA kernel achieves up to 1.4× speedup and 15% energy reduction compared to existing methods.\n- Demonstrates that the optimized kernel can accelerate vision-transformer model inference.\n\nWeaknesses\n- The paper does not compare its CUDA kernel with other structured matrix optimization kernels (e.g., Monarch matrices, Butterfly Transform).\n- The study assumes W (weight matrix) is already in Kronecker-sparse format, but does not discuss the cost of converting a dense matrix to this structure. If the transformation cost is high, the practical benefit of the proposed optimization may be reduced in self-attention operation (ex, Q @ K^T, Attention-Score @ V).\n- Only the speed difference by unit operation is shown and the End-to-End model (ViT) latency results are not provided.", "other_comments_or_suggestions": "- None", "questions_for_authors": "- Can you provide end-to-end latency comparison results measured in ViT?\n- Can you provide a latency comparison with other structured matrix optimization kernels (e.g., Monarch matrices, Butterfly Transform) rather than a comparison with the code implemented in pytorch?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1tCwo3qDXI", "reviewer": "Reviewer_XhSi", "summary": "This paper presents the first energy and time benchmarks for the multiplication of Kronecker-sparse matrices. These benchmarks reveal that specialized sparse matrix multiplication implementations spend up to 50% of run time on memory rewrite operations. As a remedy, the authors propose a new tiling strategy for Kronecker-sparse matrix multiplication achieving a median speed-up of 1.4x while also cutting energy consumption by 15%.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We address your points below.\n\n1. >Are there any other sparse related works relevant for this benchmark?\n\nThe benchmark includes all the relevant baselines we are aware of. The revision will include an additional discussion clarifying how our work relates to a few other areas of sparse matrix research (such as sparse 3d convolutions or sparse tensor compilers, as suggested by other reviewers), even though these related works do not offer relevant implementations to be included in the benchmark.\n\n2. > On the potential existence of a sparsity threshold under which previous baselines are better\n\nYou are absolutely right, thank you for mentioning that. The 12% of the cases where the baselines are still better correspond to patterns that have a high density of nonzero or a small value of the proposed heuristic (the ratio (b+c)/bc). The revision will contain a plot that shows how the speedup increases with the sparsity level (percentage of zeros).\n\nWe also thank the reviewer for spotting typos.", "claims_and_evidence": "That authors appear to be making a few main claims. \n\n1. The first claim is that previous approaches to handling Kronecker-sparse matrix multiplications spend up to 50% of their total runtime on memory rewriting operations. This is supported empirically by a set of results reported by the authors.\n\n2. The second claim is that the time spent on memory rewriting operations can be attributes to the structure adapted by most such algorithms, which can be dropped by a factor of three by tiling, essentially coalescing the read/write operations. The authors have supported this claim analytically.\n\n3. The third claim concerns empirical improvements, where the authors have claimed they new Kronecker-sparse multiplication algorithm to lead to speeds ups of up to 1.4x and energy savings of 15%, as well as improve the efficiency of neural network inference. The authors have support all of these claims empirically.", "methods_and_evaluation_criteria": "Yes, the methods and evaluation criteria make sense. The only thing is I was expecting to see comparison to more sparse matrix multiplication baselines but I am not very familiar with the area, so maybe such specialized algorithms do not exist or aren't used in practice?", "theoretical_claims": "Yes, I went over proofs and theoretical claims and they appear to make sense to me", "experimental_designs_or_analyses": "I found the experimental design and analysis to be sound. \n\nThe authors start with a hypothesis that current matrix multiplication algorithms used with Kronecker-sparse matrices spend a significant portion of their run time on memory rewriting, a hypothesis which they empirically validate on 600 different sparsity patterns. They then devised a theoretically-equivalent algorithm with a theoretically lower number of memory rewriting operations. They then validate that they new algorithm is indeed faster than the current used algorithms as well as more energy efficient, due to having to perform less energy on memory rewritings.", "relation_to_broader_scientific_literature": "The contribution put forth in this paper is very timely and application to machine learning community, but I suspect also to the algorithms and scientific computing communities. This work is essentially on how to speed up matrix multiplication when we know of, or impose, a specific sparsity structure on one of our input matrices.", "essential_references_not_discussed": "Not aware of any.", "other_strengths_and_weaknesses": "- I found the paper to be very well written despite the very technical subject matter. I greatly appreciated the top-down manner through which the problem was decomposed, the hypothesis validated and the method proposed to tackle the shortcoming of the current approaches.\n\n- I like how the authors did not stop at comparing their algorithm to previously proposed algorithms in a synthetic setting, but rather also showed that practical settings, such as inference in ViT architectures can benefit greatly in terms of speedup.\n\n- To my knowledge, this contribution is novel, and I foresee it having a sizable impact in the community's effort to improve the efficiency of inference in deep models.\n\n- I was expected to see quite a bit more in terms of related work given that sparsity is currently a highly-researched area", "other_comments_or_suggestions": "Typos: Section 5 326-363 line/row -> column?", "questions_for_authors": "The only question that I have is: What is the point of diminishing returns? My understanding is there must be a threshold whereby if the matrix does not exhibit enough sparsity then the baselines might perform better? Am I correct in my assumption? I believe this might also be suggested by the The third column of Table 3 that  shows that kernel is slower than all other baselines 12% of the tested patterns? Any idea why that might be the case?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces a specialized CUDA kernel for accelerating Kronecker-sparse matrix multiplications, a structured form of sparsity increasingly used in neural networks. The authors identify a key inefficiency in existing implementations — excessive memory rewriting — and propose a new tiling strategy to reduce this overhead. The result is a median 1.4× speedup and 15% energy savings, validated across over 600 sparsity patterns and demonstrated in accelerating inference for Vision Transformers (ViTs) on NVIDIA GPUs. Through the rebuttal, the authors provided detailed clarifications and extended their implementation to OpenCL to address generalization concerns. 

Overall, the paper makes a technically sound and practically useful contribution to structured sparsity and efficient inference. It is well-executed, timely, and relevant. While some concerns remain about the generality and integration with broader toolchains, the core idea is novel and impactful. 

Given the strength of the engineering contribution, the quality of the benchmarks, and the clarity of the rebuttal, this paper meets the bar for acceptance. Minor revision would further strengthen the work for publication.
Decision: Accept (poster)