code: def backward(ctx, do):
      q, k, v, o, m, l, bias = ctx.saved_tensors
      ori_do_size = len(do.size())
      if ori_do_size == 5:
          do = rearrange(do, "1 a b c d -> a b c d")
      do = do.contiguous()
      dq = torch.zeros_like(q, dtype=torch.float32)
      dk = torch.empty_like(k)
      dv = torch.empty_like(v)
      # bias.dtype
      Z, H, N_CTX, H_DIM = (
          q.shape[-4],
          q.shape[-3],
          q.shape[-2],
          q.shape[-1],
      )
      dp = torch.zeros(
          (Z, H, N_CTX, N_CTX), dtype=torch.float32, device="cuda"
      )
  
      do_scaled = torch.empty_like(do)
      delta = torch.empty_like(l)
      mask = ctx.mask
      inf = ctx.inf
  
      BLOCK = 128
      BLOCK_HEADDIM = max(triton.next_power_of_2(H_DIM), 16)
      grid = (triton.cdiv(N_CTX, BLOCK) * Z * H, 1)
      _bwd_preprocess[grid](
          o,
          do,
          l,
          do_scaled,
          delta,
          o.stride(0),
          o.stride(1),
          o.stride(2),
          o.stride(3),
          do.stride(0),
          do.stride(1),
          do.stride(2),
          do.stride(3),
          BLOCK_M=BLOCK,
          D_HEAD=BLOCK_HEADDIM,
      )
  
      if bias is not None:
          assert bias.dtype in [q.dtype, torch.float]
          assert bias.is_cuda
          assert bias.dim() == 4
          assert bias.stride(-1) == 1
          bias = bias.expand(Z, H, N_CTX, N_CTX)
  
      # if mask is not None:
      #    mask = mask.expand(Z, H, N_CTX, N_CTX)
  
      bias_strides = (
          (
              bias.stride(0),
              bias.stride(1),
              bias.stride(2),
              bias.stride(3),
          )
          if bias is not None
          else (0, 0, 0, 0)
      )
      mask_strides = (
          (
              mask.stride(0),
              mask.stride(1),
              mask.stride(2),
              mask.stride(3),
          )
          if mask is not None
          else (0, 0, 0, 0)
      )
  
      # BLOCK_M, BLOCK_N = 128, 64
      BLOCK_M, BLOCK_N, num_warps, num_stages = schedule_triton_mha(
          list(q.shape), fwd=False
      )
      # grid = lambda META: (triton.cdiv(N_CTX, META["BLOCK_N"]), Z * H)
      # grid = lambda META: (Z * H, triton.cdiv(N_CTX, META["BLOCK_N"]))
      # grid = lambda META: (triton.cdiv(N_CTX, META["BLOCK_N"]) if META["SEQUENCE_PARALLEL"] else 1,
      #            Z * H)
      grid = lambda META: (Z * H,)
      _bwd_kernel[grid](
          q,
          k,
          v,
          mask,
          bias,
          ctx.sm_scale,
          o,
          do_scaled,
          dq,
          dk,
          dv,
          dp,
          l,
          m,
          delta,
          q.stride(0),
          q.stride(1),
          q.stride(2),
          q.stride(3),
          k.stride(0),
          k.stride(1),
          k.stride(2),
          k.stride(3),
          v.stride(0),
          v.stride(1),
          v.stride(2),
          v.stride(3),
          *mask_strides,
          *bias_strides,
          dp.stride(0),
          dp.stride(1),
          dp.stride(2),
          dp.stride(3),
          do.stride(0),
          do.stride(1),
          do.stride(2),
          do.stride(3),
          dq.stride(0),
          dq.stride(1),
          dq.stride(2),
          dq.stride(3),
          dk.stride(0),
          dk.stride(1),
          dk.stride(2),
          dk.stride(3),
          dv.stride(0),
          dv.stride(1),
          dv.stride(2),
          dv.stride(3),
          q.shape[0],
          q.shape[1],
          q.shape[2],
          q.shape[3],
          # ctx.grid[0], # to delete
          inf=inf,
          BLOCK_M=BLOCK_M,
          BLOCK_N=BLOCK_N,
          BLOCK_DMODEL=ctx.BLOCK_DMODEL,
          use_mask=(mask != None),
          use_bias=(bias != None),
          num_warps=num_warps,
          num_stages=num_stages,
          SEQUENCE_PARALLEL=False,
      )
      dB = None
      if bias is not None:
          dB = torch.sum(dp, dim=-4, keepdim=True)
          if len(bias.size()) == 4:
              dB = rearrange(dB, "b2 h n d -> 1 b2 h n d")
      # print(h.asm["ttgir"])
  
      if ori_do_size == 5:
          dq = rearrange(dq, "b2 h n d -> 1 b2 h n d")
          dk = rearrange(dk, "b2 h n d -> 1 b2 h n d")
          dv = rearrange(dv, "b2 h n d -> 1 b2 h n d")
  
      return dq, dk, dv, None, dB, None, None
  
file_path: apex/apex/contrib/openfold_triton/mha.py
imports: import math
  from typing import Optional
  import torch
  import triton
  from einops import rearrange
  from apex.contrib.openfold_triton._mha_kernel import _attention_core, _bwd_kernel, _bwd_preprocess
label_desc: Use image node_id 2 for calling the FusedAttenionCoreFunc obj's underlying member method code with example usage: obj.backward(ctx, do) and returns: dq,  dk,  dv,  None,  dB,  None,  None