用于进行注意力运算的方法、设备和存储介质与流程

本公开的实施例总体涉及注意力运算领域，并且更具体地涉及一种用于进行注意力运算的方法、设备和存储介质。

背景技术：

1、注意力（attention）运算越来越广泛地应用于深度神经网络中。传统的注意力运算方案中，存在运算性能较低，或者是用于运算的存储器资源利用率较低，导致运算性能下降等问题。

技术实现思路

1、针对上述问题，本公开提供了一种用于进行注意力运算的方法、设备和存储介质，能够显著提高注意力运算的效率。

2、根据本公开的第一方面，提供一种用于进行注意力运算的方法。该方法包括：确定查询张量的关于序列长度（sequence length）的第一维度的第一数值、以及关于批量大小（batch size）的第二维度的第二数值；根据第一数值以及第二数值，确定目标维度；以及基于目标维度，将查询张量划分为多个第一子张量，以便多个处理器核中的每一个处理器核分别针对多个第一子张量中的一个第一子张量进行注意力计算。

3、在一些实施例中，确定查询张量的关于序列长度的第一维度的第一数值、以及关于批量大小的第二维度的第二数值包括：基于查询张量的序列长度与处理器核的关于数据对齐的最小数据粒度的比值，确定第一数值；以及基于查询张量的批量大小与查询张量的注意力头的数量（head number）的乘积，确定第二数值。

4、在一些实施例中，根据第一数值以及第二数值确定目标维度包括：确定第一数值以及第二数值是否均小于处理器核的数量；以及响应于确定第一数值以及第二数值均小于处理器核的数量，确定目标维度为第一维度和第二维度。

5、在一些实施例中，根据第一数值以及第二数值确定目标维度还包括：响应于确定第一数值以及第二数值不均小于处理器核的数量，确定第一数值以及第二数值中较大的一个所对应的维度为目标维度。

6、在一些实施例中，多个处理器核中的每一个处理器核分别针对多个第一子张量中的一个第一子张量进行注意力计算包括：根据处理器核的共享存储器的大小，针对第一子张量的目标划分维度将第一子张量划分为多个第一组块；针对键张量的目标划分维度，将键张量划分为多个第二组块；以及针对值张量的目标划分维度将值张量划分为多个第三组块；以及基于多个第一组块、多个第二组块以及多个第三组块进行注意力计算。

7、在一些实施例中，第二组块在键张量的目标划分维度上所对应的维度值与第三组块在值张量的目标划分维度上所对应的维度值相等。

8、在一些实施例中，基于多个第一组块、多个第二组块以及多个第三组块进行注意力计算包括：将多个第一组块中的每一个第一组块分别与多个第二组块中的对应第二组块依次进行矩阵乘运算，以便得到与多个第二组块对应的多个第一中间张量；以及根据多个第一中间张量以及多个第三组块，确定用于进行softmax运算的分子以及分母，以便根据该分子以及该分母确定softmax运算的结果。

9、在一些实施例中，根据多个第一中间张量以及多个第三组块确定用于进行softmax运算的分子以及分母包括：依次针对多个第一中间张量中的每一个第一中间张量执行以下操作：确定当前最大值张量，当前最大值张量为当前第一中间张量以及当前第一中间张量之前的第一中间张量中最大的第一中间张量；针对当前第一中间张量与当前最大值张量之间的差值进行指数运算，以确定第二中间张量；以及根据前一个第一中间张量对应的s向量、当前最大值张量、第二中间张量以及当前第一中间张量的前一个第一中间张量，确定当前第一中间张量对应的s向量，以便确定用于进行softmax运算的分母。

10、在一些实施例中，根据多个第一中间张量以及多个第三组块确定用于进行softmax运算的分子以及分母包括：针对多个第一中间张量中的最后一个第一中间张量，更新s向量；以及输出更新后的s向量，以便作为用于进行softmax运算的分母。

11、在一些实施例中，根据前一个第一中间张量对应的s向量、当前最大值张量、第二中间张量以及当前第一中间张量的前一个第一中间张量，确定当前第一中间张量对应的s向量包括：针对前一个最大张量与当前最大值张量之间的差值进行指数运算，以便得到第三中间张量；以及将当前的第二中间张量以及当前第二中间张量之前的第二中间张量的累加和，以及前一个第一中间张量对应的s向量与第三中间张量的乘积进行加和，以作为当前第一中间张量对应的s向量。

