code: def test_launch_attn_softmax():
      batch_size, from_len = kt.bs_sl()
      is_dec_self_attn = random.choice([True, False])
      if is_dec_self_attn:
          to_len = from_len
          is_dec_self_attn_infer = random.choice([True, False])
      else:
          _, to_len = kt.bs_sl(batch_size)
          is_dec_self_attn_infer = False
  
      if is_dec_self_attn_infer:
          to_len = from_len
          from_len = 1
          beam_size = random.choice([3, 4, 5])
          batch_size *= beam_size
  
      nhead = kt.nhead
      print(
          "(batch_size, nhead, from_len, to_len, is_dec_self_attn,"
          f" is_dec_self_attn_infer): ({batch_size}, {nhead}, {from_len}, {to_len},"
          f" {is_dec_self_attn}, {is_dec_self_attn_infer})"
      )
  
      inp = kt.rand((batch_size, nhead, from_len, to_len))
      if is_dec_self_attn:
          mask = kt.dec_self_attn_mask(to_len) * -1e8
          mask = mask.unsqueeze(0).unsqueeze(
              0
          )  # [1, 1, to_len, to_len]
      else:
          mask = kt.attn_mask(batch_size, to_len) * -1e8
          mask = mask.unsqueeze(1).unsqueeze(
              1
          )  # [batch_size, 1, 1, to_len]
      inp_list = [inp.clone() for _ in range(8)]
      tt = {"repeat_idx": 0}
  
      if kt.dtype == torch.float:
          func = cuda_module.torch_launch_attn_softmax_fp32
      else:
          func = cuda_module.torch_launch_attn_softmax_fp16
  
      def custom():
          inp_dup = inp_list[tt["repeat_idx"]]
          tt["repeat_idx"] += 1
          func(
              inp_dup,
              mask,
              batch_size,
              nhead,
              from_len,
              to_len,
              is_dec_self_attn,
              is_dec_self_attn and (not is_dec_self_attn_infer),
          )
          return [
              inp_dup,
          ]
  
      def baseline():
          f_inp = inp.clone().to(torch.float32)
          f_mask = mask.to(torch.float32)
          if not is_dec_self_attn_infer:
              res = functional.softmax(
                  f_inp + f_mask, dim=-1, dtype=torch.float32
              )
          else:
              res = functional.softmax(
                  f_inp, dim=-1, dtype=torch.float32
              )
          return kt.norm_res_list(res)
  
      return custom, baseline
  
imports: import random
  import torch
  from torch._C import ScriptModule, dtype
  from torch.nn import functional
  from lightseq.training.ops.pytorch.builder import KernelBuilder, AdamBuilder
  from tests.util import TestDecorator
  from torch_crf import CRF
label_desc: Use image node_id 5 for calling a global function with example usage: test_launch_attn_softmax() and returns: custom,  baseline, list, kt