Segmented Operations using Matrix Multiplications

TL;DR

This paper introduces MMV-RAM, a new computational model combining matrix and vector units to optimize AI operations, demonstrating theoretical speed-ups and practical implementations on AI accelerators for diverse workloads.

Contribution

The paper proposes MMV-RAM, a theoretical model extending Vector-RAM with matrix units, and develops algorithms that leverage this model for improved AI kernel performance.

Findings

Theoretical analysis shows speed-ups over vector-only approaches.

Algorithms effectively implement segmented scan and sum using MMUs.

Practical implementation on Ascend 910B accelerators demonstrates real-world benefits.

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