code: def _attention(Q, K, V, bias1, bias2):
      assert Q.shape[-3] > 16, "seq_len must be greater than 16"
      O = torch.empty_like(Q, dtype=Q.dtype)
      assert get_accelerator().on_accelerator(Q), "Q must be on cuda"
      assert get_accelerator().on_accelerator(K), "K must be on cuda"
      assert get_accelerator().on_accelerator(V), "V must be on cuda"
      assert get_accelerator().on_accelerator(
          bias1
      ), "bias1 must be on cuda"
      assert get_accelerator().on_accelerator(
          bias2
      ), "bias2 must be on cuda"
      global kernel_
      if kernel_ is None:
          kernel_ = EvoformerAttnBuilder().load()
      nheads = Q.shape[-2]
      nq = (Q.shape[-3] + 31) // 32 * 32
      nb = np.prod(Q.shape[:-3])
      lse = torch.empty(
          (nb, nheads, nq), dtype=torch.float32, device=Q.device
      )
      kernel_.attention(Q, K, V, bias1, bias2, O, lse)
      return O, lse
  
file_path: DeepSpeed/deepspeed/ops/deepspeed4science/evoformer_attn.py
imports: import torch
  import numpy
  from deepspeed.ops.op_builder import EvoformerAttnBuilder
  from deepspeed.accelerator import get_accelerator
label_desc: Use image node_id 1 for calling a global function with example usage: _attention(Q, K, V, bias1, bias2) and returns: O,  lse