code: def test_flash_attention(
      Z, H, N_CTX, D_HEAD, seq_par, causal, mode, dtype_str
  ):
      stream = torch.cuda.Stream()
      torch.cuda.set_stream(stream)
      is_backward = mode == "backward"
      capability = torch.cuda.get_device_capability()
      if capability[0] < 8:
          pytest.skip(
              "Flash attention only supported for compute capability < 80"
          )
      torch.manual_seed(20)
      dtype = {
          "float16": torch.float16,
          "bfloat16": torch.bfloat16,
          "float32": torch.float32,
      }[dtype_str]
      # init data
      if dtype_str == "float32":
          N_CTX = 1024
          D_HEAD = 16
      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
      # benchmark
      fn = lambda: triton.ops.attention(
          q, k, v, causal, sm_scale, seq_par
      )
      if is_backward:
          o = fn()
          do = torch.randn_like(o)
          fn = lambda: o.backward(do, retain_graph=True)
      ms = triton.testing.do_bench_cudagraph(fn)
      # compute flops
      flops_per_matmul = 2.0 * Z * H * N_CTX * N_CTX * D_HEAD * 0.5
      total_flops = 2 * flops_per_matmul
      if is_backward:
          total_flops *= 2.5  # 2.0(bwd) + 0.5(recompute)
      cur_gpu_perf = total_flops / ms * 1e-9
      # maximum flops
      cur_sm_clock = nvsmi(["clocks.current.sm"])[0]
      max_gpu_perf = get_max_tensorcore_tflops(
          dtype, clock_rate=cur_sm_clock * 1e3
      )
      cur_gpu_util = cur_gpu_perf / max_gpu_perf
      ref_gpu_util = flash_attention_data[DEVICE_NAME][
          (Z, H, N_CTX, D_HEAD, seq_par, causal, mode, dtype_str)
      ]
      print_perf(ms, cur_gpu_util, ref_gpu_util)
      triton.testing.assert_close(
          cur_gpu_util, ref_gpu_util, atol=0.02, rtol=0.01
      )
  
file_path: triton/python/test/regression/test_performance.py
imports: import pytest
  import torch
  import triton
  import triton.language
  import triton.ops
  from triton.testing import get_dram_gbps, get_max_tensorcore_tflops, nvsmi
label_desc: Use image node_id 5 for calling a global function with example usage: test_flash_attention(Z, H, N_CTX, D_HEAD, seq_par, causal, mode, dtype_str) without return types