Insum: Sparse GPU Kernels Simplified and Optimized with Indirect Einsums

TL;DR

This work tackles the challenge of crafting high-performance sparse GPU kernels by recasting sparse computations as format-conscious indirect Einsums and implementing Insum, a compiler that lowers these expressions to dense PyTorch code and leverages the PyTorch compiler with Tensor Core support. It introduces fixed-length GroupCOO and BlockGroupCOO formats to enable efficient indirect indexing, and extends TorchInductor with a tl.dot path and Lazy Broadcasting to achieve fully fused gather–Einsum–scatter kernels. Across structured and unstructured SpMM, point-cloud sparse convolution, and equivariant tensor products, the approach delivers 1.14x–3.81x speedups and reduces lines of code by 202x–4491x compared to hand-written implementations, while matching or surpassing state-of-the-art hand-optimized kernels. The method significantly lowers engineering effort and generalizes across diverse sparse workloads, representing a practical and scalable path toward GPU-accelerated sparse computation.

Abstract

Programming high-performance sparse GPU kernels is notoriously difficult, requiring both substantial effort and deep expertise. Sparse compilers aim to simplify this process, but existing systems fall short in two key ways. First, they are primarily designed for CPUs and rarely produce high-performance GPU code. Second, when computations involve both sparse and dense regions, these compilers often fail to optimize the dense portions effectively. In this paper, we propose a new approach for expressing sparse computations. We start from format-agnostic Einsums over sparse tensors and rewrite them into format-conscious indirect Einsums, which explicitly encode format information by mapping sparse data and metadata onto dense tensor operations through indirect indexing. To execute indirect Einsums, we introduce the Insum compiler, which generates efficient GPU code for these Einsums by lowering to the PyTorch compiler, extended to better support Tensor Core-enabled indirect Einsums. We also present two fixed-length sparse formats, GroupCOO and BlockGroupCOO, designed to fit naturally with indirect Einsums. Our approach achieves 1.14x to 3.81x speedups across a range of sparse GPU applications while reducing lines of code by 202x to 4491x compared to hand-written implementations.

Figures (13)

• Figure 1: SpMM example ($C_{m,n} = A_{m,k} * B_{k,n}$) and various sparse formats for the matrix $A$: COO, ELL, and CSR.
• Figure 2: COO SpMM : $C_{AM_p, n} = AV_p * B_{AK_p, n}$
• Figure 3: Overview of our system. Given a sparse problem, users express it in an indirect Einsum. This expression is then translated into PyTorch program using Insum compiler, and then lowered into a Triton kernel via our extended torch.compile.
• Figure 4: GroupCOO format with grouping on the M dimension by group sizes 2 and 3. Grouping with the maximum number of nonzero per row (size = 3) yields the ELL format.
• Figure 5: BlockCOO SpMM: $C_{AM_p,\,bm,\,n} = AV_{p,\,bm,\,bk} * B_{AK_p,\,bk,\,n}$, where bm and bk index within each block.