Research map / Architectures and efficiency / RNNs and normalization
Fast Inference with Kronecker-Sparse Matrices
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a fused, output-stationary GPU kernel for Kronecker-sparse matrix multiplication that avoids the input/output tensor permutations of existing Monarch-style implementations, yielding a median 1.4x FP32 speedup over the best baseline across 600 KS patterns.
Problem
Kronecker-sparse matrix multiplication on GPUs (used in Butterfly and Monarch layers) is slowed by memory-bound tensor permutations, which can take up to 50% of runtime in existing implementations.
Approach
Reformulates KS matmul so each thread block computes one dense sub-product directly from non-contiguous global memory tiles, fusing the permute-GEMM-permute pipeline into one output-stationary kernel that cuts global memory traffic about threefold.
Key results
- In FP32 across 600 KS patterns, the fused kernel wins on 85.5% of patterns with a median 1.39x speedup over the best baseline.
- Energy consumption is reduced by a median 15% in FP32, with reductions on 72% of patterns.
- End-to-end FP32 latency falls by 22% on ViT-S/16 (1.28x vs dense) and 16% on GPT-2 Medium (1.19x vs dense).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Fuse the input/output permutations with GEMM into one output-stationary kernel | KS matrix multiplication · 600 KS patterns, FP32 | speedup over best baseline: 1.39 | Helped |
| Fuse the input/output permutations with GEMM into one output-stationary kernel | KS matrix multiplication · 600 KS patterns, FP32 | energy consumption: 15% | Helped |
| Replace dense linear layers with KS factors and the fused kernel | ViT-S/16 · ViT-S/16, FP32 | latency: 0.89 -> 0.78 | Helped |
| Replace dense linear layers with KS factors and the fused kernel | GPT-2 Medium · GPT-2 Medium, FP32 | latency: 0.88 -> 0.84 | Helped |
| Switch from batch-size-first to batch-size-last layout (ablation) | Fused output-stationary kernel · KS matrix multiplication | speedup: 2 | Helped |
Surprising findings
- Switching the baseline to batch-size-last does not change BMM runtime, while the fused kernel gains about 2x from batch-size-last layout.
- Gains are smaller in FP16/FP8 where tensor cores are used, according to the authors' limitations discussion.
Limitations
- Gains are reported mainly in FP32; advantage is smaller in FP16 where tensor cores are used.
- Batch-size-last layout gives the best performance but PyTorch defaults to batch-size-first, which requires costly transpositions.
- Kernel performance depends on batch size, pattern and hand-tuned hyperparameters.
Experimental setup
Benchmarks: 600 KS patterns benchmark (time and energy) ViT-S/16 inference latency GPT-2 Medium inference latency
Metrics: median speedup win rate latency energy consumption
Baselines: BMM (Monarch batched GEMM) BSR (block-sparse) einsum dense linear CSR sparse linear
Builds on: Monarch Butterfly matrices BSR sparse format CUTLASS
Closest papers by meaning
- Giga-scale Kernel Matrix-Vector Multiplication on GPU (NeurIPS 2022 · similarity 0.88)
- Kernel Methods Through the Roof: Handling Billions of Points Efficiently (NeurIPS 2020 · similarity 0.86)
- Scalable Gaussian Processes with Latent Kronecker Structure (ICML 2025 · similarity 0.84)
- Improved Algorithms for Kernel Matrix-Vector Multiplication Under Sparsity Assumptions (ICLR 2025 · similarity 0.83)
- Schwarz–Schur Involution: Lightspeed Differentiable Sparse Linear Solvers (ICML 2025 · similarity 0.83)
- Searching for Efficient Linear Layers over a Continuous Space of Structured Matrices (NeurIPS 2024 · similarity 0.83)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.