code: def test_fused_softmax():
      bert = BertModel.from_pretrained("bert-base-cased").cuda().half()
      tokenizer = BertTokenizer.from_pretrained("bert-base-cased")
      test_text = (
          "Hello. How are you? I am fine thank you and you? yes Good. "
          "hi hi hi hi hi hi hi hi hi hi hi hi hi"  # 32
      )
  
      tokens = tokenizer(
          [test_text] * 4,
          return_tensors="pt",
      )
  
      embedding_output = bert.embeddings(
          input_ids=tokens["input_ids"].cuda(),
          position_ids=None,
          token_type_ids=tokens["token_type_ids"].cuda(),
          inputs_embeds=None,
          past_key_values_length=0,
      )
  
      # (bsz, 1, 1, seq_len)
      mask = bert.get_extended_attention_mask(
          attention_mask=tokens["attention_mask"].cuda(),
          input_shape=tokens["input_ids"].shape,
          device=bert.device,
      )
      # (bsz, 1, seq_len, seq_len)
      mask = mask.repeat(1, 1, mask.size()[-1], 1)
  
      attention = bert.encoder.layer[0].attention.self
      key_layer = attention.transpose_for_scores(
          attention.key(embedding_output)
      )
      query_layer = attention.transpose_for_scores(
          attention.query(embedding_output)
      )
  
      attention_scores = torch.matmul(
          query_layer, key_layer.transpose(-1, -2)
      )
      attention_scores /= math.sqrt(key_layer.size()[-1])
  
      fused_softmax = (
          FusedScaleMaskSoftmax(
              input_in_fp16=True,
              input_in_bf16=False,
              mask_func=attention_mask_func,
              scale=None,
              softmax_in_fp32=False,
              attn_mask_type=AttnMaskType.padding,
              scaled_masked_softmax_fusion=True,
          )
          .cuda()
          .half()
      )
  
      fused_softmax_output = fused_softmax(
          attention_scores,
          (mask != 0),
      )
  
      torch_softmax = (
          FusedScaleMaskSoftmax(
              input_in_fp16=True,
              input_in_bf16=False,
              mask_func=attention_mask_func,
              scale=None,
              softmax_in_fp32=False,
              attn_mask_type=AttnMaskType.padding,
              scaled_masked_softmax_fusion=False,
          )
          .cuda()
          .half()
      )
  
      torch_softmax_output = torch_softmax(
          attention_scores,
          (mask != 0),
      )
  
      test_result = (fused_softmax_output - torch_softmax_output).abs()
  
      while test_result.dim() != 1:
          test_result = test_result.mean(dim=-1)
  
      diff = test_result.mean(dim=-1)
  
      if diff <= 1e-3:
          print(
              f"\n[Success] test_fused_softmax"
              f"\n > mean_difference={diff}"
              f"\n > fused_values={fused_softmax_output[-1][-1][-1][:5].tolist()}"
              f"\n > torch_values={torch_softmax_output[-1][-1][-1][:5].tolist()}"
          )
      else:
          print(
              f"\n[Fail] test_fused_softmax"
              f"\n > mean_difference={diff}, "
              f"\n > fused_values={fused_softmax_output[-1][-1][-1][:5].tolist()}, "
              f"\n > torch_values={torch_softmax_output[-1][-1][-1][:5].tolist()}"
          )
  
file_path: megatron-lm/megatron/fused_kernels/tests/test_fused_kernels.py
imports: import math
  import torch
  from torch.nn import LayerNorm
  from megatron.model.enums import AttnMaskType
  from megatron.model.fused_layer_norm import MixedFusedLayerNorm
  from megatron.model.fused_softmax import FusedScaleMaskSoftmax
  from megatron.model.utils import attention_mask_func
  from megatron.fused_kernels import load
label_desc: Use image node_id 2 for calling a global function with example usage: test_fused_softmax() without return types