Sparse GPU Kernels for Deep Learning

Trevor Gale; Matei Zaharia; Cliff Young; Erich Elsen

Source-linked AI summary

Sparse GPU Kernels for Deep Learning

Trevor Gale, Matei Zaharia, Cliff Young, Erich Elsen

arXiv:2006.10901v2Jun 18, 2020cs.LGcs.DCstat.ML

TL;DR

Deep-learning sparsity often fails to deliver GPU speedups because existing sparse kernels target much higher sparsity levels. This paper studies deep-learning sparse matrices and develops CSR-based SpMM and SDDMM kernels using application-specific properties and tiling techniques. The resulting sparse Transformer and MobileNet models achieve 1.2–2.1× speedups and up to 12.8× memory savings without sacrificing accuracy.

Problem
Existing GPU sparse kernels often cannot outperform dense computation at the moderate sparsity levels found in deep neural networks.
Method
The paper studies deep-learning sparse-matrix properties and develops CSR-based GPU kernels for SpMM and SDDMM using 1-dimensional tiling and related memory-alignment techniques.
Results
Sparse Transformer and MobileNet models achieve 1.2–2.1× speedups and up to 12.8× memory savings while matching dense-model accuracy.
Takeaways & Limitations
The kernels provide application-level acceleration and memory reductions for sparse deep-learning models without sacrificing accuracy.
Takeaways & Limitations
On large problems, kernel performance is limited by shared-memory bandwidth.

Abstract

from arXiv · show

Scientific workloads have traditionally exploited high levels of sparsity to accelerate computation and reduce memory requirements. While deep neural networks can be made sparse, achieving practical speedups on GPUs is difficult because these applications have relatively moderate levels of sparsity that are not sufficient for existing sparse kernels to outperform their dense counterparts. In this work, we study sparse matrices from deep learning applications and identify favorable properties that can be exploited to accelerate computation. Based on these insights, we develop high-performance GPU kernels for two sparse matrix operations widely applicable in neural networks: sparse matrix-dense matrix multiplication and sampled dense-dense matrix multiplication. Our kernels reach 27% of single-precision peak on Nvidia V100 GPUs. Using our kernels, we demonstrate sparse Transformer and MobileNet models that achieve 1.2-2.1x speedups and up to 12.8x memory savings without sacrificing accuracy.

I. INTRODUCTION

Deep-learning sparsity is difficult to accelerate on GPUs because moderate sparsity defeats kernels optimized for extremely sparse scientific matrices. The paper studies deep-learning sparsity and develops GPU techniques and kernels that improve sparse computation without imposing structured nonzeros.

Motivation: Moderate sparsity in deep neural networks prevents existing GPU sparse kernels from consistently outperforming dense computation.Existing kernels primarily target scientific matrices with 99%+ sparsity.
Approach: The paper targets sparse matrix–matrix multiplication and sampled dense–dense matrix multiplication directly on standard CSR-format matrices without restricting nonzero locations.This preserves unstructured sparsity while targeting operations widely used in neural networks.
Contributions: The authors identify favorable deep-learning matrix properties and introduce 1-dimensional tiling, subwarp tiling, and reverseoffset memory alignment.These techniques facilitate operand reuse, support vector memory instructions on misaligned sparse data, and enable specialized kernel variants.
Results: 3.58× and 2.19× geometric mean speedups over Nvidia cuSPARSE are achieved for SpMM and SDDMM, respectively, on Nvidia V100 GPUs.Top-performing problems reach 27% of single-precision peak.
Results: Sparse Transformer and MobileNet models achieve 1.2–2.1× end-to-end speedups and up to 12.8× memory savings while matching dense-model accuracy.The reported gains demonstrate application-level benefits from the specialized kernels.

II. SPARSE MATRICES IN DEEP LEARNING

The paper characterizes deep-learning sparse matrices and uses their properties to motivate GPU kernel design. Compared with scientific matrices, deep-learning matrices are less sparse, have longer rows, and exhibit more uniform row lengths.

Dataset: The dataset contains sparse ResNet-50 and Transformer weight matrices trained with four sparsity-inducing algorithms.Transformer models are restricted to those achieving above 20 BLEU on WMT14 English-to-German.
Metrics: The analysis measures row-length coefficient of variation, average row length, and sparsity.Coefficient of variation is standard deviation divided by mean row length; higher values indicate load imbalance.
Implications: Longer rows amortize startup and one-time costs, while lower row-length variation reduces load imbalance across rows.Lower sparsity can also increase opportunities to reuse operands through caches when nonzeros align across rows.
Results: Deep-learning matrices are 13.4× less sparse, have 2.3× longer rows, and have 25× less variation in row length than scientific matrices on average.The comparison uses the SuiteSparse Matrix Collection, containing 2,833 matrices from varied scientific workloads.
GPU Background: The GPU background describes thread blocks, shared memory, registers, occupancy, caches, and coalesced global-memory accesses relevant to kernel execution.The paper explains these architectural mechanisms before presenting its sparse operations and kernels.

A. Sparse Matrix–Matrix Multiplication Operation

