code: def bench_flash_attention(
      BATCH,
      H,
      N_CTX,
      D_HEAD,
      causal,
      mode,
      provider,
      dtype=torch.float16,
      device="cuda",
  ):
      assert mode in ["fwd", "bwd"]
      warmup = 25
      rep = 100
      if provider == "triton":
          q = torch.randn(
              (BATCH, H, N_CTX, D_HEAD),
              dtype=dtype,
              device="cuda",
              requires_grad=True,
          )
          k = torch.randn(
              (BATCH, H, N_CTX, D_HEAD),
              dtype=dtype,
              device="cuda",
              requires_grad=True,
          )
          if mode == "fwd" and TORCH_HAS_FP8:
              q = q.to(torch.float8_e5m2)
              k = k.to(torch.float8_e5m2)
          v = torch.randn(
              (BATCH, H, N_CTX, D_HEAD),
              dtype=dtype,
              device="cuda",
              requires_grad=True,
          )
          sm_scale = 1.3
          fn = lambda: attention(q, k, v, causal, sm_scale)
          if mode == "bwd":
              o = fn()
              do = torch.randn_like(o)
              fn = lambda: o.backward(do, retain_graph=True)
          ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
      if provider == "flash":
          qkv = torch.randn(
              (BATCH, N_CTX, 3, H, D_HEAD),
              dtype=dtype,
              device=device,
              requires_grad=True,
          )
          fn = lambda: flash_attn_func(qkv, causal=causal)
          if mode == "bwd":
              o = fn()
              do = torch.randn_like(o)
              fn = lambda: o.backward(do, retain_graph=True)
          ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
      flops_per_matmul = 2.0 * BATCH * H * N_CTX * N_CTX * D_HEAD
      total_flops = 2 * flops_per_matmul
      if causal:
          total_flops *= 0.5
      if mode == "bwd":
          total_flops *= 2.5  # 2.0(bwd) + 0.5(recompute)
      return total_flops / ms * 1e-9
  
file_path: triton/python/tutorials/06-fused-attention.py
imports: import pytest
  import torch
  import triton
  import triton.language
label_desc: Use image node_id 8 for calling a global function with example usage: bench_flash_attention(BATCH, H, N_CTX, D_HEAD, causal, mode, provider, dtype, device) and returns: unknown