超越FLOPs：COSMA如何基于通信下界构建并行矩阵乘法
内容提要
COSMA基于通信下界设计并行矩阵乘法。针对计算量增大后数据搬运成为瓶颈、有限内存限制数据复用的问题，利用三维计算空间与红蓝卵石模型推导近最优分块策略：优先保留方形输出块、流式读取输入，并权衡二维与三维划分。实验表明，相比ScaLAPACK等最高加速12.8倍，达到平台峰值性能的88%。
延伸解读
通信下界为何是平方根而非线性
文章指出，若快速内存能容纳三个边长为t的方形块，则加载Θ(t²)个数据可支持Θ(t³)次乘加，每个数据支持Θ(t)次乘加，因此复用效率与√S成正比，而非S。这解释了I/O下界中指数1/2的来源。但文章强调，方形块例子仅说明指数，严格证明需处理不规则块、部分和、重计算及跨阶段复用，COSMA通过红蓝卵石模型和划分使分析更精确。
保留输出块并流式读取输入的策略
文章提出另一种顺序调度：在快速内存中保留a×b的输出块，流式读取A的列段和B的行段。对于固定输出面积ab=R，由算术几何平均不等式，当a=b=√R时复用效率最高。因此大部分内存应存放接近方形的输出块，辅以较小的输入缓冲区。每个输出块在遍历完整个k维度后才写回，避免反复写入部分和造成额外通信。
二维与三维划分的权衡
当处理器数量增加时，仅划分i、j维度会使输出块变小、输入复用下降；同时划分ℓ维度可保留较大的局部输出块，但需跨处理器归约部分和。文章指出，选择二维还是三维划分应依据问题形状、处理器数量和可用内存。此外，处理器数量并非越多越好，例如65个处理器难以排成规则网格，而64个可组成4×4×4网格，可能通信更少。
实验结论的适用边界
文章报告COSMA在Piz Daint的CPU分区上相比ScaLAPACK、CARMA和CTF最高加速12.8倍，平均2.2倍，达到平台峰值性能的88%。但这些结果是针对2019年论文的特定平台、基线和实验，并非对现代GPU或任意矩阵形状的保证。文章强调，其价值在于从下界、分解到实现的完整链条，而非具体加速数字的普适性。
Q&A
COSMA 是什么？它主要解决什么问题？
COSMA 是一种基于通信下界设计的并行矩阵乘法算法。它针对计算量增大后数据搬运成为瓶颈、有限内存限制数据复用的问题，利用三维计算空间与红蓝卵石模型推导近最优分块策略，以最小化通信量。
为什么并行矩阵乘法中通信会成为瓶颈？
因为每个处理器必须接收输入数据，部分结果可能需要在处理器间合并。随着算术速度提升，数据移动（通信）可能成为限制性能的主要因素。
COSMA 如何利用三维计算空间来优化矩阵乘法？
COSMA 将矩阵乘法视为一个 m×n×k 的三维计算空间，每个标量乘加对应一个三元组 (i,j,ℓ)。处理器接收一个 a×b×c 的子盒，所需数据是子盒在三个坐标平面上的投影：A 为 a×c，B 为 c×b，C 的部分和为 a×b。局部计算量为 abc，边界数据量为 ac+bc+ab。通过最大化计算量与边界之比来实现数据复用。
COSMA 中提到的“方形输出块”策略是什么？为什么优先保留方形输出块？
该策略是在快速内存中保留一个 a×b 的输出块，同时流式读取 A 的列段和 B 的行段。对于固定输出面积 ab=R，根据算术-几何平均不等式，当 a=b=√R 时复用效率最高。因此，大部分内存应存储尽可能接近方形的输出块，并辅以较小的输入缓冲区。
在并行矩阵乘法中，何时应该拆分第三个维度（ℓ 维度）？
当处理器数量增加导致输出块缩小、输入复用下降时，可以考虑拆分 ℓ 维度。这样多个处理器可以计算同一输出块的不同贡献，保持较大的局部输出块，但需要通信来归约部分和。是否拆分取决于问题形状、处理器数量和可用内存。
COSMA 的实验结果如何？相比其他库有哪些优势？
在 Piz Daint 的 CPU 分区上，COSMA 与 ScaLAPACK、CARMA 和 CTF 相比，最高加速 12.8 倍，平均加速 2.2 倍，性能达到平台峰值计算速率的 88%。这些结果是针对 2019 年论文的特定平台、基线和实验，并非对所有现代 GPU 或矩阵形状的保证。
COSMA 对推理性能研究有什么启示？
COSMA 建议了一种估计性能极限的方法：描述允许的标量计算和复用，在内存约束下推导最小通信量，然后与设备带宽关联。但该论证仍假设经典矩阵乘法。如果融合多个 Transformer 矩阵乘法，一个算子的输出可能保留在快速内存中供下一个使用，因此不能机械地添加孤立算子的强制写回成本，下界必须描述允许的融合范围和驻留状态。
