摘要
arXiv:2505.06302v1 宣告类型: cross
摘要: 计算密集型张量操作占大型语言模型(LLMs)和深度神经网络中计算的大约90%以上。自动并有效地利用硬件原语生成高性能张量操作对于不断发展的硬件架构(如RISC-V、ARM和GPU)至关重要,因为手动优化实现至少需要几个月时间,并且缺乏可移植性。LLMs 在生成高层次语言代码方面表现出色,但在充分理解硬件特性并生成高性能张量操作方面存在困难。我们引入了一个带有一行用户提示的张量操作自动生成框架(QiMeng-TensorOp),该框架使LLMs能够自动利用硬件特性,利用硬件原语生成张量操作,并在多种硬件上进行参数调优以实现最佳性能。在各种硬件平台、最先进的LLMs和典型张量操作上的实验结果表明,QiMeng-TensorOp 有效释放了各种硬件平台的计算能力,并自动生成了高性能的张量操作。与vanilla LLMs 相比,QiMeng-TensorOp 的性能提高了多达1291倍。即使与人类专家相比,QiMeng-TensorOp 也能在RISC-V CPU 上达到OpenBLAS的251%,在NVIDIA GPU 上达到cuBLAS的124%。此外,与人类专家相比,QiMeng-TensorOp 还将开发成本减少了200倍。