SparseFlow: A Sparse Inference Acceleration Library for Spiking Neural Networks, Achieving Up to 90x Performance Improvement

SparseFlow is a high-performance sparse inference acceleration library specifically designed for Spiking Neural Networks (SNNs). By leveraging the inherent high sparsity of Leaky Integrate-and-Fire (LIF) neuron outputs, it achieves up to 90x inference acceleration. Its core innovations include a two-stage sparse computing architecture and dynamic blocking strategy, providing an efficient engineering solution for brain-inspired computing and neuromorphic computing, and helping SNNs move from the lab to practical applications.

As the third generation of neural networks, Spiking Neural Networks (SNNs) are based on the pulse signal transmission mechanism of biological nervous systems. Their event-driven computing mode theoretically has extremely high energy efficiency, making them suitable for edge computing and neuromorphic chip deployment. However, in practical engineering, traditional dense convolution operators (such as cuDNN) perform full computations on all-zero pulse data blocks, leading to computational resource waste. How to use sparsity to accelerate inference is a key bottleneck for the practical application of SNNs.

SparseFlow adopts a two-stage sparse computing architecture: the first stage uses a lightweight pre-scan to identify non-zero block indices, and the second stage only performs convolution on non-zero blocks. The dynamic blocking strategy adaptively selects block sizes (16x16, 8x8, 4x4) based on the height of the feature map, ensuring optimal performance across different network layers and input sizes. Users can replace their existing SNN networks with the sparse accelerated version using just one line of code.

1. Automated Operator Replacement Framework: Builds a computation graph via torch.fx symbolic tracing, uses BFS traversal to identify convolution layers that need optimization, supports transparent layer penetration (skipping Dropout, Pooling, etc.), and falls back to forward hook linear search if symbolic tracing fails; 2. Triton GPU Kernel: Writes high-performance GPU kernels (pre-scan, 3x3/1x1 sparse convolution, etc.) based on Triton, uses scatter mode and atomic addition to accumulate results, and provides an nn.Module wrapper layer for seamless integration with PyTorch.

SparseFlow was tested on ResNet architectures (ResNet34/50/101/152), and the acceleration effect increases with network depth and sparsity: the layer1.0.conv2 layer (98.5% sparsity) achieves 13.1x acceleration; the layer2.1.conv2 layer (100% sparsity) reaches 72.2x; some layers have a maximum acceleration ratio of 90x. Deep networks benefit more from sparse acceleration because their pulses are sparser.

• Minimal Integration: No need to modify model definitions or training code; adding two lines of code after model creation completes the acceleration; - Intelligent Device Selection: The benchmark script automatically selects the idle GPU with the largest memory, and also supports manual specification of GPU IDs, ensuring code correctness and portability.

SparseFlow solves the key problem of SNNs moving from theory to practice and provides a reference for the software stack of neuromorphic computing chips. It represents the direction of deep learning hardware-software co-optimization and is of great significance for edge AI and low-power computing (mobile, IoT, autonomous driving). Conclusion: SparseFlow transforms the sparsity advantage of SNNs into actual performance improvements, promoting SNNs to play a greater role in edge AI and other fields.