code: def test_fused_upper_triangle_mask_softmax():
      gpt = GPT2Model.from_pretrained("gpt2").cuda().half()
      tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
      test_text = (
          "Hello. How are you? I am fine thank you and you? yes Good. "
          "hi hi hi hi hi hi hi"  # 24
      )
  
      tokens = tokenizer(
          [test_text] * 4,
          return_tensors="pt",
      )
  
      attention_mask = tokens["attention_mask"].cuda()
      attention_mask = attention_mask.view(attention_mask.size(0), -1)
      attention_mask = attention_mask[:, None, None, :]
      attention_mask = (1.0 - attention_mask) * -10000.0
      attention_mask = attention_mask.repeat(
          1, 1, attention_mask.size()[-1], 1
      )
      attn = gpt.h[0]
  
      hidden_states = gpt.wte(tokens["input_ids"].cuda())
      q, k, v = attn.attn.c_attn(hidden_states).split(768, dim=-1)
      q = attn.attn._split_heads(
          q, attn.attn.num_heads, attn.attn.head_dim
      )
      k = attn.attn._split_heads(
          k, attn.attn.num_heads, attn.attn.head_dim
      )
      attn_weights = torch.matmul(q, k.transpose(-1, -2))
  
      sq, sk = q.size(-2), k.size(-2)
      causal_mask = attn.attn.bias[:, :, sk - sq : sk, :sk].bool()
      total_mask = ~(causal_mask & (attention_mask == 0))
      """
      tensor([[[[False,  True,  True,  ...,  True,  True,  True],
                [False, False,  True,  ...,  True,  True,  True],
                [False, False, False,  ...,  True,  True,  True],
                ...,
                [False, False, False,  ..., False,  True,  True],
                [False, False, False,  ..., False, False,  True],
                [False, False, False,  ..., False, False, False]]]
      """
  
      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.causal,
              scaled_masked_softmax_fusion=True,
          )
          .cuda()
          .half()
      )
  
      fused_softmax_output = fused_softmax(
          attn_weights,
          total_mask,
      )
  
      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.causal,
              scaled_masked_softmax_fusion=False,
          )
          .cuda()
          .half()
      )
  
      torch_softmax_output = torch_softmax(
          attn_weights,
          total_mask,
      )
  
      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_upper_triangle_mask_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_upper_triangle_mask_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 3 for calling a global function with example usage: test_fused_upper_triangle_mask_softmax() without return types