code: def test02():
      # try cosine flash attention
      import time
  
      torch.backends.cuda.matmul.allow_tf32 = True
      torch.backends.cudnn.allow_tf32 = True
      torch.backends.cudnn.benchmark = True
      print("testing cosine flash attention...")
      DIM = 1024
      SEQLEN = 4096
      BS = 16
  
      print(" softmax (vanilla) first...")
      model = BasicTransformerBlock(
          dim=DIM,
          n_heads=16,
          d_head=64,
          dropout=0.0,
          context_dim=None,
          attn_mode="softmax",
      ).cuda()
      try:
          x = torch.randn(BS, SEQLEN, DIM).cuda()
          tic = time.time()
          y = model(x)
          toc = time.time()
          print(y.shape, toc - tic)
      except RuntimeError as e:
          # likely oom
          print(str(e))
  
      print("\n now flash-cosine...")
      model = BasicTransformerBlock(
          dim=DIM,
          n_heads=16,
          d_head=64,
          dropout=0.0,
          context_dim=None,
          attn_mode="flash-cosine",
      ).cuda()
      x = torch.randn(BS, SEQLEN, DIM).cuda()
      tic = time.time()
      y = model(x)
      toc = time.time()
      print(y.shape, toc - tic)
      print("done with test02.\n")
  
file_path: generative-models/scripts/tests/attention.py
imports: import einops
  import torch
  import torch.nn.functional
  import torch.utils.benchmark
  from torch.backends.cuda import SDPBackend
  from sgm.modules.attention import BasicTransformerBlock, SpatialTransformer
label_desc: Use image node_id 5 for calling a global function with example usage: test02() without return types