code: def attention_ref(
      q,
      k,
      v,
      query_padding_mask=None,
      key_padding_mask=None,
      attn_bias=None,
      dropout_p=0.0,
      dropout_mask=None,
      causal=False,
      window_size=(-1, -1),  # -1 means infinite window size
      upcast=True,
      reorder_ops=False,
  ):
      """
      Arguments:
          q: (batch_size, seqlen_q, nheads, head_dim)
          k: (batch_size, seqlen_k, nheads_k, head_dim)
          v: (batch_size, seqlen_k, nheads_k, head_dim)
          query_padding_mask: (batch_size, seqlen_q)
          key_padding_mask: (batch_size, seqlen_k)
          attn_bias: broadcastable to (batch_size, nheads, seqlen_q, seqlen_k)
          dropout_p: float
          dropout_mask: (batch_size, nheads, seqlen_q, seqlen_k)
          causal: whether to apply causal masking
          window_size: (int, int), left and right window size
          upcast: whether to cast all inputs to fp32, do all computation in fp32, then cast
              output back to fp16/bf16.
          reorder_ops: whether to change the order of operations (scaling k instead of scaling k, etc.)
              without changing the math. This is to estimate the numerical error from operation
              reordering.
      Output:
          output: (batch_size, seqlen_q, nheads, head_dim)
          attention: (batch_size, nheads, seqlen_q, seqlen_k), softmax after dropout
      """
      if causal:
          window_size = (window_size[0], 0)
      dtype_og = q.dtype
      if upcast:
          q, k, v = q.float(), k.float(), v.float()
      seqlen_q, seqlen_k = q.shape[1], k.shape[1]
      k = repeat(
          k, "b s h d -> b s (h g) d", g=q.shape[2] // k.shape[2]
      )
      v = repeat(
          v, "b s h d -> b s (h g) d", g=q.shape[2] // v.shape[2]
      )
      d = q.shape[-1]
      if not reorder_ops:
          scores = torch.einsum("bthd,bshd->bhts", q / math.sqrt(d), k)
      else:
          scores = torch.einsum("bthd,bshd->bhts", q, k / math.sqrt(d))
      if key_padding_mask is not None:
          scores.masked_fill_(
              rearrange(~key_padding_mask, "b s -> b 1 1 s"),
              float("-inf"),
          )
      if window_size[0] >= 0 or window_size[1] >= 0:
          local_mask = construct_local_mask(
              seqlen_q,
              seqlen_k,
              window_size,
              query_padding_mask,
              key_padding_mask,
              q.device,
          )
          scores.masked_fill_(local_mask, float("-inf"))
      if attn_bias is not None:
          scores = scores + attn_bias
      attention = torch.softmax(scores, dim=-1).to(v.dtype)
      # Some rows might be completely masked out so we fill them with zero instead of NaN
      if window_size[0] >= 0 or window_size[1] >= 0:
          attention = attention.masked_fill(
              torch.all(local_mask, dim=-1, keepdim=True), 0.0
          )
      # We want to mask here so that the attention matrix doesn't have any NaNs
      # Otherwise we'll get NaN in dV
      if query_padding_mask is not None:
          attention = attention.masked_fill(
              rearrange(~query_padding_mask, "b s -> b 1 s 1"), 0.0
          )
      dropout_scaling = 1.0 / (1 - dropout_p)
      # attention_drop = attention.masked_fill(~dropout_mask, 0.0) * dropout_scaling
      # output = torch.einsum('bhts,bshd->bthd', attention_drop , v)
      if dropout_mask is not None:
          attention_drop = attention.masked_fill(~dropout_mask, 0.0)
      else:
          attention_drop = attention
      output = torch.einsum(
          "bhts,bshd->bthd", attention_drop, v * dropout_scaling
      )
      if query_padding_mask is not None:
          output.masked_fill_(
              rearrange(~query_padding_mask, "b s -> b s 1 1"), 0.0
          )
      return output.to(dtype=dtype_og), attention.to(dtype=dtype_og)
  
imports: import math
  import pytest
  import torch
  import torch.nn.functional
  from einops import rearrange, repeat
  from flash_attn import flash_attn_func, flash_attn_kvpacked_func, flash_attn_qkvpacked_func, flash_attn_varlen_func, flash_attn_varlen_kvpacked_func, flash_attn_varlen_qkvpacked_func, flash_attn_with_kvcache
  from flash_attn.bert_padding import pad_input, unpad_input
  from flash_attn.flash_attn_interface import _get_block_size_n
  from flash_attn.layers.rotary import apply_rotary_emb
label_desc: Use image node_id 5 for calling a global function with example usage: attention_ref(q, k, v, query_padding_mask, key_padding_mask, attn_bias, dropout_p, dropout_mask, causal, window_size, upcast, reorder_ops) and returns: output,  attention