Title: Segmented Operations using Matrix Multiplications Show Chinese title

Title: 使用矩阵乘法的分段操作

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. Show Chinese abstract

Abstract: 专门执行小矩阵乘法作为原始操作的计算单元通常存在于现代AI加速器中。 然而，这些矩阵乘法单元（MMUs）在许多基本的深度学习操作中除了密集矩阵乘法之外常常被低估使用。 巧合的是，这种架构缺乏严格的理论计算模型阻碍了算法设计。 在本工作中，我们提出了MMV-RAM，这是一种计算模型，它明智地扩展了向量-RAM模型，增加了一个MMU。 我们提供了详细的理论分析，并在矩阵和向量单元之间仔细平衡计算能力，这是由电路复杂性下界决定的，即奇偶性不在AC{[0]}中。 在MMV-RAM的基础上，我们继续进行算法设计，从两个基本的并行操作开始：分段扫描和求和。 通过将它们表示为基本并行原语的组合（例如，分段求和可以简化为：扫描、压缩和向量微分），我们可以利用MMU执行推测性分块计算，最终导致相对于仅向量方法的可证明理论加速。 这些结果也适用于其他常见的AI内核，包括密集矩阵乘积和稀疏矩阵-向量乘积。 作为一个案例研究，我们在Ascend 910B AI加速器上实现了所提出的算法，该加速器包含矩阵和向量核心。 我们在来自各种应用的合成和真实世界数据集上评估了这些实现，包括大型语言模型。