code: def test_launch_attn_softmax_bw_new():
      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))
      inp_grad = kt.rand((batch_size, nhead, from_len, to_len))
  
      if kt.dtype == torch.float:
          base_func = cuda_module.torch_launch_attn_softmax_bw_fp32
      else:
          base_func = cuda_module.torch_launch_attn_softmax_bw_fp16
  
      if kt.dtype == torch.float:
          cust_func = cuda_module.torch_launch_attn_softmax_bw_new_fp32
      else:
          cust_func = cuda_module.torch_launch_attn_softmax_bw_new_fp16
  
      inp_grad_dup = inp_grad.clone()
  
      def custom():
          out_grad_dup = out_grad.clone()
          cust_func(
              inp_grad_dup,
              out_grad_dup,
              soft_inp,
              batch_size * nhead * from_len,
              to_len,
          )
          return [
              inp_grad_dup,
          ]
  
      def baseline():
          base_out_grad = out_grad.clone()
          base_func(
              base_out_grad,
              soft_inp,
              batch_size * nhead * from_len,
              to_len,
          )
          return [
              base_out_grad,
          ]
  
      return custom, baseline
  
imports: import sys
  from __init__ import lightseq_dir
  import random
  import torch
  from torch._C import ScriptModule, dtype
  from torch.nn import functional
  from csrc.pytorch.builder import CudaKernelBuilder
  from csrc.tests.util import TestDecorator
  from torch_crf import CRF
  from transformers import LlamaModel
  from transformers.activations import SiLUActivation
label_desc: Use image node_id 8 for calling a global function with example usage: test_launch_attn_softmax_bw_new() and returns: custom,  baseline, list, list