ScaledMaskedSoftmax & ScaledMaskedSoftmaxGrad

算子IR：

REG_OP(ScaledMaskedSoftmax)
    .INPUT(x, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16}))
    .OPTIONAL_INPUT(mask, TensorType({DT_BOOL, DT_UINT1}))
    .OUTPUT(y, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16}))
    .ATTR(scale, Float, 1.0)
    .ATTR(fixed_triu_mask, Bool, false)
    .OP_END_FACTORY_REG(ScaledMaskedSoftmax)

torch_npu接口：

torch_npu.npu_scaled_masked_softmax(x, mask, scale, fixed_triu_mask)

参数说明：

模型中替换代码：

if self.input_in_float16 and self.softmax_in_fp32:
    input = input.float()

if self.scale is not None:
    input = input * self.scale
mask_output = self.mask_func(input, mask) if mask is not None else input
probs = torch.nn.Softmax(dim=-1)(mask_output)

if self.input_in_float16 and self.softmax_in_fp32:
    if self.input_in_fp16:
        probs = probs.half()
    else:
        probs = probs.bfloat16()

替换为：

probs = torch_npu.npu_scaled_masked_softmax(input , mask, self.scale, fixed_triu_mask )

算子的计算逻辑如下：

if fixed_triu_mask:
    mask = torch.triu(mask.shape, diagonal=1)
y = torch.softmax((x * scale).masked_fill(mask, -inf), dim=-1)