code: def test_attention_gpt(proj_shape, dtype, dropout):
      (B, S, H, D_HEAD) = proj_shape
      D = H * D_HEAD
  
      q = torch.randn(
          (B, S, H, D_HEAD),
          dtype=dtype,
          device="cuda",
          requires_grad=True,
      )
      k = torch.randn(
          (B, S, H, D_HEAD),
          dtype=dtype,
          device="cuda",
          requires_grad=True,
      )
      v = torch.randn(
          (B, S, H, D_HEAD),
          dtype=dtype,
          device="cuda",
          requires_grad=True,
      )
  
      mask = [
          torch.ones(S - i, dtype=torch.bool, device="cuda")
          for i in range(B)
      ]
      mask = torch.nn.utils.rnn.pad_sequence(mask, batch_first=True)
  
      attn = ColoAttention(D, H, dropout=dropout)
      y = attn(
          q,
          k,
          v,
          attn_mask=mask,
          attn_mask_type=AttnMaskType.paddedcausal,
      )
  
      assert list(y.shape) == [B, S, D]
  
      out_ref = attention_ref(q, k, v, mask, causal=True)
  
      # check gradients
      dy = torch.rand_like(y)
      grad_q, grad_k, grad_v = torch.autograd.grad(y, (q, k, v), dy)
      grad_ref_q, grad_ref_k, grad_ref_v = torch.autograd.grad(
          out_ref, (q, k, v), dy
      )
  
      torch.allclose(
          y, out_ref, atol=1e-7
      ), f"{(y - out_ref).abs().max()}"
      torch.allclose(
          grad_q, grad_ref_q, atol=1e-7
      ), f"{(grad_q - grad_ref_q).abs().max()}"
      torch.allclose(
          grad_k, grad_ref_k, atol=1e-7
      ), f"{(grad_k - grad_ref_k).abs().max()}"
      torch.allclose(
          grad_v, grad_ref_v, atol=1e-7
      ), f"{(grad_v - grad_ref_v).abs().max()}"
  
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 2 for calling a global function with example usage: test_attention_gpt(proj_shape, dtype, dropout) without return types