code: def benchmark_attn():
      # Lets define a helpful benchmarking function:
      # https://pytorch.org/tutorials/intermediate/scaled_dot_product_attention_tutorial.html
      device = "cuda" if torch.cuda.is_available() else "cpu"
  
      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
  
      # Lets define the hyper-parameters of our input
      batch_size = 32
      max_sequence_len = 1024
      num_heads = 32
      embed_dimension = 32
  
      dtype = torch.float16
  
      query = torch.rand(
          batch_size,
          num_heads,
          max_sequence_len,
          embed_dimension,
          device=device,
          dtype=dtype,
      )
      key = torch.rand(
          batch_size,
          num_heads,
          max_sequence_len,
          embed_dimension,
          device=device,
          dtype=dtype,
      )
      value = torch.rand(
          batch_size,
          num_heads,
          max_sequence_len,
          embed_dimension,
          device=device,
          dtype=dtype,
      )
  
      print(f"q/k/v shape:", query.shape, key.shape, value.shape)
  
      # Lets explore the speed of each of the 3 implementations
      from torch.backends.cuda import SDPBackend, sdp_kernel
  
      # Helpful arguments mapper
      backend_map = {
          SDPBackend.MATH: {
              "enable_math": True,
              "enable_flash": False,
              "enable_mem_efficient": False,
          },
          SDPBackend.FLASH_ATTENTION: {
              "enable_math": False,
              "enable_flash": True,
              "enable_mem_efficient": False,
          },
          SDPBackend.EFFICIENT_ATTENTION: {
              "enable_math": False,
              "enable_flash": False,
              "enable_mem_efficient": True,
          },
      }
  
      from torch.profiler import (
          ProfilerActivity,
          profile,
          record_function,
      )
  
      activities = [ProfilerActivity.CPU, ProfilerActivity.CUDA]
  
      print(
          f"The default implementation runs in {benchmark_torch_function_in_microseconds(F.scaled_dot_product_attention, query, key, value):.3f} microseconds"
      )
      with profile(
          activities=activities,
          record_shapes=False,
          profile_memory=True,
      ) as prof:
          with record_function("Default detailed stats"):
              for _ in range(25):
                  o = F.scaled_dot_product_attention(query, key, value)
      print(
          prof.key_averages().table(
              sort_by="cuda_time_total", row_limit=10
          )
      )
  
      print(
          f"The math implementation runs in {benchmark_torch_function_in_microseconds(F.scaled_dot_product_attention, query, key, value):.3f} microseconds"
      )
      with sdp_kernel(**backend_map[SDPBackend.MATH]):
          with profile(
              activities=activities,
              record_shapes=False,
              profile_memory=True,
          ) as prof:
              with record_function("Math implmentation stats"):
                  for _ in range(25):
                      o = F.scaled_dot_product_attention(
                          query, key, value
                      )
          print(
              prof.key_averages().table(
                  sort_by="cuda_time_total", row_limit=10
              )
          )
  
      with sdp_kernel(**backend_map[SDPBackend.FLASH_ATTENTION]):
          try:
              print(
                  f"The flash attention implementation runs in {benchmark_torch_function_in_microseconds(F.scaled_dot_product_attention, query, key, value):.3f} microseconds"
              )
          except RuntimeError:
              print(
                  "FlashAttention is not supported. See warnings for reasons."
              )
          with profile(
              activities=activities,
              record_shapes=False,
              profile_memory=True,
          ) as prof:
              with record_function("FlashAttention stats"):
                  for _ in range(25):
                      o = F.scaled_dot_product_attention(
                          query, key, value
                      )
          print(
              prof.key_averages().table(
                  sort_by="cuda_time_total", row_limit=10
              )
          )
  
      with sdp_kernel(**backend_map[SDPBackend.EFFICIENT_ATTENTION]):
          try:
              print(
                  f"The memory efficient implementation runs in {benchmark_torch_function_in_microseconds(F.scaled_dot_product_attention, query, key, value):.3f} microseconds"
              )
          except RuntimeError:
              print(
                  "EfficientAttention is not supported. See warnings for reasons."
              )
          with profile(
              activities=activities,
              record_shapes=False,
              profile_memory=True,
          ) as prof:
              with record_function("EfficientAttention stats"):
                  for _ in range(25):
                      o = F.scaled_dot_product_attention(
                          query, key, value
                      )
          print(
              prof.key_averages().table(
                  sort_by="cuda_time_total", row_limit=10
              )
          )
  
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 1 for calling a global function with example usage: benchmark_attn() and returns: unknown