code: def test_launch_attn_softmax_bw():
      nhead = kt.nhead
      batch_size, from_len = kt.bs_sl()
      _, to_len = kt.bs_sl(batch_size)
      print(
          "(batch_size, nhead, from_len, to_len): "
          f"({batch_size}, {nhead}, {from_len}, {to_len})"
      )
  
      out_grad = kt.rand((batch_size, nhead, from_len, to_len))
      soft_inp = kt.rand((batch_size, nhead, from_len, to_len))
  
      if kt.dtype == torch.float:
          func = cuda_module.torch_launch_attn_softmax_bw_fp32
      else:
          func = cuda_module.torch_launch_attn_softmax_bw_fp16
  
      def custom():
          out_grad_dup = out_grad.clone()
          func(
              out_grad_dup,
              soft_inp,
              batch_size * nhead * from_len,
              to_len,
          )
          return [
              out_grad_dup,
          ]
  
      def baseline():
          f_soft_inp = soft_inp.to(dtype=torch.float)
          f_out_grad = out_grad.clone().to(dtype=torch.float)
          tsum = f_out_grad * f_soft_inp
          # [b, nh, s, 1]
          tsum = tsum.sum(dim=-1, keepdim=True)
          res = f_soft_inp * (f_out_grad - tsum)
          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 7 for calling a global function with example usage: test_launch_attn_softmax_bw() and returns: custom,  baseline, list, kt