arXiv:2510.08874cs.DCcs.AI2025-10被引 1

一种通用分布式矩阵乘法算法,支持任意划分方式,无需重分布数据。

Slicing Is All You Need: Towards A Universal One-Sided Algorithm for Distributed Matrix Multiplication

  • 用切片计算需相乘的局部块,直接执行或优化后执行。
  • 在多种划分和复制因子下性能媲美PyTorch DTensor。
  • 适合需要灵活支持不同划分的科学计算与AI场景。

科学、数据分析和人工智能工作负载中的许多重要应用依赖于分布式矩阵乘法。现有算法针对不同问题规模和划分方式(如1D、2D、1.5D、2.5D)设计,但仅支持部分划分组合。缺乏对应算法时需重新分配数据,增加通信开销。本文提出一种通用的一侧式分布式矩阵乘法算法,可支持所有划分组合与复制因子。算法通过切片(索引运算)确定必须相乘的重叠块集合,该列表可直接执行,或重排并降低为优化IR以最大化重叠。我们使用基于C++的PGAS编程框架实现,利用节点内互连进行直接GPU-to-GPU通信。在多种划分和复制因子下评估性能,结果表明本方法在性能上可与高度优化的PyTorch DTensor相媲美。

原文摘要 · Abstract (English)

Many important applications across science, data analytics, and AI workloads depend on distributed matrix multiplication. Prior work has developed a large array of algorithms suitable for different problem sizes and partitionings including 1D, 2D, 1.5D, and 2.5D algorithms. A limitation of current work is that existing algorithms are limited to a subset of partitionings. Multiple algorithm implementations are required to support the full space of possible partitionings. If no algorithm implementation is available for a particular set of partitionings, one or more operands must be redistributed, increasing communication costs. This paper presents a universal one-sided algorithm for distributed matrix multiplication that supports all combinations of partitionings and replication factors. Our algorithm uses slicing (index arithmetic) to compute the sets of overlapping tiles that must be multiplied together. This list of local matrix multiplies can then either be executed directly, or reordered and lowered to an optimized IR to maximize overlap. We implement our algorithm using a high-level C++-based PGAS programming framework that performs direct GPU-to-GPU communication using intra-node interconnects. We evaluate performance for a wide variety of partitionings and replication factors, finding that our work is competitive with PyTorch DTensor, a highly optimized distributed tensor library targeting AI models.

矩阵乘法分布式计算通用算法高性能

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。