code: def forward(
      ctx,
      q,
      k,
      v,
      mask=None,
      bias=None,
      inf=1000000000.0,
      is_training=True,
  ):
      q_ori_size = len(q.size())
      if q_ori_size == 5:
          q = rearrange(q, "1 b2 h n d -> (1 b2) h n d")
          k = rearrange(k, "1 b2 h n d -> (1 b2) h n d")
          v = rearrange(v, "1 b2 h n d -> (1 b2) h n d")
      if bias is not None:
          if len(bias.size()) == 5:
              bias = rearrange(bias, "1 b2 h n d -> (1 b2) h n d")
  
      if mask is not None and len(mask.size()) == 5:
          mask = rearrange(mask, "1 b 1 1 e -> b 1 1 e")
  
      batch = 1
      sm_scale = 1.0 / math.sqrt(q.size(-1))
      # q *= sm_scale
      # shape constraints
      Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
      assert Lq == Lk and Lk == Lv
  
      if not is_training:
          Lk = max(triton.next_power_of_2(Lk), 16)
  
      assert Lk in {16, 32, 64, 128}
  
      o = torch.empty_like(q)
  
      Z, H, N_CTX, H_DIM = q.shape
      grid = lambda META: (triton.cdiv(N_CTX, META["BLOCK_M"]), Z * H)
      l = torch.empty(
          (q.shape[-4], q.shape[-3], q.shape[-2]),
          device=q.device,
          dtype=torch.float32,
      )
      m = torch.empty(
          (q.shape[-4], q.shape[-3], q.shape[-2]),
          device=q.device,
          dtype=torch.float32,
      )
      # BLOCK_M, BLOCK_N, num_warps, num_stages  = 64, 64, 2, 3
      BLOCK_M, BLOCK_N, num_warps, num_stages = schedule_triton_mha(
          list(q.shape), fwd=True
      )
      if bias != None:
          bias = bias.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)
      )
      if mask != None:
          mask = mask.expand(-1, q.shape[1], q.shape[2], -1)
      mask_strides = (
          (
              mask.stride(0),
              mask.stride(1),
              mask.stride(2),
              mask.stride(3),
          )
          if mask is not None
          else (0, 0, 0, 0)
      )
  
      _attention_core[grid](
          q,
          k,
          v,
          mask,
          bias,
          sm_scale,
          l,
          m,
          o,
          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),
          o.stride(0),
          o.stride(1),
          o.stride(2),
          o.stride(3),
          *bias_strides,
          *mask_strides,
          q.shape[0],
          q.shape[1],
          q.shape[2],
          q.shape[3],
          batch,  # 256 8 128 1
          inf=inf,
          IS_TRAINING=is_training,
          BLOCK_M=BLOCK_M,
          BLOCK_N=BLOCK_N,
          BLOCK_DMODEL=Lk,
          use_mask=(mask != None),
          use_bias=(bias != None),
          num_warps=num_warps,
          num_stages=num_stages,
      )
      o = o.contiguous()
      # print(h.asm["ttgir"])
      if is_training:
          ctx.save_for_backward(q, k, v, o, m, l, bias)
          ctx.grid = grid
          ctx.sm_scale = sm_scale
          ctx.BLOCK_DMODEL = Lk
          ctx.mask = mask
          ctx.inf = inf
      if q_ori_size == 5:
          o = rearrange(o, "a b c d -> 1 a b c d")
      return o
  
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 1 for calling the FusedAttenionCoreFunc obj's underlying member method code with example usage: obj.forward(ctx, q, k, v, mask, bias, inf, is_training) and returns: o