code: def bench_flash_attention(
      BATCH,
      H,
      N_CTX,
      D_HEAD,
      mode,
      casual,
      seq_par,
      provider,
      dtype=torch.float16,
      device="cuda",
  ):
      assert mode in ["fwd", "bwd"]
      warmup = 25
      rep = 100
      sm_scale = 1.3
      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,
      )
      if provider == "triton":
          fn = lambda: triton.ops.attention(
              q, k, v, casual, sm_scale, seq_par
          )
          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)
          fn = lambda: flash_attn_func(
              q,
              k,
              v,
              dropout_p=0.0,
              softmax_scale=sm_scale,
              causal=casual,
          )
          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
  
imports: import pytest
  import torch
  import triton
  import triton.ops
label_desc: Use image node_id 2 for calling a global function with example usage: bench_flash_attention(BATCH, H, N_CTX, D_HEAD, mode, casual, seq_par, provider, dtype, device) and returns: ms, ms