code: def benchmark_transformer_blocks():
      device = "cuda" if torch.cuda.is_available() else "cpu"
      import torch.utils.benchmark as benchmark
  
      def benchmark_torch_function_in_microseconds(f, *args, **kwargs):
          t0 = benchmark.Timer(
              stmt="f(*args, **kwargs)",
              globals={"args": args, "kwargs": kwargs, "f": f},
          )
          return t0.blocked_autorange().mean * 1e6
  
      checkpoint = True
      compile = False
  
      batch_size = 32
      h, w = 64, 64
      context_len = 77
      embed_dimension = 1024
      context_dim = 1024
      d_head = 64
  
      transformer_depth = 4
  
      n_heads = embed_dimension // d_head
  
      dtype = torch.float16
  
      model_native = SpatialTransformer(
          embed_dimension,
          n_heads,
          d_head,
          context_dim=context_dim,
          use_linear=True,
          use_checkpoint=checkpoint,
          attn_type="softmax",
          depth=transformer_depth,
          sdp_backend=SDPBackend.FLASH_ATTENTION,
      ).to(device)
      model_efficient_attn = SpatialTransformer(
          embed_dimension,
          n_heads,
          d_head,
          context_dim=context_dim,
          use_linear=True,
          depth=transformer_depth,
          use_checkpoint=checkpoint,
          attn_type="softmax-xformers",
      ).to(device)
      if not checkpoint and compile:
          print("compiling models")
          model_native = torch.compile(model_native)
          model_efficient_attn = torch.compile(model_efficient_attn)
  
      x = torch.rand(
          batch_size, embed_dimension, h, w, device=device, dtype=dtype
      )
      c = torch.rand(
          batch_size,
          context_len,
          context_dim,
          device=device,
          dtype=dtype,
      )
  
      from torch.profiler import (
          ProfilerActivity,
          profile,
          record_function,
      )
  
      activities = [ProfilerActivity.CPU, ProfilerActivity.CUDA]
  
      with torch.autocast("cuda"):
          print(
              f"The native model runs in {benchmark_torch_function_in_microseconds(model_native.forward, x, c):.3f} microseconds"
          )
          print(
              f"The efficientattn model runs in {benchmark_torch_function_in_microseconds(model_efficient_attn.forward, x, c):.3f} microseconds"
          )
  
          print(75 * "+")
          print("NATIVE")
          print(75 * "+")
          torch.cuda.reset_peak_memory_stats()
          with profile(
              activities=activities,
              record_shapes=False,
              profile_memory=True,
          ) as prof:
              with record_function("NativeAttention stats"):
                  for _ in range(25):
                      model_native(x, c)
          print(
              prof.key_averages().table(
                  sort_by="cuda_time_total", row_limit=10
              )
          )
          print(
              torch.cuda.max_memory_allocated() * 1e-9,
              "GB used by native block",
          )
  
          print(75 * "+")
          print("Xformers")
          print(75 * "+")
          torch.cuda.reset_peak_memory_stats()
          with profile(
              activities=activities,
              record_shapes=False,
              profile_memory=True,
          ) as prof:
              with record_function("xformers stats"):
                  for _ in range(25):
                      model_efficient_attn(x, c)
          print(
              prof.key_averages().table(
                  sort_by="cuda_time_total", row_limit=10
              )
          )
          print(
              torch.cuda.max_memory_allocated() * 1e-9,
              "GB used by xformers block",
          )
  
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 3 for calling a global function with example usage: benchmark_transformer_blocks() and returns: unknown