STen: Productive and Efficient Sparsity in PyTorch

As deep learning models grow, sparsity is becoming an increasingly critical component of deep neural networks, enabling improved performance and reduced storage. However, existing frameworks offer poor support for sparsity. Specialized sparsity engines focus exclusively on sparse inference, while general frameworks primarily focus on sparse tensors in classical formats and neglect the broader sparsification pipeline necessary for using sparse models, especially during training. Further, existing frameworks are not easily extensible: adding a new sparse tensor format or operator is challenging and time-consuming. To address this, we propose STen, a sparsity programming model and interface for PyTorch whose key design insight is the decoupling of sparsity layouts, operators, and sparsifiers into composable, first-class abstractions that users can independently define and combine. An automatic dispatch mechanism selects the best available sparse implementation and transparently falls back to dense execution, allowing STen to support virtually all sparsification methods while enabling rapid prototyping without sacrificing performance for optimized paths. We demonstrate the versatility of STen by expressing existing sparsification techniques within its abstraction, achieving a code size reduction of over 2×. Finally, we develop a novel, high-performance grouped n : m sparsity layout for CPU inference at moderate sparsity, accelerating end-to-end BERT BASE inference by up to 3.2×. STen brings high performance and ease of use, making sparsity readily accessible for existing PyTorch models.