Segmented Operations using Matrix Multiplications

Aleksandros Sobczyk, Giuseppe Sorrentino, Anastasios Zouzias

TL;DR: We propose the MMV-RAM model of computation for algorithm analysis on modern AI accelerators.

Abstract: Specialized computational units that perform small matrix multiplications as primitive operations are typically present in modern AI accelerators. However, these Matrix Multiplication Units (MMUs) are often underutilized for many fundamental deep learning operations besides dense matrix multiplications. Coincidentally, the lack of a rigorous theoretical model of computation for such architectures obstructs algorithmic design. In this work, we propose MMV-RAM, a computational model which judiciously extends the Vector-RAM model with an additional MMU. We provide a detailed theoretical analysis and carefully balance the computational power between the matrix and vector units, guided by the circuit complexity lower bound that parity is not in AC[0]. In MMV-RAM, we proceed to algorithm design, starting with two fundamental parallel operations: *segmented scan* and *sum*. By expressing them as compositions of elementary parallel primitives (e.g., seg. sum reduces to: scan, compress, and vector differentiation), we can exploit MMUs to perform *speculative* blocked computations, ultimately leading to *provable theoretical speed-ups* against vector-only approaches. These results extend to other ubiquitous AI kernels, including dense matrix product, and sparse matrix-vector product. As a case study, we implemented the proposed algorithms on the Ascend 910B AI accelerator, which contains matrix and vector cores. We evaluate these implementations on synthetic and real-world datasets from various applications, including Large Language Models.

Primary Area: Theory->Optimization

Keywords: Parallel algorithms, AI accelerators, Models of computation, Matrix multiplication, Segmented operations, MMV-RAM

Submission Number: 14669