code: def attention_ref(q, k, v, attn_mask=None, causal=False):
      """
      attention output of the control group
      """
      dtype_og = q.dtype
      seqlen_q, seqlen_k = q.shape[1], k.shape[1]
      d = q.shape[-1]
      scale = 1.0 / math.sqrt(d)
      scores = torch.einsum("bthd,bshd->bhts", q * scale, k)
  
      if attn_mask is not None:
          scores.masked_fill_(
              rearrange(~attn_mask, "b s -> b 1 1 s"), float("-inf")
          )
      if causal:
          causal_mask = torch.triu(
              torch.ones(
                  seqlen_q, seqlen_k, dtype=torch.bool, device=q.device
              ),
              1,
          )
          scores.masked_fill_(causal_mask, float("-inf"))
      attention = torch.softmax(scores, dim=-1)
  
      output = torch.einsum("bhts,bshd->bthd", attention, v)
      output = rearrange(output, "b s h d -> b s (h d)")
  
      # Modify the data at the positions of the mask to 0
      if attn_mask is not None:
          output.masked_fill_(
              rearrange(~attn_mask, "b s -> b s 1"), 0.0
          )
  
      return output.to(dtype=dtype_og)
  
file_path: ColossalAI/tests/test_utils/test_flash_attention.py
imports: import math
  import pytest
  import torch
  from einops import rearrange
  from colossalai.kernel.cuda_native.mha.flash_attn_2 import HAS_FLASH_ATTN
  from colossalai.kernel.cuda_native.mha.mem_eff_attn import HAS_MEM_EFF_ATTN
  from colossalai.testing import clear_cache_before_run, parameterize
label_desc: Use image node_id 1 for calling a global function with example usage: attention_ref(q, k, v, attn_mask, causal) and returns: output