Spiking Transformer Acceleration on RISC-V Streaming Processors

Transformers built on conventional Artificial Neural Networks have represented the key innovation leading to the powerful AI models we have access to today. Spiking Neural Networks offer a more energy-efficient option than standard neural networks because the neurons inside such networks exchange spikes, which are binary signals. Recent work has shown that spiking Transformers obtain strong results thanks to sparse binary computation and generally present better energy efficiency. A consequence of operating with sparse binary spike tensors is, for instance, the elimination of multiplication, that can be replaced by addition and mask operations. However, this sparsity is unstructured and known only at run time. Dedicated neuromorphic chips exploit it but are tied to specific models, and despite it, CPUs and GPUs execute spike tensors as dense operations. In this work, we profile the attention module of the Spiking Transformer with Spatial-Temporal Attention on CIFAR-100 inputs and find its query, key and value spike tensors are about 96.87% sparse. These findings allow us to design accelerated kernels for the Snitch cluster, an eight-core RISC-V cluster with hardware streams for indirect accesses and vectorization support. They cover every operation of the attention block, and we evaluate them in cycle-accurate simulation, both individually and chained into the complete block. Over a plain-C reference, the Sparse Matrix-Vector Addition outperforms the baseline for up to 34.02× at fp16x4. At fp16x4, the sparse engine is faster than a dense matrix multiplication of the same lane width up to a density of 0.75, and slower above it. A programmable streaming cluster can thus turn the sparsity of a spiking Transformer into measured speedups without fixing the model in hardware, so it can follow the model as it evolves.