code: def _bwd_kernel(
      Q,
      K,
      V,
      Mask,
      Bias,
      sm_scale,
      Out,
      DO,
      DQ,
      DK,
      DV,
      DP,
      L,
      M,
      D,
      stride_qz,
      stride_qh,
      stride_qm,
      stride_qk,
      stride_kz,
      stride_kh,
      stride_kn,
      stride_kk,
      stride_vz,
      stride_vh,
      stride_vk,
      stride_vn,
      stride_mz,
      stride_mh,
      stride_mm,
      stride_mn,
      stride_bz,
      stride_bh,
      stride_bm,
      stride_bn,
      stride_dpz,
      stride_dph,
      stride_dpm,
      stride_dpn,
      stride_dob,
      stride_doh,
      stride_dom,
      stride_dok,
      stride_dqb,
      stride_dqh,
      stride_dqm,
      stride_dqk,
      stride_dkb,
      stride_dkh,
      stride_dkn,
      stride_dkk,
      stride_dvb,
      stride_dvh,
      stride_dvn,
      stride_dvk,
      Z,
      H,
      N_CTX,
      H_DIM,
      # num_block,
      inf: tl.constexpr,
      BLOCK_M: tl.constexpr,
      BLOCK_DMODEL: tl.constexpr,
      BLOCK_N: tl.constexpr,
      use_mask: tl.constexpr,
      use_bias: tl.constexpr,
      EVEN_M: tl.constexpr,
      EVEN_N: tl.constexpr,
      EVEN_HEADDIM: tl.constexpr,
      SEQUENCE_PARALLEL: tl.constexpr,
  ):
      off_hz = tl.program_id(0)
      off_b = off_hz // H
      off_h = off_hz % H
  
      # offset pointers for batch/head
      Q += off_b * stride_qz + off_h * stride_qh
      K += off_b * stride_kz + off_h * stride_kh
      V += off_b * stride_vz + off_h * stride_vh
      DO += off_b * stride_dob + off_h * stride_doh
      DQ += off_b * stride_dqb + off_h * stride_dqh
      DK += off_b * stride_dkb + off_h * stride_dkh
      DV += off_b * stride_dvb + off_h * stride_dvh
      DP += off_b * stride_dpz + off_h * stride_dph
  
      if use_bias:
          Bias += off_b * stride_bz + off_h * stride_bh
      if use_mask:
          # offs_base_mask = off_b * N_CTX
          Mask += off_b * stride_mz + off_h * stride_mh
  
      num_block_n = tl.cdiv(N_CTX, BLOCK_N)
      for start_n in range(0, num_block_n):
          # lo = start_n * BLOCK_M
          lo = 0
          # initialize row/col offsets
          offs_qm = lo + tl.arange(0, BLOCK_M)
          offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)  # BLOCK_M
          offs_m = tl.arange(0, BLOCK_M)  # BLOCK_N
          offs_k = tl.arange(0, BLOCK_DMODEL)
          # initialize pointers to value-like data
          q_ptrs = Q + (
              offs_qm[:, None] * stride_qm + offs_k[None, :] * stride_qk
          )
          k_ptrs = K + (
              offs_n[:, None] * stride_kn + offs_k[None, :] * stride_kk
          )
          v_ptrs = V + (
              offs_n[:, None] * stride_vk + offs_k[None, :] * stride_vn
          )
          do_ptrs = DO + (
              offs_qm[:, None] * stride_dom
              + offs_k[None, :] * stride_dok
          )
          dq_ptrs = DQ + (
              offs_qm[:, None] * stride_dqm
              + offs_k[None, :] * stride_dqk
          )
          dp_ptrs = DP + (
              offs_qm[:, None] * stride_dpm
              + offs_n[None, :] * stride_dpn
          )
          if use_bias:
              b_ptrs = Bias + (
                  offs_qm[:, None] * stride_bm
                  + offs_n[None, :] * stride_bn
              )
          if use_mask:
              mask_ptrs = Mask + (
                  offs_qm[:, None] * stride_mm
                  + offs_n[None, :] * stride_mn
              )
          # pointer to row-wise quantities in value-like data
          D_ptrs = D + off_hz * N_CTX
          m_ptrs = M + off_hz * N_CTX
          # initialize dv amd dk
          dv = tl.zeros(
              [BLOCK_N, BLOCK_DMODEL], dtype=tl.float32
          )  # BLOCK_M
          dk = tl.zeros(
              [BLOCK_N, BLOCK_DMODEL], dtype=tl.float32
          )  # BLOCK_M
          # k and v stay in SRAM throughout
          if EVEN_N & EVEN_M:
              if EVEN_HEADDIM:
                  k = tl.load(k_ptrs)
                  v = tl.load(v_ptrs)
              else:
                  k = tl.load(
                      k_ptrs, mask=offs_k[None, :] < H_DIM, other=0.0
                  )
                  v = tl.load(
                      v_ptrs, mask=offs_k[None, :] < H_DIM, other=0.0
                  )
          else:
              if EVEN_HEADDIM:
                  k = tl.load(
                      k_ptrs, mask=offs_n[:, None] < N_CTX, other=0.0
                  )
                  v = tl.load(
                      v_ptrs, mask=offs_n[:, None] < N_CTX, other=0.0
                  )
              else:
                  k = tl.load(
                      k_ptrs,
                      mask=(offs_n[:, None] < N_CTX)
                      & (offs_k[None, :] < H_DIM),
                      other=0.0,
                  )
                  v = tl.load(
                      v_ptrs,
                      mask=(offs_n[:, None] < N_CTX)
                      & (offs_k[None, :] < H_DIM),
                      other=0.0,
                  )
          # loop over rows
          num_block_m = tl.cdiv(N_CTX, BLOCK_M)
          for start_m in range(lo, num_block_m * BLOCK_M, BLOCK_M):
              start_m = tl.multiple_of(start_m, BLOCK_M)
              offs_m_curr = start_m + offs_m
              # load q, k, v, do on-chip
              if EVEN_M & EVEN_HEADDIM:
                  q = tl.load(q_ptrs)
              else:
                  if EVEN_HEADDIM:
                      q = tl.load(
                          q_ptrs,
                          mask=offs_m_curr[:, None] < N_CTX,
                          other=0.0,
                      )
                  else:
                      q = tl.load(
                          q_ptrs,
                          mask=(offs_m_curr[:, None] < N_CTX)
                          & (offs_k[None, :] < H_DIM),
                          other=0.0,
                      )
              # recompute p = softmax(qk, dim=-1).T
              # NOTE: `do` is pre-divided by `l`; no normalization here
              qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
              qk += tl.dot(q, tl.trans(k))
  
              if use_bias:
                  tl.debug_barrier()  # Race condition otherwise
                  if EVEN_M & EVEN_N:
                      bias = tl.load(b_ptrs).to(tl.float32)
                  else:
                      bias = tl.load(
                          b_ptrs,
                          mask=(offs_m_curr[:, None] < N_CTX)
                          & (offs_n[None, :] < N_CTX),
                          other=0.0,
                      ).to(tl.float32)
                  qk = qk * sm_scale + bias
  
              if use_mask:
                  # tl.debug_barrier()  # Race condition otherwise
                  # qk = tl.where(offs_m_curr[:, None] >= N_CTX, float("-1e20"), qk)
                  # qk = tl.where(offs_n[None, :] >= N_CTX, float("-1e20"), qk)
                  # mask_data = tl.load(Mask + offs_base_mask + offs_n)
                  # qk = tl.where(mask_data[None, :] == 0., float("-1e20"), qk)
                  if EVEN_M & EVEN_N:
                      mask_data = tl.load(mask_ptrs).to(tl.float32)
                  else:
                      mask_data = tl.load(
                          mask_ptrs,
                          mask=(offs_m_curr[:, None] < N_CTX)
                          & (offs_n[None, :] < N_CTX),
                          other=0.0,
                      ).to(tl.float32)
  
                  qk += tl.where(mask_data == 0.0, -inf, 0.0)
                  # qk = tl.where(mask_data == 0., -inf, qk)
  
              m = tl.load(m_ptrs + offs_m_curr)
              if use_bias:
                  p = tl.exp(qk - m[:, None])
              else:
                  p = tl.exp(qk * sm_scale - m[:, None])
              # compute dv
              if EVEN_M & EVEN_HEADDIM:
                  do = tl.load(do_ptrs)  # .to(tl.float32)
              else:
                  do = tl.load(
                      do_ptrs,
                      mask=(offs_m_curr[:, None] < N_CTX)
                      & (offs_k[None, :] < H_DIM),
                      other=0.0,
                  )
  
              dv += tl.dot(tl.trans(p.to(Q.dtype.element_ty)), do)
              # compute dp = dot(v, do)
              Di = tl.load(D_ptrs + offs_m_curr)
  
              dp = (
                  tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
                  - Di[:, None]
              )
              dp += tl.dot(do, tl.trans(v))
  
              # compute ds = p * (dp - delta[:, None])
              ds = p * dp
              if use_bias:
                  tl.store(dp_ptrs, ds)
              ds = ds * sm_scale
  
              # compute dk = dot(ds.T, q)
              dk += tl.dot(tl.trans(ds.to(Q.dtype.element_ty)), q)
  
              # compute dq
              # can we remove .to(tl.float32)
              if (
                  EVEN_M & EVEN_HEADDIM
              ):  # Race condition if we just do EVEN_M
                  dq = tl.load(dq_ptrs).to(tl.float32)
                  dq += tl.dot(ds.to(Q.dtype.element_ty), k)
                  tl.store(dq_ptrs, dq)
              else:
                  if EVEN_HEADDIM:
                      dq = tl.load(
                          dq_ptrs,
                          mask=offs_m_curr[:, None] < N_CTX,
                          other=0.0,
                      ).to(tl.float32)
                      dq += tl.dot(ds.to(Q.dtype.element_ty), k)
                      tl.store(
                          dq_ptrs, dq, mask=offs_m_curr[:, None] < N_CTX
                      )
                  else:
                      dq = tl.load(
                          dq_ptrs,
                          mask=(offs_m_curr[:, None] < N_CTX)
                          & (offs_k[None, :] < H_DIM),
                          other=0.0,
                      ).to(tl.float32)
                      dq += tl.dot(ds.to(Q.dtype.element_ty), k)
                      tl.store(
                          dq_ptrs,
                          dq,
                          mask=(offs_m_curr[:, None] < N_CTX)
                          & (offs_k[None, :] < H_DIM),
                      )
              # increment pointers
              dq_ptrs += BLOCK_M * stride_dqm
              q_ptrs += BLOCK_M * stride_qm
              do_ptrs += BLOCK_M * stride_dom
  
              dp_ptrs += BLOCK_M * stride_dpm
              if use_bias:
                  b_ptrs += BLOCK_M * stride_bm
              if use_mask:
                  mask_ptrs += BLOCK_M * stride_mm
          # write-back
          dv_ptrs = DV + (
              offs_n[:, None] * stride_dvn
              + offs_k[None, :] * stride_dvk
          )
          dk_ptrs = DK + (
              offs_n[:, None] * stride_dkn
              + offs_k[None, :] * stride_dkk
          )
  
          if EVEN_N & EVEN_M:
              if EVEN_HEADDIM:
                  tl.store(dv_ptrs, dv)
                  tl.store(dk_ptrs, dk)
              else:
                  tl.store(dv_ptrs, dv, mask=offs_k[None, :] < H_DIM)
                  tl.store(dk_ptrs, dk, mask=offs_k[None, :] < H_DIM)
          else:
              if EVEN_HEADDIM:
                  tl.store(dv_ptrs, dv, mask=offs_n[:, None] < N_CTX)
                  tl.store(dk_ptrs, dk, mask=offs_n[:, None] < N_CTX)
              else:
                  tl.store(
                      dv_ptrs,
                      dv,
                      mask=(offs_n[:, None] < N_CTX)
                      & (offs_k[None, :] < H_DIM),
                  )
                  tl.store(
                      dk_ptrs,
                      dk,
                      mask=(offs_n[:, None] < N_CTX)
                      & (offs_k[None, :] < H_DIM),
                  )
  
file_path: apex/apex/contrib/openfold_triton/_mha_kernel.py
imports: import triton
  import triton.language
label_desc: Use image node_id 6 for calling a global function with example usage: _bwd_kernel(Q, K, V, Mask, Bias, sm_scale, Out, DO, DQ, DK, DV, DP, L, M, D, stride_qz, stride_qh, stride_qm, stride_qk, stride_kz, stride_kh, stride_kn, stride_kk, stride_vz, stride_vh, stride_vk, stride_vn, stride_mz, stride_mh, stride_mm, stride_mn, stride_bz, stride_bh, stride_bm, stride_bn, stride_dpz, stride_dph, stride_dpm, stride_dpn, stride_dob, stride_doh, stride_dom, stride_dok, stride_dqb, stride_dqh, stride_dqm, stride_dqk, stride_dkb, stride_dkh, stride_dkn, stride_dkk, stride_dvb, stride_dvh, stride_dvn, stride_dvk, Z, H, N_CTX, H_DIM, inf, BLOCK_M, BLOCK_DMODEL, BLOCK_N, use_mask, use_bias, EVEN_M, EVEN_N, EVEN_HEADDIM, SEQUENCE_PARALLEL) without return types