class_docstr: fused operation: scaling + mask + softmax
  
  Arguments:
      input_in_fp16: flag to indicate if input in fp16 data format.
      input_in_bf16: flag to indicate if input in bf16 data format.
      attn_mask_type: attention mask type (pad or causal)
      scaled_masked_softmax_fusion: flag to indicate user want to use softmax fusion
      mask_func: mask function to be applied.
      softmax_in_fp32: if true, softmax in performed at fp32 precision.
      scale: scaling factor used in input tensor scaling.
code: def forward_fused_softmax(self, input, mask):
      b, np, sq, sk = input.size()
      scale = self.scale if self.scale is not None else 1.0
  
      if self.attn_mask_type == AttnMaskType.causal:
          assert sq == sk, "causal mask is only for self attention"
  
          # input is 3D tensor (attn_batches, sq, sk)
          input = input.view(-1, sq, sk)
          probs = ScaledUpperTriangMaskedSoftmax.apply(input, scale)
          return probs.view(b, np, sq, sk)
      else:
          # input is 4D tensor (b, np, sq, sk)
          if mask is not None:
              return ScaledMaskedSoftmax.apply(input, mask, scale)
          else:
              return ScaledSoftmax.apply(input, scale)
  
imports: import torch
  import torch.nn
  from megatron.model.enums import AttnMaskType
label_desc: Use image node_id 4 for calling the FusedScaleMaskSoftmax obj's underlying member method code with example usage: obj.forward_fused_softmax(input, mask) and returns: probs, ScaledMaskedSoftmax, ScaledSoftmax