Slicing Is All You Need: Towards a Universal One-Sided Distributed MatMul (arxiv.org)

Researchers introduced a universal one-sided algorithm for distributed matrix multiplication that handles all combinations of partitionings (1D, 2D, 1.5D, 2.5D) and replication factors using a single method. Instead of maintaining separate implementations or paying the cost of operand redistribution when an uncommon partitioning appears, their approach uses slicing (index arithmetic) to enumerate the overlapping tiles that must be multiplied locally. Those local multiplies can be executed directly or reordered and lowered to an optimized IR to maximize computation/communication overlap. The team implemented the method in a high-level C++ PGAS framework with direct GPU-to-GPU communication over intra-node interconnects and evaluated it across many partitionings and replication factors. Results show competitive performance with PyTorch DTensor, suggesting the approach can match specialized libraries while offering far greater flexibility. For the AI/ML community, this reduces engineering complexity and potential communication overhead in distributed training and inference, enabling a single, portable primitive that adapts to diverse sharding/replication setups without costly data movement.