Exploring Sparse Matrix Multiplication Kernels on the Cerebras CS-3
In recent years, novel AI accelerators have emerged as promising alternatives to GPUs for AI model training and inference. One such accelerator, the Cerebras CS-3, has demonstrated strong performance on machine learning as well as scientific applications, such as molecular dynamics and seismic simulations. While the benefits of Cerebras systems for dense workloads have been well demonstrated, their potential for sparse workloads is not yet understood, particularly for large matrices that cannot fit on the device when represented in dense format. Yet, many applications, such as linear solvers and graph neural networks, rely on large sparse matrices. In this work, we make a step toward a better understanding of the use of CS-3 platforms for sparse operations. To this end, we explore two key sparse linear algebra kernels, sparse-dense matrix multiplication (SpMM) and sampled dense-dense matrix multiplication (SDDMM), on the Cerebras CS-3. We propose low-level CS-3 designs for these operations and optimize them to improve I/O performance, memory footprint, and scalability to large matrices. We evaluate speedup relative to the CPU. The results show that our CS-3 kernels can outperform the CPU by up to 100$\times$ for SpMM on 90\% sparse matrices, with performance improving as matrix dimensionality increases. SDDMM on the CS-3 can outperform the CPU by up to 20$\times$ on 90\% sparse matrices. However, as sparsity increases beyond 99\%, our kernels suffer performance degradation, approaching the performance of state-of-the-art CPU libraries or even underperforming them.