GenTen 中面向性能可移植的稠密 MTTKRP
数学软件
2025-10-17 v1
摘要
我们通过引入一种在现代 CPU 和 GPU 架构上具有加速性、稠密 matricized tensor times Khatri-Rao product (MTTKRP) 的方法来扩展 GenTen 张量分解软件包,这是标准 polyadic (CP) 张量分解的核心内核。与 Tensor Toolbox、TensorLy 等基于矩阵乘法的 MTTKRP 内核相比,这些方法显式地形成 Khatri-Rao 矩阵,我们发展了一种基于元素级并行化的矩阵自由方法,其内存开销随秩 R 增长,类似于张量形状之和 O(R(n+m+k)),而矩阵基方法的内存开销则随张量形状的乘积 O(R(mnk)) 增长。对于我们研究的最大规模问题,秩为 2000 的 MTTKRP,较小的增长率导致矩阵自由的内存开销仅为矩阵基方法的 2%,实现了 50 倍的改进。实际情况下,减少的内存影响意味着我们的矩阵自由 MTTKRP 可在单个 NVIDIA H100 上计算秩为 2000 的张量分解,而矩阵基 MTTKRP 需要 6 个 H100。我们还将优化的矩阵自由 MTTKRP与不同设备上的基准矩阵自由实现进行比较,显示在 Intel 8480+ CPU 上实现 3 倍单设备加速,在 H100 GPU 上实现 11 倍加速。除了数值结果外,我们提供了针对理想多级缓存机器的细粒度性能模型,将分析性能预测与实证结果进行比较,并提供了一个受启发的启发式方法,用于选择算法的超参数。
引用
@article{arxiv.2510.14891,
title = {A Performance Portable Matrix Free Dense MTTKRP in GenTen},
author = {Gabriel Kosmacher and Eric T. Phipps and Sivasankaran Rajamanickam},
journal= {arXiv preprint arXiv:2510.14891},
year = {2025}
}
备注
10 pages, 5 figures, 4 tables, for implementation see https://github.com/sandialabs/GenTen