Thank you for your review.

Regarding your questions

1. End-to-end latency results in ViT

Table 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.

2. Comparison to other “existing Kronecker-sparse matrix calculation techniques” / “structured matrix optimization kernels” such as Monarch Transform or Butterfly Transform

The 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.

Regarding the other points you mentioned

3. Further strengthen the study by extending it to other hardwares

We 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.

4. CUDA kernel code not included

The final version will include a template of the kernel code and a link to the open-source (non-anonymous) repository.

5. Computational cost of converting a dense matrix to the butterfly structure

Although 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].

References

[1] Monarch: Expressive Structured Matrices for Efficient and Accurate Training, Dao et al, PMLR 2022.

[2] Monarch mixer: A simple sub-quadratic GEMM-based architecture. Fu et al. NeurIPS, 2023.

[3] Butterfly transform:An efficient FFT based neural architecture design. Vahid et al. CVPR, 2020.

[4] Learning fast algorithms for linear transforms using butterfly factorizations. Dao et al. ICML, 2019

[5] Butterfly factorization with error guarantees. Le et al., preprint, 2024.

Thank you for your review.

Regarding your questions

1. How did we measure the time spent on memory rewritings

We 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).

2. Implications on other transformer sizes, e.g. larger transformers like GPT/Llama

While 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.

3. Cache hit ratio

We measured the cache hit as suggested by the reviewer and will add a discussion about it.

Regarding the other points you mentioned

4. How Kronecker-sparse matrices are used in practice (motivation)

Kronecker-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.

5. Relation to similar issues in sparse 3d convolutions

A 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).

6. Specify whether the kernel is input/output or weight-stationary

The 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.

• 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.
• 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.

7. Relation to sparse tensor compilers

These 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.

8. On adding results to show how the sparsity level affects speedups

Thanks 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.

References

[1] TorchSparse: Efficient Point Cloud Inference Engine. Tang et al 2022

[2] Torchsparse++: Efficient point cloud engine. Tang et al CVPR 2023

[3] SpConv, Yan 2022

[4] SECOND: Sparsely Embedded Convolutional Detection. Yan et al 2018

[5] github.com/NVIDIA/cutlass/blob/main/media/docs/efficient_gemm.md

Thank you for your review. We address your points below.

1. Are there any other sparse related works relevant for this benchmark?

The 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.

2. On the potential existence of a sparsity threshold under which previous baselines are better

You 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).

We also thank the reviewer for spotting typos.

Thank you for your review.

Regarding your questions

1. Extending the kernel to other hardwares / Limited Hardware Scope

We 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.

2. Extending the heuristic beyond Kronecker-sparsity

The 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.

3. Influence of batch-size-last versus batch-size-first on other neural network operations

We 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.

4. Measuring speedups in other architectures, beyond Vision Transformers

There 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.

5. Generalisation of the speedups to other hardwares

The 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.

6. Kronecker-sparsity versus other forms of sparsity

Structured 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.

Regarding the other points you mentioned

7. Formal proof of the algorithm equivalence, background on permutation shuffles matrix

Thank you for the suggestions, the revision will include both.

8. Energy measurements on V100 while time measurement on A100

The pyJoules package used in the paper to measure energy consumption is unfortunately not yet compatible with A100 GPUs.

9. On the robustness of the benchmark for different matrix sizes

We 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].

References

[1] Deformable Butterfly: A Highly Structured and Sparse Linear Transform. Lin et al NIPS 2022

[2] Does a sparse ReLU network training problem always admit an optimum? Le et al NIPS 2023.

[3] The Llama 3 Herd of Models, 2024.