Shfl-BW: Accelerating Deep Neural Network Inference with Tensor-Core Aware Weight Pruning

TL;DR

This paper introduces Shfl-BW, a novel weight pruning pattern that enables efficient tensor-core utilization in DNN inference, achieving significant speedups while maintaining accuracy.

Contribution

The paper proposes Shfl-BW, a flexible sparse pattern that balances accuracy and efficiency, optimized GPU kernels for tensor-core acceleration in DNNs.

Findings

Achieves up to 4.18x speedup on GPU layers with 75% sparsity.

Maintains high accuracy with minimal loss while accelerating Transformer layers.

Demonstrates state-of-the-art speed-accuracy trade-offs in GPU DNN inference.

Abstract

Weight pruning in deep neural networks (DNNs) can reduce storage and computation cost, but struggles to bring practical speedup to the model inference time. Tensor-cores can significantly boost the throughput of GPUs on dense computation, but exploiting tensor-cores for sparse DNNs is very challenging. Compared to existing CUDA-cores, tensor-cores require higher data reuse and matrix-shaped instruction granularity, both difficult to yield from sparse DNN kernels. Existing pruning approaches fail to balance the demands of accuracy and efficiency: random sparsity preserves the model quality well but prohibits tensor-core acceleration, while highly-structured block-wise sparsity can exploit tensor-cores but suffers from severe accuracy loss. In this work, we propose a novel sparse pattern, Shuffled Block-wise sparsity (Shfl-BW), designed to efficiently utilize tensor-cores while minimizing the constraints on the weight structure. Our insight is that row- and column-wise permutation provides abundant flexibility for the weight structure, while introduces negligible overheads using our GPU kernel designs. We optimize the GPU kernels for Shfl-BW in linear and convolution layers. Evaluations show that our techniques can achieve the state-of-the-art speed-accuracy trade-offs on GPUs. For example, with small accuracy loss, we can accelerate the computation-intensive layers of Transformer by 1.81, 4.18 and 1.90 times on NVIDIA V100, T4 and A100 GPUs respectively at 75% sparsity.