code: def softmax_forward(
      self, attention_scores, attention_mask, query_layer, key_layer
  ):
      attention_scores = attention_scores / math.sqrt(
          self.attention_head_size
      )
      # causal mask
      query_length, key_length = query_layer.size(-2), key_layer.size(
          -2
      )
      causal_mask = self.bias[
          :, :, key_length - query_length : key_length, :key_length
      ].bool()
      attention_scores = torch.where(
          causal_mask,
          attention_scores,
          self.masked_bias.to(attention_scores),
      )
      if attention_mask is not None:
          # Apply the attention mask
          attention_scores = attention_scores + attention_mask
      attention_scores = self.softmax(attention_scores)
      return attention_scores
  
imports: import math
  import torch
  from torch import Tensor
  from torch import nn
  from colossalai import kernel
  from colossalai import nn
  from colossalai.kernel.cuda_native.scaled_softmax import AttnMaskType
  from colossalai.legacy.core import global_context
  from colossalai.legacy.nn.layer import Linear1D_Col, Linear1D_Row
  from colossalai.legacy.nn.layer.base_layer import ParallelLayer
  from colossalai.legacy.nn.layer.utils import ACT2FN, divide
  from colossalai.legacy.utils.activation_checkpoint import checkpoint
  from colossalai.utils import checkpoint
label_desc: Use image node_id 2 for calling the GPTSelfAttention1D obj's underlying member method code with example usage: obj.softmax_forward(attention_scores, attention_mask, query_layer, key_layer) and returns: attention_scores