一种注意力机制计算方法、计算系统及存储介质技术方案

技术编号：42036579 阅读：24 留言：0更新日期：2024-07-16 23:22

本发明专利技术公开了一种注意力机制计算方法、计算系统及存储介质，该计算方法包括前向传播阶段以及反向传播阶段；在前向传播阶段，将QKV映射、更新KV缓存以及旋转位置编码的算子进行融合，并调整融合后算子的内部计算流程，以减小访存的时间开销；在反向传播阶段，调整计算顺序以及访存顺序，以提高反向传播过程的效率。该方法通过优化注意力计算的前向传播和反向传播过程，分别采用算子融合和计算过程重排的策略，减少了计算中的访存次数和访存量，通过加速大语言模型中的注意力计算过程，提升了大语言模型的训练和推理性能。

全部详细技术资料下载

【技术实现步骤摘要】

本专利技术涉及人工智能领域，尤其涉及一种注意力机制计算方法、计算系统及存储介质。

技术介绍

1、大语言模型(llm)的影响力与日俱增，它有非常好的处理文本的能力，通过训练或微调模型的参数能够适用于各种场景，包括机器翻译、文本生成、智能助手、聊天机器人等。然而一个大语言模型的参数量远大于传统的神经网络，通常达到了数十亿级别，导致其训练和推理的耗时非常长，如何加速大语言模型的计算过程是关键。现有的大语言模型如llama2，opt，gpt的计算过程主要分为预填充(prefill)和解码(decode)两个阶段，在prefill阶段或是输入长文本的decode阶段时，注意力计算的占比都非常高，往往会达到整体计算过程的1/3甚至更多。就以llama2-7b为例，整个网络结构由32个transformerblock+归一化层+线性层组成，transformerblock(变换器块)则主要包括注意力计算和ffn(feed-forward neural network，前馈神经网络)计算。

2、现有attention(注意力机制)计算加速优...

【技术保护点】

1.一种注意力机制计算方法，其特征在于，包括前向传播阶段以及反向传播阶段；

2.根据权利要求1所述的一种注意力机制计算方法，其特征在于：融合后算子的内部计算流程包括三个阶段：

3.根据权利要求2所述的一种注意力机制计算方法，其特征在于：在第一阶段和第二阶段中，输入矩阵X暂存在片上；在第二阶段中，矩阵K和矩阵V暂存在片上；在第三阶段中，矩阵Q和矩阵K暂存在片上。

4.根据权利要求3所述的一种注意力机制计算方法，其特征在于：

5.根据权利要求3所述的一种注意力机制计算方法，其特征在于：片上的存储空间包括共享内存。