The paper defines SpMM and SDDMM operations for deep-learning workloads, then designs a GPU SpMM kernel around one-dimensional output tiling to handle widely varying dense-matrix widths.

SpMM computes AB⇒C with sparse A in CSR format, while SDDMM computes AB ⊙C⇒D with sparse C and D.
Deep-learning applications use SDDMM for sparse-weight gradients and Transformer sparse attention, with the paper specializing its implementation to AB^T ⊙I[C]⇒D.
The implementation stores dense matrices row-major and sparse matrices in CSR format to enable coalesced accesses into inputs and outputs.
Hierarchical 1-Dimensional Tiling: One-dimensional tiling partitions the output into independent tiles processed by separate thread blocks, loading sparse strips and dense tiles to accumulate vector-matrix products.
Hierarchical 1-Dimensional Tiling: Varying dense-matrix column counts motivate one-dimensional tiling across RNN batches, Transformer tokens, convolutions, and EfficientNet spatial dimensions.
Hierarchical 1-Dimensional Tiling: The scheme supports specialized tile sizes, higher occupancy for small dimensions, and aggressive loop unrolling through fixed-size blocks.

B. Vector Memory Operations

Vector memory instructions can reduce bandwidth and instruction costs, but sparse rows create width and alignment problems; subwarp tiling and ROMA address these constraints.

Vector memory instructions mitigate bandwidth bottlenecks and reduce the instruction count needed to express computations.
Naive vector loads waste work when sparse rows or dense-matrix widths are smaller than the vector request, limiting their utility to very large problems.
Sparse-row starts lack alignment guarantees because arbitrary row lengths determine their addresses, even when dense and output rows are aligned.
Subwarp Tiling: Subwarp tiling assigns warp subsets to independent one-dimensional output tiles, reducing access width constraints and spreading threads across more rows.
Subwarp Tiling: Variable-length sparse rows can cause warp divergence, creating load imbalance within a warp.
Reverse Offset Memory Alignment: Reverse offset memory alignment backs each row’s address to the nearest vector-width boundary and masks values from the previous row during the first iteration.

C. Row Swizzle Load Balancing

Sparse GPU workloads suffer imbalance both across processing elements and within warps. Row swizzle remaps work at two levels to distribute computation more evenly.

Load imbalance arises between warps or thread blocks and within warps or thread blocks, causing idle SMs, divergence, and inefficient resource use.
Row swizzle remaps work assignments so each processing element receives approximately equal work, addressing both imbalance sources.
Row Binning: Row binning alters tile mappings so SMs receive roughly equal work based on thread-block scheduling.
Figure 5 contrasts subwarp tiling’s split warp accesses with reverse offset alignment’s backed-up row starts and first-iteration masking.
Row Bundling: Row bundling changes tile mappings so subwarps receive approximately equal work when warps span multiple sparse-matrix rows.

1) Volta Thread Block Scheduler:

The scheduler-aware row-binning heuristic orders sparse row bundles to balance work across warps and streaming multiprocessors, with sorting overhead amortized across repeated computations.

Rows are sorted by decreasing length and grouped into consecutive bundles of similarly sized rows.
The first wave of bundles is scheduled round-robin across SMs, while later bundles execute in decreasing order of heaviness as resources free.
The heuristic avoids requiring a target bundle size or knowledge of kernel-selection heuristics.
Sorting row indices can be amortized because DNN sparse-matrix topology is typically updated infrequently.
Figure 6 summarizes the row-swizzle scheme, while Figure 7 evaluates its performance as load imbalance increases.

D. Implementation Details

The implementation combines low-level memory and tiling optimizations with row swizzling to improve sparse-kernel efficiency across irregular workloads and precisions.

The kernels load sparse values and indices into shared memory, scaling indices before storage to avoid repeated index-processing work.
Residue handling accumulates products left after full nonzero tiles, which matters because sparse row lengths are rarely divisible by the tile size.
Zeroing shared-memory buffers enables 128-bit loads for residual values and indices despite residual counts not being divisible by four.
Mixed-precision kernels accept FP16 data and 16-bit indices, compute with FP32 fused multiply-adds, and convert outputs back to FP16.
Row swizzling reorders rows, groups similar lengths into bundles, and processes bundles by decreasing size to balance work across warps and SMs.

VI. SAMPLED DENSE–DENSE MATRIX MULTIPLICATION

The SDDMM kernel adapts one-dimensional tiling to sparse output strips, avoids uncoalesced accesses to the transposed operand, and uses subwarps for broader vectorization support.

SDDMM maps thread blocks to consecutive nonzero strips, distributing work directly over the sparse output.
The kernel assigns each thread a portion of all outputs in its tile, avoiding strided accesses to the row-major right-hand matrix.
An explicit shared-memory transpose would reduce register use but double shared-memory consumption and shrink L1 capacity, so it is not used.
Subwarp processing extends vector loads and stores across more problem configurations, while sparse-matrix accesses remain scalar.
Load balancing is less critical for SDDMM because every dot product has equal length; dynamic parallelism could help at very high sparsity.

A. Kernel Benchmarks