12、在一些实施例中，根据多个第一中间张量以及多个第三组块确定用于进行softmax运算的分子以及分母还包括：依次针对多个第一中间张量中的每一个第一中间张量执行以下操作：根据前一个第一中间张量所对应的第四中间张量与第三中间张量的乘积，确定第五中间张量；将第二中间张量与对应的第三组块进行矩阵乘操作，以便得到第六中间张量；以及根据第五中间张量与第六中间张量的和，确定当前第一中间张量对应的第四中间张量，以便确定用于进行softmax运算的分子。

13、在一些实施例中，根据该分子以及该分母确定softmax运算的结果包括：在确定多个第一中间张量中的最后一个第一中间张量对应的第四中间张量之后，根据第四中间张量与更新后的s向量的商确定softmax运算的结果。

14、在一些实施例中，该方法还包括：将第一中间张量存储至共享存储器；以及将第二中间张量存储至共享存储器。

15、根据本公开的第二方面，提供一种电子设备，该电子设备包括：至少一个处理器；以及与该至少一个处理器通信连接的存储器；该存储器存储有可被该至少一个处理器执行的指令，该指令被该至少一个处理器执行，以使该至少一个处理器能够执行根据本公开的第一方面的方法。

16、根据本公开的第三方面，提供一种计算机可读存储介质，该计算机可读存储介质上存储有计算机程序，该计算机程序被机器执行时实现根据本公开的第一方面的方法。

17、应当理解，本部分所描述的内容并非旨在标识本公开的实施例的关键或重要特征，也不用于限制本公开的范围。本公开的其它特征将通过以下的说明书而变得容易理解。

技术特征：

1.一种用于进行注意力运算的方法，其特征在于，包括：

2.如权利要求1所述的方法，其特征在于，确定查询张量的关于序列长度的第一维度的第一数值、以及关于批量大小的第二维度的第二数值包括：

3. 如权利要求2所述的方法，其特征在于，根据第一数值以及第二数值确定目标维度包括：

4.如权利要求3所述的方法，其特征在于，根据第一数值以及第二数值确定目标维度还包括：

5.如权利要求1所述的方法，其特征在于，多个处理器核中的每一个处理器核分别针对多个第一子张量中的一个第一子张量进行注意力计算包括：

6. 如权利要求5所述的方法，其特征在于，基于多个第一组块、多个第二组块以及多个第三组块进行注意力计算包括：

7.如权利要求6所述的方法，其特征在于，根据多个第一中间张量以及多个第三组块确定用于进行softmax运算的分子以及分母包括：

8. 如权利要求7所述的方法，其特征在于，根据多个第一中间张量以及多个第三组块确定用于进行softmax运算的分子以及分母包括：

9. 如权利要求7所述的方法，其特征在于，根据前一个第一中间张量对应的s向量、当前最大值张量、第二中间张量以及当前第一中间张量的前一个第一中间张量，确定当前第一中间张量对应的s向量包括：

10.如权利要求9所述的方法，其特征在于，根据多个第一中间张量以及多个第三组块确定用于进行softmax运算的分子以及分母还包括：

11.如权利要求10所述的方法，其特征在于，根据所述分子以及所述分母确定softmax运算的结果包括：

12. 如权利要求11所述的方法，其特征在于，还包括：

13.如权利要求5所述的方法，其特征在于，第二组块在键张量的目标划分维度上所对应的维度值与第三组块在值张量的目标划分维度上所对应的维度值相等。

14. 一种电子设备，其特征在于，包括：

15.一种计算机可读存储介质，计算机可读存储介质上存储有计算机程序，其特征在于，所述计算机程序被机器执行时实现根据权利要求1至13中任一项所述的方法。

技术总结
本公开提供一种用于进行注意力运算的方法、设备和存储介质。其中方法包括确定查询张量的关于序列长度的第一维度的第一数值、以及关于批量大小的第二维度的第二数值；根据第一数值以及第二数值，确定目标维度；以及基于目标维度，将查询张量划分为多个第一子张量，以便多个处理器核中的每一个处理器核分别针对多个第一子张量中的一个第一子张量进行注意力计算。本公开的方案能够显著提高注意力运算的效率。

技术研发人员：请求不公布姓名,请求不公布姓名,请求不公布姓名,请求不公布姓名
受保护的技术使用者：北京壁仞科技开发有限公司
技术研发日：
技术公布日：2024/3/17