code: def test_op(Z, H, N_CTX, D_HEAD, dtype, causal, seq_par):
      import os
  
      enable_tma = os.environ.get("ENABLE_TMA", "not found").lower()
      if enable_tma in ["on", "true", "1"]:
          if dtype == torch.bfloat16:
              pytest.skip("bfloat16 tma not support currently")
  
      capability = torch.cuda.get_device_capability()
      interpreter = os.environ.get("TRITON_INTERPRET", "not found") in [
          "on",
          "true",
          "1",
      ]
      if not interpreter and capability[0] < 8:
          pytest.skip(
              "Flash attention only supported for compute capability >= 80"
          )
      torch.manual_seed(20)
      q = (
          torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda")
          .normal_(mean=0.0, std=0.5)
          .requires_grad_()
      )
      k = (
          torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda")
          .normal_(mean=0.0, std=0.5)
          .requires_grad_()
      )
      v = (
          torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda")
          .normal_(mean=0.0, std=0.5)
          .requires_grad_()
      )
      sm_scale = 0.5
      dout = torch.randn_like(q)
      # reference implementation
      M = torch.tril(torch.ones((N_CTX, N_CTX), device="cuda"))
      p = torch.matmul(q, k.transpose(2, 3)) * sm_scale
      if causal:
          p[:, :, M == 0] = float("-inf")
      p = torch.softmax(p.float(), dim=-1).to(dtype)
      # p = torch.exp(p)
      ref_out = torch.matmul(p, v)
      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
      # # triton implementation
      tri_out = triton.ops.attention(q, k, v, causal, sm_scale, seq_par)
      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
      # compare
      atol = 1e-1 if dtype == torch.bfloat16 else 1e-2
      torch.testing.assert_close(
          torch.nn.functional.normalize(torch.flatten(ref_out), dim=0),
          torch.nn.functional.normalize(torch.flatten(tri_out), dim=0),
          atol=atol,
          rtol=0,
      )
      torch.testing.assert_close(
          torch.nn.functional.normalize(torch.flatten(ref_dv), dim=0),
          torch.nn.functional.normalize(torch.flatten(tri_dv), dim=0),
          atol=atol,
          rtol=0,
      )
      torch.testing.assert_close(
          torch.nn.functional.normalize(torch.flatten(ref_dk), dim=0),
          torch.nn.functional.normalize(torch.flatten(tri_dk), dim=0),
          atol=atol,
          rtol=0,
      )
      torch.testing.assert_close(
          torch.nn.functional.normalize(torch.flatten(ref_dq), dim=0),
          torch.nn.functional.normalize(torch.flatten(tri_dq), dim=0),
          atol=atol,
          rtol=0,
      )
  
imports: import pytest
  import torch
  import triton
  import triton.ops
label_desc: Use image node_id 1 for calling a global function with example usage: test_op(Z, H, N_CTX, D_HEAD, dtype, causal, seq_par) without return types