ScaledMaskedSoftmax算子使用指南

算子IR及torch_npu接口参数

• 算子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)

• 参数定义：

  名称	类型	dtype	Shape要求	默认值
  x	输入	bfloat16, float16, float32	必须为4维, 且后两维都需要在[32, 4096]范围内，且能被32整除	/
  mask	输入	bool	必须为4维，且后两维和x一致，且能被广播成x的shape	/
  scale	属性	float	对输入x缩放	1.0
  fixed_triu_mask	属性	bool	是否生成可用的上三角bool掩码	False

模型中替换代码及算子计算逻辑

• 模型中替换代码：

  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)

• 计算流程图为：
  图1 流程图
  

算子替换的模型中小算子

使用限制

1. 输入x的shape限制如下：
   1. 必须为4维。
   2. 第三维的取值需要在[32, 4096]范围内。
   3. 第四维的取值需要在[32, 4096]范围内。
   4. 第三维的取值需要能被32整除。
   5. 第四维的取值需要能被32整除。
2. 输入mask的shape限制如下：
   1. 必须为4维。
   2. 后两维必须与x的后两维相等。
   3. 前两维需要能被广播成x的前两维。

已支持模型典型Case

如下case均包含fp16、fp32、bf16。