Exploring Sparse Matrix Multiplication Kernels on the Cerebras CS-3

Abstract: In recent years, novel AI accelerators have emerged as promising alternatives to GPU for AI model training and inference tasks. One such accelerator, the Cerebras CS-3, achieves strong performance on large model training as well as scientific applications like molecular dynamics simulations. While dense compute workloads have been thoroughly explored for the CS-3, its potential for sparse workloads has not been fully examined. Applications requiring sparse linear algebra kernels, such as GNNs, linear solvers, and recommendation systems, could achieve good performance on a dataflow accelerator like the CS-3. In this work, 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 kernel designs for these operations and optimize our designs to improve I/O performance, memory footprint, and scalability to large matrices. Our evaluation examines memory footprint and SpMM/SDDMM speedup relative to CPU. The evaluation suggests that the CS-3 can outperform CPU by 100$\times$ for SpMM with 90\% sparse matrices with performance improving as sparse matrix dimensionality increases. SDDMM on CS-3 can outperform CPU 20$\times$ for 90\% sparse matrices. We additionally find that as sparsity increases to beyond 99\%, the CS-3 suffers from performance degradation that makes it slower than CPU for SpMM.

• The paper introduces low-level, device-optimized implementations for SpMM and SDDMM on the Cerebras CS-3, achieving speedups of up to 100x and 20x respectively.
• The study employs a SELLPACK-like storage format to improve I/O bandwidth and parallel routing, balancing trade-offs between communication overhead and memory efficiency.
• The empirical evaluation highlights performance degradation in hypersparse regimes and underscores the need for further algorithmic and architectural optimizations.

Sparse Matrix Multiplication Kernels on Cerebras CS-3: An Analytical Essay

Introduction and Context

The emergence of AI accelerators, particularly wafer-scale engines like the Cerebras CS-3, has shifted paradigms in deep learning and scientific computing, primarily through substantial advances in compute density and on-chip memory capacity. While extensive evaluation of such architectures for dense operations is available, their efficacy and constraints in the context of sparse linear algebraic kernels—a foundation of applications such as GNNs, pruned DNNs, and scientific simulations—have not been similarly scrutinized. The paper "Exploring Sparse Matrix Multiplication Kernels on the Cerebras CS-3" (2604.27985) undertakes a comprehensive investigation into the two pivotal kernels: sparse-dense matrix multiplication (SpMM) and sampled dense-dense matrix multiplication (SDDMM), proposing low-level device-optimized implementations and systematically evaluating performance and scalability properties.

Figure 1: The Cerebras Wafer-Scale Engine 3 (WSE-3, or CS-3) architectural diagram, highlighting wafer boundaries, I/O channels, and hierarchical PE structure.

Challenges with Sparse Computation on Wafer-Scale Architectures

Sparse computation on massively parallel, spatial architectures like CS-3 faces divergent challenges compared to conventional GPU/CPU architectures:

• Memory Footprint and Data Movement: Standard dense storage exhausts on-chip memory at moderate matrix sizes due to quadratic scaling, especially problematic for real-world GNN graphs. Sparse formats (CSR, SELLPACK variants) are essential but introduce complexity in data flow partitioning and pipelining.
• Communication Bottlenecks: The efficiency of sparse kernels is tightly bounded by the host-to-device and intra-wafer data routing performance, necessitating careful alignment of storage formats to architectural features for maximum I/O utilization.
• PE Allocation and Synchronization: Fine-grained mapping and synchronization of router, worker, and accumulator tasks on the CS-3 mesh is non-trivial, with design trade-offs between parallelism and local memory constraints.

  Figure 2: Timing and scaling behavior of GNN, trained using dense-dense matmul over synthetic sparse matrices, underscores quadratic explosion of memory requirements and rapid device limits for increasing N.

Kernel Design and Implementation

SpMM Kernel Evolution

The initial SpMM design leverages a hierarchical routing strategy: the sparse matrix $A$ (in CSR), decomposed into index-value pairs, is streamed via a single I/O channel then distributed to rows of worker PEs, each maintaining a slab of the dense matrix $H$. Despite exploiting the PE mesh for pipelined partial sum accumulation and efficient multiplication, the design is bottlenecked by serialization at the top-level router, constraining data ingress bandwidth.

Figure 3: Initial SpMM kernel schematic in CSL, illustrating router, worker, and accumulator PE roles and data propagation through the mesh.

To alleviate this bottleneck, the authors propose a SELLPACK-like matrix storage format, enabling immediate distribution of nonzero index-value pairs to the appropriate router, thus facilitating concurrent utilization of multiple wafer I/O channels and mitigating north-south data movement pressure.

Figure 4: Comparison of matrix storage formats—dense, CSR, and SELLPACK-like—indicating explicit marking of end-of-row and null entries for efficient streaming and padding.

Figure 5: Optimized SpMM with SELLPACK-like format, supporting parallel routing and processing and reducing synchronization.

To further enhance device utilization, multiple accumulator rows are introduced, enabling buffering of output $Y$ on-chip and amortizing device-to-host communication over larger sub-blocks.

