code: def benchmark():
      ours_fw = []
      ours_bw = []
      baseline_fw = []
      baseline_bw = []
      for batch in range(1, 17):
          Q = torch.randn(
              batch,
              N,
              seq_len,
              heads,
              dim,
              dtype=dtype,
              device="cuda",
              requires_grad=True,
          )
          K = torch.randn(
              batch,
              N,
              seq_len,
              heads,
              dim,
              dtype=dtype,
              device="cuda",
              requires_grad=True,
          )
          V = torch.randn(
              batch,
              N,
              seq_len,
              heads,
              dim,
              dtype=dtype,
              device="cuda",
              requires_grad=True,
          )
          bias1 = torch.randn(
              batch,
              N,
              1,
              1,
              seq_len,
              dtype=dtype,
              device="cuda",
              requires_grad=False,
          )
          bias2 = torch.randn(
              batch,
              1,
              heads,
              seq_len,
              seq_len,
              dtype=dtype,
              device="cuda",
              requires_grad=True,
          )
          # warm up
          DS4Sci_EvoformerAttention(Q, K, V, [bias1, bias2])
          with cuda_timer(ours_fw):
              out = DS4Sci_EvoformerAttention(Q, K, V, [bias1, bias2])
          d_out = torch.rand_like(out)
          with cuda_timer(ours_bw):
              out.backward(d_out)
          # warm up
          attention_reference(Q, K, V, [bias1, bias2], 1 / (dim**0.5))
          with cuda_timer(baseline_fw):
              ref_out = attention_reference(
                  Q, K, V, [bias1, bias2], 1 / (dim**0.5)
              )
          with cuda_timer(baseline_bw):
              ref_out.backward(d_out)
  
      print(
          f"batch size\tours (FW)\tbaseline (FW)\tours (BW)\tbaseline (BW)"
      )
      for i in range(len(ours_fw)):
          print(
              f"{i+1}\t{ours_fw[i]}\t{baseline_fw[i]}\t{ours_bw[i]}\t{baseline_bw[i]}"
          )
  
file_path: deepspeed/tests/benchmarks/DS4Sci_EvoformerAttention_bench.py
imports: import contextlib
  import torch
  from typing import List
  from torch.nn import functional
  from deepspeed.ops.deepspeed4science import DS4Sci_EvoformerAttention
  from deepspeed.accelerator import get_accelerator
label_desc: Use image node_id 3 for calling a global function with example usage: benchmark() without return types