Title: Segmented Operations using Matrix Multiplications
Primary area: infrastructure, software libraries, hardware, systems, etc.
Keywords: AI accelerators, Algorithm design, Parallel algorithms, Matrix multiplication, Segmented operations, Prefix sum, Models of computation, MMV-RAM

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]. Given 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 both matrix and vector cores. We evaluate these implementations on synthetic and real-world datasets from various applications, including Large Language Models.

TL;DR: We introduce the MMV-RAM model to enable algorithmic design and analysis on modern AI accelerators