快速推断：Kronecker稀疏矩阵
机器学习
2025-06-16 v3
摘要
Kronecker稀疏(KS)矩阵——其支持由身份矩阵和全1矩阵块的Kronecker积构成——构成了Butterfly和Monarch矩阵的结构，具有更高效模型的潜力。然而，现有GPU核函数进行KS矩阵乘法时存在高数据移动成本，其中最高可达50%的时间花在内存绑定的张量置换上。我们提出一种融合的、面向输出的GPU核函数，消除这些开销，使全局内存流量减少三倍。在600种KS模式中，本核函数在FP32下实现中位数加速率为x1.4，能源消耗降低15%。基于KS模式参数的简单启发式方法可预测何时本方法优于现有方法。我们在github.com/PascalCarrivain/ksmm上发布了所有代码，包括兼容PyTorch的KSLinear层，演示在FP32下ViT-S/16端到端延迟降低最高达22%，GPT-2 medium降低16%。
引用
@article{arxiv.2405.15013,
title = {Fast Inference with Kronecker-Sparse Matrices},
author = {Antoine Gonon and Léon Zheng and Pascal Carrivain and Quoc-Tung Le},
journal= {arXiv preprint arXiv:2405.15013},
year = {2025}
}