code: def backward(ctx, dout):
      qkv, S_dmask = ctx.saved_tensors
      batch_size = ctx.cu_seqlens.numel() - 1
      if batch_size < 4:
          dqkv, dp, _ = mha.bwd_nl(
              dout,
              qkv,
              S_dmask,
              ctx.cu_seqlens,
              ctx.p_dropout,
              ctx.max_s,
              ctx.zero_tensors,
          )
      else:
          dqkv, dp = mha.bwd(
              dout,
              qkv,
              S_dmask,
              ctx.cu_seqlens,
              ctx.p_dropout,
              ctx.max_s,
              ctx.zero_tensors,
          )
  
      return dqkv, None, None, None, None, None
  
file_path: apex/apex/contrib/fmha/fmha.py
imports: import torch
  import torch.nn.functional
  import fmhalib
label_desc: Use image node_id 2 for calling the FMHAFun obj's underlying member method code with example usage: obj.backward(ctx, dout) and returns: dqkv,  None,  None,  None,  None,  None