On DNN and recurrent-network workloads, the kernels consistently outperform cuSPARSE and often other sparse implementations, while ablations show that optimization benefits depend on problem configuration.

A. Kernel Benchmarks: 3.58× geometric-mean SpMM speedup and 2.19× SDDMM speedup over cuSPARSE were measured across the DNN sparse-matrix dataset.
A. Kernel Benchmarks: 27.3% of single-precision peak was reached for SpMM and 26.2% for SDDMM, with cuSPARSE outperformed on 99.75% and 93.34% of problems respectively.
A. Kernel Benchmarks: 5.97× mixed-precision SpMM geometric-mean speedup over cuSPARSE accompanied a 5.57TFLOPs peak throughput and 99.7% problem win rate.
2) Sparse Recurrent Neural Networks:: 1.56×, 1.59×, and 3.47× SpMM geometric-mean speedups were achieved over MergeSpmm, ASpT, and cuSPARSE, respectively.
B. Ablation Study: Row swizzling and residue unrolling remain robust across configurations, whereas vector memory instructions benefit compute-heavy problems more than small ones.
B. Ablation Study: SDDMM scalar memory operations can outperform vector operations on small DNN weight matrices, while vector kernels reached 2.45× speedup on RNN problems.
B. Ablation Study: The kernels also benefit from favorable data layouts and an efficient implementation enabled by one-dimensional tiling.

C. Application: Sparse Transformer

Sparse Transformer attention reduces quadratic attention computation by evaluating only selected query-key interactions, implemented through SDDMM followed by SpMM. The evaluated model uses a structured local-and-random attention mask for long sequences.

Application: Sparse Transformer: Transformer attention computes token similarities through QK^T, whose cost and memory grow quadratically with sequence length.Sparse attention computes only a subset of QK^T outputs before multiplying by V.
Application: Sparse Transformer: Unstructured sparse attention maps to an SDDMM followed by an SpMM.The sparse output of the attention similarity computation is multiplied by V.
Application: Sparse Transformer: The Transformer experiment uses sequence length 12,288, three layers, eight attention heads, hidden dimension 1,024, and filter size 4,096.Models were trained with batch size 8 for 140,000 steps, while sparsity was simulated during training and converted to sparse form for benchmarking.
Application: Sparse Transformer: The sparse attention mask combines a dense diagonal band with random off-diagonal sparsity sampled inversely with distance from the diagonal.The dense band has size 256, off-diagonal sparsity is 95%, and the mask is shared across attention heads and layers.

2) Results & Analysis:

The sparse MobileNet evaluation compares pruned 1×1 convolutions against dense baselines, while the broader results include kernel ablations and Transformer attention connectivity. Sparse models improve the accuracy-throughput tradeoff, with additional gains constrained by widened models and remaining dense operations.

Results & Analysis: Sparse MobileNetV1 models use magnitude pruning at 90% sparsity for 1×1 convolutions, while the first layer remains dense.The first layer was left dense because its activation matrix made it bandwidth bound and less responsive to weight sparsity.
Results & Analysis: Inference fuses batch normalization, bias, and ReLU around depthwise and 1×1 convolutions, with cuBLAS serving the dense 1×1 baselines.The sparse 1×1 convolutions use fused bias and ReLU kernels.
Results & Analysis: Pruning-induced accuracy loss requires sparse models to be wider than accuracy-matched dense models, increasing the cost of nonsparse operations.Depthwise convolutions become a significant bottleneck after the 1×1 convolutions are pruned.
Results & Analysis: The Transformer connectivity comparison contrasts dense all-to-all attention with sparse attention using a dense band and distance-weighted random off-diagonal sparsity.The upper diagonal is masked so tokens attend only to preceding tokens.

IX. DISCUSSION & CONCLUSION

The discussion identifies practical constraints and future opportunities for sparse deep-learning kernels. Despite bottlenecks and limited access to dedicated matrix-multiply hardware, the kernels support sparse Transformer and MobileNet models with substantial speed and memory benefits without accuracy loss.

Discussion & Conclusion: Training DNNs also requires transposed sparse computation, which is difficult to fuse into CSR SpMM but can use cached offsets and indices with an argsort-based transpose.Alternative sparse formats may support both transposed and non-transposed computation.
Discussion & Conclusion: Sparse MobileNetV1 models achieve 21-24% speedups at a given accuracy, or 1.1% higher accuracy at the same throughput.The corresponding tradeoff curves are reported across model sizes.
Discussion & Conclusion: Shared memory bandwidth limits kernel performance on large problems.Reusing right-input values across multiple rows is proposed as one way to alleviate this bottleneck.
Discussion & Conclusion: The kernels cannot use dedicated matrix-multiply hardware for unstructured sparsity, although unpacking sparse tiles could potentially enable those operations.Vector and block sparsity can exploit such hardware despite model-quality loss.
Discussion & Conclusion: 1.2-2.1× speedups and up to 12.8× memory savings are demonstrated for sparse Transformer and MobileNet models without sacrificing accuracy.The work attributes these applications to favorable properties found in deep-learning sparse matrices.