# SPLAT: A framework for optimised GPU code-generation for SParse reguLar ATtention

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 62 references that are in the map.
- First posted on arXiv 23 Jul 2024.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Long Range Arena: A Benchmark for Efficient Transformers](https://mlsanity.com/paper/2011.04006/) (2020)
- [SCROLLS: Standardized CompaRison Over Long Language Sequences](https://mlsanity.com/paper/2201.03533/) (2022)
- [Blockwise Self-Attention for Long Document Understanding](https://mlsanity.com/paper/1911.02972/) (2019)
- [Sparse GPU Kernels for Deep Learning](https://mlsanity.com/paper/2006.10901/) (2020)
- [Long-Short Transformer: Efficient Transformers for Language and Vision](https://mlsanity.com/paper/2107.02192/) (2021)
- [Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions](https://mlsanity.com/paper/1802.04730/) (2018)
- [Monarch Mixer: A Simple Sub-Quadratic GEMM-Based Architecture](https://mlsanity.com/paper/2310.12109/) (2023)
- [Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Code](https://mlsanity.com/paper/1804.10694/) (2018)
- [SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning](https://mlsanity.com/paper/2207.04606/) (2022)
- [A Systematic Survey of General Sparse Matrix-Matrix Multiplication](https://mlsanity.com/paper/2002.11273/) (2020)
- [Heuristic Adaptability to Input Dynamics for SpMM on GPUs](https://mlsanity.com/paper/2202.08556/) (2022)
- [Sympiler: Transforming Sparse Matrix Codes by Decoupling Symbolic Analysis](https://mlsanity.com/paper/1705.06575/) (2017)

## Built on by (most cited first, 5 of 5)

- [HASTE: Hardware-Aware Dynamic Sparse Training for Large Output Spaces](https://mlsanity.com/paper/2606.01117/) (2026)
- [Accelerating Sparse Transformer Inference on GPU](https://mlsanity.com/paper/2506.06095/) (2025)
- [DCC: Data-Centric Compilation of Machine Learning Kernels for Processing-In-Memory Architectures](https://mlsanity.com/paper/2511.15503/) (2025)
- [Spira: Exploiting Voxel Data Structural Properties for Efficient Sparse Convolution in Point Cloud Networks](https://mlsanity.com/paper/2511.20834/) (2025)
- [TempoNet: Slack-Quantized Transformer-Guided Reinforcement Scheduler for Adaptive Deadline-Centric Real-Time Dispatchs](https://mlsanity.com/paper/2602.18109/) (2026)

## Abstract (arXiv)

Multi-head-self-attention (MHSA) mechanisms achieve state-of-the-art (SOTA) performance across natural language processing and vision tasks. However, their quadratic dependence on sequence lengths has bottlenecked inference speeds. To circumvent this bottleneck, researchers have proposed various sparse-MHSA models, where a subset of full attention is computed. Despite their promise, current sparse libraries and compilers do not support high-performance implementations for diverse sparse-MHSA patterns due to the underlying sparse formats they operate on. These formats, which are typically designed for high-performance & scientific computing applications, are either curated for extreme amounts of random sparsity (<1% non-zero values), or specific sparsity patterns. However, the sparsity patterns in sparse-MHSA are moderately sparse (10-50% non-zero values) and varied, resulting in existing sparse-formats trading off generality for performance. We bridge this gap, achieving both generality and performance, by proposing a novel sparse format: affine-compressed-sparse-row (ACSR) and supporting code-generation scheme, SPLAT, that generates high-performance implementations for diverse sparse-MHSA patterns on GPUs. Core to our proposed format and code generation algorithm is the observation that common sparse-MHSA patterns have uniquely regular geometric properties. These properties, which can be analyzed just-in-time, expose novel optimizations and tiling strategies that SPLAT exploits to generate high-performance implementations for diverse patterns. To demonstrate SPLAT's efficacy, we use it to generate code for various sparse-MHSA models, achieving geomean speedups of 2.05x and 4.05x over hand-written kernels written in triton and TVM respectively on A100 GPUs. Moreover, its interfaces are intuitive and easy to use with existing implementations of MHSA in JAX.