code: def bench_flash_attention(
      BATCH,
      H,
      N_CTX,
      D_HEAD,
      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,
          )
          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, 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)
          return ms
      if provider == "flash":
          lengths = torch.full(
              (BATCH,), fill_value=N_CTX, device=device
          )
          cu_seqlens = torch.zeros(
              (BATCH + 1,), device=device, dtype=torch.int32
          )
          cu_seqlens[1:] = lengths.cumsum(0)
          qkv = torch.randn(
              (BATCH * N_CTX, 3, H, D_HEAD),
              dtype=dtype,
              device=device,
              requires_grad=True,
          )
          fn = lambda: flash_attn_func(
              qkv, cu_seqlens, 0.0, N_CTX, causal=True
          )
          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)
          return ms
  
file_path: triton/python/test/unit/hopper/test_flashattention.py
imports: import pytest
  import torch
  import triton
  import triton.language
label_desc: Use image node_id 5 for calling a global function with example usage: bench_flash_attention(BATCH, H, N_CTX, D_HEAD, mode, provider, dtype, device) and returns: ms, ms