GPU 上的快速 Kronecker 矩阵乘法
分布式、并行与集群计算
2024-02-29 v3
摘要
Kronecker 矩阵乘法是若干较小矩阵的 Kronecker 乘积与一个矩阵的乘法。Kron-Matmul 是许多科学和机器学习计算的核心操作。SOTA 的 Kron-Matmul 实现利用了现有的张量代数操作，如矩阵乘法、转置和张量矩阵乘法。然而，这种设计选择阻碍了若干 Kron-Matmul 特定的优化，从而留下了显著的性能提升空间。为了解决这个问题，我们提出了 FastKron，一种在单 GPU 和多 GPU 上高效的 Kron-Matmul 技术。FastKron 独立于线性代数操作，从而为 Kron-Matmul 实现了若干新的优化。因此，它在 1 个和 16 个 GPU 上分别比现有实现快 40.7 倍和 7.85 倍。
关键词
引用
@article{arxiv.2401.10187,
title = {Fast Kronecker Matrix-Matrix Multiplication on GPUs},
author = {Abhinav Jangda and Mohit Yadav},
journal= {arXiv preprint arXiv:2401.10187},
year = {2024}
}
备注
Accepted at PPoPP 2024