Figure 6: Schematic for buffered multi-accumulator output on-chip, reducing serialization and increasing communication-compute overlap.

SDDMM Kernel

The SDDMM kernel is mapped similarly, with routers distributing dense $B$ and $C$ tiles along orthogonal network axes, feeding workers responsible for nonzero tile subsets of $A$.

Figure 7: SDDMM task mapping on CS-3, where nonzeros in $A$ trigger corresponding sampled $B C$ multiplies across the PE grid.

Empirical Evaluation

Storage Format Efficiency

The analysis of the SELLPACK-like format reveals a density-dependent trade-off: at extreme sparsity (density $\leq 10^{-4}$), padding and end-of-row markers inflate the transmitted element count—up to 80$\times$ relative to pure nonzero count in CSR. However, for sub-99% sparsity, overhead moderates to $H$01.5$H$1, a necessary compromise for high I/O bandwidth parallelization.

Figure 8: Ratio of elements sent via SELLPACK-like format vs. true nonzero count across density and matrix size, highlighting increasing overhead with greater sparsity.

Performance: SpMM

CS-3 yields speedups of up to 100$H$2 over a high-end CPU when operating on moderate-to-high density matrices—specifically, at 90% sparsity. The substantial speedup is a direct consequence of improved compute bandwidth utilization enabled by parallelized routing and tiling, as well as effective PE occupation.

However, performance rapidly deteriorates in the hypersparse regime ($H$399% sparsity), where communication overheads and underutilized compute resources dominate, causing CS-3 throughput to fall below CPU baselines. Notably, scaling benefits are consistent with larger N, and parameter choices for PE allocation (myc, mcpp) show robustness in maintaining speedup.

Figure 9: Comparative SpMM results showing CS-3 vs. CPU speedup as a function of density and matrix size across kernel parameterizations.

Performance: SDDMM

For SDDMM, CS-3 delivers up to 20$H$4 speedup on 90% sparse matrices compared to CPU. Unlike SpMM, SDDMM's performance scaling with density is less steep, as data movement for output $H$5 buffers increases for higher max_nonzeros (mnz). Reducing mnz can yield further improvements by minimizing device-host transfers.

Figure 10: SDDMM speedup (CS-3 vs. CPU) as a function of density, N, and mnz; higher density and lower mnz are associated with greater speedup.

Discussion and Implications

The study establishes several key points:

• Wafer-scale engines are highly advantageous for workloads with moderate sparsity and massive matrix scale, particularly in regimes where dense storage is infeasible and the number of nonzeros remains sufficiently high.
• For hypersparse matrices, communication costs and poor compute resource utilization become the primary bottlenecks, suggesting the need for further algorithmic and architectural adaptation.
• The parallelized SELLPACK-inspired storage format is validated as a pragmatic trade-off between communication efficiency and memory overhead, especially in the context of spatial architectures where maximization of I/O bandwidth is paramount.
• The design space for PE allocation (myc, mcpp, accumulator rows) reveals significant flexibility, making it feasible to co-locate multiple sparse kernels (e.g., GCN layers) while maintaining high performance.
• The study confirms prior distributed-memory and GPU-centric observations—namely, load balancing, communication minimization, and storage format co-design are essential in modern sparse kernel acceleration [spmm-distr1, spmm-gpu1].

Theoretical Ramifications

This work advances the understanding of spatial dataflow hardware for sparse workloads: it empirically quantifies the interplay between architectural constraints (I/O channels, on-chip memory, routing granularity), storage format inflation, and achievable parallelism. The results also imply that, as hardware evolves for extreme scale, the canonical wisdom for distributed sparse computation—favoring higher-dimensional decompositions and communication-avoiding schemes—translates to on-chip networks in wafer-scale environments.

Practical Outlook and Future Directions

The findings have direct implications for practitioners deploying large-scale GNNs and other sparse-structured models on AI accelerators. For optimal utilization of wafer-scale devices, workloads should be structured to maximize nonzero activity per task and consider storage format inflations during pipeline design.

The authors highlight future opportunities including:

• Adapting the kernel to further exploit hypersparse matrix patterns through more sophisticated dataflow scheduling, load balancing, and format adaptation.
• Investigating alternative sparse storage formats with lower padding and marker overhead while still supporting high I/O parallelism.
• Tighter fusion of GEMM, SDDMM, and SpMM on-chip to amortize data movement, potentially improving both compute utilization and memory efficiency for complex model architectures (e.g., GAT layers).

Conclusion

"Exploring Sparse Matrix Multiplication Kernels on the Cerebras CS-3" (2604.27985) demonstrates that wafer-scale spatial accelerators can provide extraordinary speedups for sparse linear algebra—up to 100$H$6 for SpMM and 20$H$7 for SDDMM—when nonzero density and matrix scale are sufficient. As density drops into the hypersparse regime, communication and memory inefficiencies become limiting factors. The work thus sets a rigorous foundation for both the further hardware-software co-design of sparse workloads on overlapping compute-communication architectures, and for new research directions in storage format engineering and algorithmic adaptation for future, more sparse-aware accelerators.