code: def attention_pytorch(qkv, dropout_p=0.0, causal=True):
      """
      Arguments:
          qkv: (batch_size, seqlen, 3, nheads, head_dim)
          dropout_p: float
      Output:
          output: (batch_size, seqlen, nheads, head_dim)
      """
      batch_size, seqlen, _, nheads, d = qkv.shape
      q, k, v = qkv.unbind(dim=2)
      q = rearrange(q, "b t h d -> (b h) t d")
      k = rearrange(k, "b s h d -> (b h) d s")
      softmax_scale = 1.0 / math.sqrt(d)
      # Preallocate attn_weights for `baddbmm`
      scores = torch.empty(
          batch_size * nheads,
          seqlen,
          seqlen,
          dtype=qkv.dtype,
          device=qkv.device,
      )
      scores = rearrange(
          torch.baddbmm(scores, q, k, beta=0, alpha=softmax_scale),
          "(b h) t s -> b h t s",
          h=nheads,
      )
      if causal:
          # "triu_tril_cuda_template" not implemented for 'BFloat16'
          # So we have to construct the mask in float
          causal_mask = torch.triu(
              torch.full(
                  (seqlen, seqlen), -10000.0, device=scores.device
              ),
              1,
          )
          # TD [2022-09-30]: Adding is faster than masked_fill_ (idk why, just better kernel I guess)
          scores = scores + causal_mask.to(dtype=scores.dtype)
      attention = torch.softmax(scores, dim=-1)
      attention_drop = F.dropout(attention, dropout_p)
      output = torch.einsum("bhts,bshd->bthd", attention_drop, v)
      return output.to(dtype=qkv.dtype)
  
file_path: flash-attention/benchmarks/benchmark_causal.py
imports: from functools import partial
  import math
  import torch
  import torch.nn
  import torch.nn.functional
  from einops import rearrange, repeat
  from flash_attn.utils.benchmark import benchmark_forward, benchmark_backward, benchmark_combined, benchmark_all, benchmark_fwd_bwd, pytorch_profiler
  from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func
  from flash_attn import flash_attn_qkvpacked_func, flash_attn_kvpacked_func
  import pickle
label_desc: Use image node_id 1 for calling a global function with example usage: attention_pytorch(qkv, dropout_p, causal) and returns: output