（beta）torch_npu.npu_fused_attention_score

接口原型

torch_npu.npu_fused_attention_score(Tensor query_layer, Tensor key_layer, Tensor value_layer, Tensor attention_mask, Scalar scale, float keep_prob, bool query_transpose=False, bool key_transpose=False, bool bmm_score_transpose_a=False, bool bmm_score_transpose_b=False, bool value_transpose=False, bool dx_transpose=False) -> Tensor

功能描述

实现“Transformer attention score”的融合计算逻辑，主要将matmul、transpose、add、softmax、dropout、batchmatmul、permute等计算进行了融合。

参数说明

• query_layer：Tensor类型，仅支持float16。
• key_layer：Tensor类型，仅支持float16。
• value_layer：Tensor类型，仅支持float16 。
• attention_mask：Tensor类型，仅支持float16 。
• scale：缩放系数，浮点数标量 。
• keep_prob：不做dropout的概率，0-1之间，浮点数。
• query_transpose：query是否做转置，bool类型，默认为False 。
• key_transpose：key是否做转置，bool类型，默认为False 。
• bmm_score_transpose_a：bmm计算中左矩阵是否做转置，bool类型，默认为False。
• bmm_score_transpose_b：bmm计算中右矩阵是否做转置，bool类型，默认为False。
• value_transpose：value是否做转置，bool类型，默认为False。
• dx_transpose：反向计算时dx是否做转置，bool类型，默认为False。

约束说明

输入tensor的格式编号必须均为29，数据类型为FP16。

支持的型号

Atlas 训练系列产品

调用示例

>>> import torch
>>> import torch_npu
>>> query_layer = torch_npu.npu_format_cast(torch.rand(24, 16, 512, 64).npu(), 29).half()
>>> key_layer = torch_npu.npu_format_cast(torch.rand(24, 16, 512, 64).npu(), 29).half()
>>> value_layer = torch_npu.npu_format_cast(torch.rand(24, 16, 512, 64).npu(), 29).half()
>>> attention_mask = torch_npu.npu_format_cast(torch.rand(24, 16, 512, 512).npu(), 29).half()
>>> scale = 0.125
>>> keep_prob = 0.5
>>> context_layer = torch_npu.npu_fused_attention_score(query_layer, key_layer, value_layer, attention_mask, scale, keep_prob)
>>> print(context_layer)
        tensor([[0.5063, 0.4900, 0.4951,  ..., 0.5493, 0.5249, 0.5400],
               [0.4844, 0.4724, 0.4927,  ..., 0.5176, 0.4702, 0.4790],
               [0.4683, 0.4771, 0.5054,  ..., 0.4917, 0.4614, 0.4739],
               ...,
               [0.5137, 0.5010, 0.5078,  ..., 0.4656, 0.4592, 0.5034],
               [0.5425, 0.5732, 0.5347,  ..., 0.5054, 0.5024, 0.4924],