code: def _flash_attn_backward(
      dout,
      q,
      k,
      v,
      out,
      softmax_lse,
      dq,
      dk,
      dv,
      cu_seqlens_q,
      cu_seqlens_k,
      max_seqlen_q,
      max_seqlen_k,
      dropout_p,
      softmax_scale,
      causal,
      generator=None,
  ):
      softmax_d = flash_attn_cuda.bwd(
          dout,
          q,
          k,
          v,
          out,
          softmax_lse,
          dq,
          dk,
          dv,
          cu_seqlens_q,
          cu_seqlens_k,
          max_seqlen_q,
          max_seqlen_k,
          dropout_p,
          softmax_scale,
          False,
          causal,
          generator,
      )
      # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any():
      #     breakpoint()
      return dq, dk, dv, softmax_d
  
imports: import torch
  import torch.nn
  import torch.nn.functional
  import flash_attn_cuda
label_desc: Use image node_id 3 for calling a global function with example usage: _flash_attn_backward(dout, q, k, v, out, softmax_lse, dq, dk, dv, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, dropout_p, softmax_scale, causal, generator) and returns: dq,  dk,  dv,  softmax_d