code: def test_op(Z, H, N_CTX, D_HEAD, dtype=torch.float16):
      torch.manual_seed(20)
      q = (
          torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda")
          .normal_(mean=0.1, std=0.2)
          .requires_grad_()
      )
      k = (
          torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda")
          .normal_(mean=0.4, std=0.2)
          .requires_grad_()
      )
      v = (
          torch.empty((Z, H, N_CTX, D_HEAD), dtype=dtype, device="cuda")
          .normal_(mean=0.3, std=0.2)
          .requires_grad_()
      )
      sm_scale = 0.2
      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
      for z in range(Z):
          for h in range(H):
              p[:, :, M == 0] = float("-inf")
      p = torch.softmax(p.float(), dim=-1).half()
      # 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 = attention(q, k, v, sm_scale)
      # print(ref_out)
      # print(tri_out)
      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
      assert torch.allclose(ref_out, tri_out, atol=1e-2, rtol=0)
      assert torch.allclose(ref_dv, tri_dv, atol=1e-2, rtol=0)
      assert torch.allclose(ref_dk, tri_dk, atol=1e-2, rtol=0)
      assert torch.allclose(ref_dq, tri_dq, atol=1e-2, rtol=0)
  
file_path: triton/python/test/unit/hopper/test_persistent_warp_specialized_fused-attention.py
imports: import pytest
  import torch
  import triton
  import triton.language
label_desc: Use image node_id 4 for calling a global function with example usage: test_op(Z, H, N_CTX, D_HEAD, dtype) without return types