code: def test_openfold_triton_mha(
      self, Z=256, H=4, N_CTX=256, D_HEAD=32, dtype=torch.float16
  ):
      One = 1
      q = (
          torch.empty(
              (One, Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda"
          )
          .normal_(mean=0.1, std=0.2)
          .requires_grad_()
      )
      k = (
          torch.empty(
              (One, Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda"
          )
          .normal_(mean=0.4, std=0.2)
          .requires_grad_()
      )
      v = (
          torch.empty(
              (One, Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda"
          )
          .normal_(mean=0.3, std=0.2)
          .requires_grad_()
      )
      bias = (
          torch.empty(
              (One, One, H, N_CTX, N_CTX), dtype=dtype, device="cuda"
          )
          .normal_(mean=0.2, std=0.2)
          .requires_grad_()
      )
      mask = (
          torch.empty(
              (One, N_CTX, One, One, N_CTX), device="cuda"
          ).normal_(mean=0, std=0.5)
          > 0
      )
      mask = mask.to(
          device=torch.device("cuda"), dtype=dtype
      ).requires_grad_(False)
  
      dout = torch.randn_like(q)
      inf = 1e9
  
      # reference implementation
      ref_out = openfold_attention_eager(q, k, v, mask, bias, inf)
      ref_out.backward(dout)
  
      ref_dv, v.grad = v.grad.clone(), None
      ref_dk, k.grad = k.grad.clone(), None
      ref_dq, q.grad = q.grad.clone(), None
      ref_dbias, bias.grad = bias.grad.clone(), None
  
      # triton implementation
      tri_out = openfold_attention_triton(
          q, k, v, mask, bias, inf, torch.is_grad_enabled()
      )
      tri_out.backward(dout)
  
      tri_dv, v.grad = v.grad.clone(), None
      tri_dk, k.grad = k.grad.clone(), None
      tri_dq, q.grad = q.grad.clone(), None
      tri_dbias, bias.grad = bias.grad.clone(), None
  
      # check results
      torch.testing.assert_close(ref_out, tri_out, atol=1e-2, rtol=0)
      torch.testing.assert_close(ref_dv, tri_dv, atol=1e-2, rtol=0)
      torch.testing.assert_close(ref_dk, tri_dk, atol=1e-2, rtol=0)
      torch.testing.assert_close(ref_dq, tri_dq, atol=1e-2, rtol=0)
      torch.testing.assert_close(
          ref_dbias, tri_dbias, atol=1e-2, rtol=0
      )
  
file_path: apex/apex/contrib/test/openfold_triton/test_openfold_mha.py
imports: import math
  import random
  from typing import Optional
  import torch
  import unittest
label_desc: Use image node_id 2 for calling the OpenfoldMhaTest obj's underlying member method code with example usage: obj.test_openfold_triton_mha(Z, H, N_CTX, D_HEAD, dtype) without return types