code: def _attention_core(
      Q,
      K,
      V,
      Mask,
      Bias,
      sm_scale,
      L,
      M,
      Out,
      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_oz,
      stride_oh,
      stride_om,
      stride_on,
      stride_bz,
      stride_bh,
      stride_bm,
      stride_bn,
      stride_mz,
      stride_mh,
      stride_mm,
      stride_mn,
      Z,
      H,
      N_CTX,
      H_DIM,
      BATCH,  # 256 8 128 32 1
      inf: tl.constexpr,
      IS_TRAINING: tl.constexpr,
      BLOCK_M: tl.constexpr,
      BLOCK_N: tl.constexpr,
      BLOCK_DMODEL: tl.constexpr,
      use_mask: tl.constexpr,
      use_bias: tl.constexpr,
      EVEN_M: tl.constexpr,
      EVEN_N: tl.constexpr,
      EVEN_HEADDIM: tl.constexpr,
  ):
      start_m = tl.program_id(0)
      off_hz = tl.program_id(1)
      off_b = off_hz // H
      off_h = off_hz % H
      # initialize offsets
      offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
      offs_n = tl.arange(0, BLOCK_N)
      offs_d = tl.arange(0, BLOCK_DMODEL)
      off_q = (
          off_b * stride_qz
          + off_h * stride_qh
          + offs_m[:, None] * stride_qm
          + offs_d[None, :] * stride_qk
      )
      off_k = (
          off_b * stride_kz
          + off_h * stride_kh
          + offs_n[None, :] * stride_kn
          + offs_d[:, None] * stride_kk
      )
      off_v = (
          off_b * stride_vz
          + off_h * stride_vh
          + offs_n[:, None] * stride_vk
          + offs_d[None, :] * stride_vn
      )
      # Initialize pointers to Q, K, V
      q_ptrs = Q + off_q
      k_ptrs = K + off_k
      v_ptrs = V + off_v
  
      # Initialize pointers to bias, mask
      if use_bias:
          batch_2 = Z // BATCH
          off_hz_bias = (off_hz // (batch_2 * H) * H) + (off_hz % H)
          offs_base_bias = (
              off_hz_bias * (N_CTX * N_CTX)
              + offs_m[:, None] * N_CTX
              + offs_n[None, :]
          )
          """
          off_b = off_hz // H
          off_h = off_hz % H
          bias_ptrs = Bias + off_b * stride_bz + off_h * stride_bh + (offs_m[:, None] * stride_bm + offs_n[None, :] * stride_bn)
          """
  
      if use_mask:
          # off_hz_mask = (off_hz // H)
          # offs_base_mask = off_hz_mask * N_CTX
          off_b = off_hz // H
          off_h = off_hz % H
          mask_ptrs = (
              Mask
              + off_b * stride_mz
              + off_h * stride_mh
              + (
                  offs_m[:, None] * stride_mm
                  + offs_n[None, :] * stride_mn
              )
          )
  
      # initialize pointer to m and l
      m_prev = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
      l_prev = tl.zeros([BLOCK_M], dtype=tl.float32)
      acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
      # load q: it will stay in SRAM throughout
      if EVEN_M & EVEN_N:
          if EVEN_HEADDIM:
              q = tl.load(q_ptrs)
          else:
              q = tl.load(
                  q_ptrs, mask=offs_d[None, :] < H_DIM, other=0.0
              )
      else:
          if EVEN_HEADDIM:
              q = tl.load(
                  q_ptrs, mask=offs_m[:, None] < N_CTX, other=0.0
              )
          else:
              q = tl.load(
                  q_ptrs,
                  mask=(offs_m[:, None] < N_CTX)
                  & (offs_d[None, :] < H_DIM),
                  other=0.0,
              )
  
      # loop over k, v and update accumulator
      #  (start_m + 1) * BLOCK_M
      for start_n in range(0, N_CTX, BLOCK_N):
          start_n = tl.multiple_of(start_n, BLOCK_N)
          # -- compute qk ----
          if (
              EVEN_N & EVEN_M
          ):  # If we just do "if EVEN_N", there seems to be some race condition
              if EVEN_HEADDIM:
                  k = tl.load(k_ptrs)
              else:
                  k = tl.load(
                      k_ptrs, mask=offs_d[:, None] < H_DIM, other=0.0
                  )
          else:
              if EVEN_HEADDIM:
                  k = tl.load(
                      k_ptrs,
                      mask=(start_n + offs_n)[None, :] < N_CTX,
                      other=0.0,
                  )
              else:
                  k = tl.load(
                      k_ptrs,
                      mask=((start_n + offs_n)[None, :] < N_CTX)
                      & (offs_d[:, None] < H_DIM),
                      other=0.0,
                  )
  
          qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
  
          # qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk, float("-inf"))
          if use_bias:
              qk += tl.dot(q * sm_scale.to(tl.bfloat16), k).to(
                  tl.bfloat16
              )
              qk += tl.where(
                  (start_n + offs_n)[None, :] < N_CTX, 0, -inf
              ).to(tl.bfloat16)
              if EVEN_M & EVEN_N:
                  bias_data = tl.load(Bias + offs_base_bias + start_n)
              else:
                  bias_load_mask = tl.zeros(
                      [BLOCK_M, BLOCK_N], dtype=tl.float32
                  )
                  bias_load_mask = tl.where(
                      offs_m[:, None] >= N_CTX, 1.0, bias_load_mask
                  )
                  bias_load_mask = tl.where(
                      (start_n + offs_n)[None, :] >= N_CTX,
                      1.0,
                      bias_load_mask,
                  )
                  bias_data = tl.load(
                      Bias + offs_base_bias + start_n,
                      mask=(bias_load_mask == 0.0),
                      other=0.0,
                  )
              qk = qk + bias_data
          else:
              qk += tl.dot(q, k)
              qk += tl.where(
                  (start_n + offs_n)[None, :] < N_CTX, 0, -inf
              )
  
          qk = qk.to(tl.bfloat16)
  
          if use_mask:
              if EVEN_M & EVEN_N:
                  mask_data = tl.load(mask_ptrs + start_n).to(tl.int32)
              else:
                  mask_data = tl.load(
                      mask_ptrs + start_n,
                      mask=(offs_m[:, None] < N_CTX)
                      & ((start_n + offs_n)[None, :] < N_CTX),
                      other=0,
                  ).to(tl.int32)
              qk += tl.where(mask_data == 0, -inf, 0.0)
  
          if use_bias:
              # compute new m
              m_curr = tl.maximum(tl.max(qk, 1), m_prev)
              # correct old l
              l_prev *= tl.exp(m_prev - m_curr)
              # attention weights
              p = tl.exp(qk - m_curr[:, None])
          else:
              m_curr = tl.maximum(tl.max(qk, 1) * sm_scale, m_prev)
              l_prev *= tl.exp(m_prev - m_curr)
              p = tl.exp(qk * sm_scale - m_curr[:, None])
  
          l_curr = tl.sum(p, 1) + l_prev
          # rescale operands of matmuls
          l_rcp = 1.0 / l_curr
          p *= l_rcp[:, None]
          acc *= (l_prev * l_rcp)[:, None]
          # update acc
          p = p.to(Q.dtype.element_ty)
  
          if (
              EVEN_N & EVEN_M
          ):  # If we just do "if EVEN_N", there seems to be some race condition
              if EVEN_HEADDIM:
                  v = tl.load(v_ptrs)
              else:
                  v = tl.load(
                      v_ptrs, mask=offs_d[None, :] < H_DIM, other=0.0
                  )
          else:
              if EVEN_HEADDIM:
                  v = tl.load(
                      v_ptrs,
                      mask=(start_n + offs_n)[:, None] < N_CTX,
                      other=0.0,
                  )
              else:
                  v = tl.load(
                      v_ptrs,
                      mask=((start_n + offs_n)[:, None] < N_CTX)
                      & (offs_d[None, :] < H_DIM),
                      other=0.0,
                  )
          acc += tl.dot(p, v)
          # update m_i and l_i
          l_prev = l_curr
          m_prev = m_curr
          # update pointers
          k_ptrs += BLOCK_N * stride_kn
          v_ptrs += BLOCK_N * stride_vk
      # rematerialize offsets to save registers
      start_m = tl.program_id(0)
      offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
      # write back l and m
      if IS_TRAINING:
          l_ptrs = L + off_hz * N_CTX + offs_m
          m_ptrs = M + off_hz * N_CTX + offs_m
          tl.store(l_ptrs, l_prev)
          tl.store(m_ptrs, m_prev)
      # initialize pointers to output
      offs_n = tl.arange(0, BLOCK_DMODEL)
      off_o = (
          off_b * stride_oz
          + off_h * stride_oh
          + offs_m[:, None] * stride_om
          + offs_n[None, :] * stride_on
      )
      out_ptrs = Out + off_o
      if EVEN_M:
          if EVEN_HEADDIM:
              tl.store(out_ptrs, acc.to(Q.dtype.element_ty))
          else:
              tl.store(
                  out_ptrs,
                  acc.to(Q.dtype.element_ty),
                  mask=offs_n[None, :] < H_DIM,
              )
      else:
          if EVEN_HEADDIM:
              tl.store(
                  out_ptrs,
                  acc.to(Q.dtype.element_ty),
                  mask=offs_m[:, None] < N_CTX,
              )
          else:
              tl.store(
                  out_ptrs,
                  acc.to(Q.dtype.element_ty),
                  mask=(offs_m[:, None] < N_CTX)
                  & (offs_n[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 3 for calling a global function with example usage: _attention_core(Q, K, V, Mask, Bias, sm_scale, L, M, Out, 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_oz, stride_oh, stride_om, stride_on, stride_bz, stride_bh, stride_bm, stride_bn, stride_mz, stride_mh, stride_mm, stride_mn, Z, H, N_CTX, H_DIM, BATCH, inf, IS_TRAINING, BLOCK_M, BLOCK_N, BLOCK_DMODEL, use_mask, use_bias, EVEN_M, EVEN_N, EVEN_HEADDIM) without return types