用空间填充曲线简化矩阵乘法,无需调参即可高效运行。
Space Filling Curves is All You Need: Communication-Avoiding Matrix Multiplication Made Simple
- 用空间填充曲线划分矩阵计算任务,自动获得良好数据局部性。
- 在多个CPU平台实现最高5.5倍性能提升,平均提速1.8倍。
- 适合追求免调参、高效率的高性能计算与深度学习开发者。
通用矩阵乘法(GEMM)是高性能计算和深度学习的核心。现有厂商库通过调整张量布局、并行策略和缓存分块来减少内存层级间的数据移动,但最优参数依赖于硬件平台和矩阵形状,导致调优困难。本文重新利用空间填充曲线(SFC),通过改进的SFC分割策略,构建了平台无关、形状无关的矩阵乘法方案,显著提升数据局部性。进一步将SFC方法扩展为通信避免(CA)算法,可严格最小化数据移动。该方案代码简洁,集成无缝,在多个CPU平台上达到业界领先性能,对多种GEMM形状的加速比最高达5.5倍(加权调和均值提速1.8倍)。我们在两个真实场景中验证其价值:一是大语言模型推理预填充阶段,速度提升最高达1.85倍;二是分布式内存矩阵乘法,加速比最高达2.2倍。
原文摘要 · Abstract (English)
General Matrix Multiplication (GEMM) is the cornerstone of HPC workloads and Deep Learning. State-of-the-art vendor libraries tune tensor layouts, parallelization schemes, and cache blocking to minimize data movement across the memory hierarchy and maximize throughput. Optimal settings for these parameters depend on the target platform and matrix shapes, making exhaustive tuning infeasible. We revisit Space Filling Curves (SFC) to alleviate this cumbersome tuning. We partition the Matrix Multiplication using advancements in SFC, and obtain platform-oblivious and shape-oblivious Matrix Multiplication schemes with high degree of data locality. We extend the SFC-based work partitioning to implement Communication-Avoiding (CA) algorithms that provably minimize data movement. The integration of CA-algorithms is seamless with compact code, achieving state-of-the-art results on multiple CPU platforms, outperforming vendor libraries up to 5.5x for a range of GEMM-shapes (1.8x Weighted Harmonic Mean speedup). We show the impact of our work on two real-world applications by leveraging our GEMM as compute backend: i) prefill of LLM inference with speedups up to 1.85x over State-Of-The-Art, and ii) distributed-memory Matrix Multiplication with speedups up to 2.2x.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。