code: def __init__(
      self,
      embed_dim,
      num_heads,
      batch_first=True,
      attention_dropout=0.0,
      causal=False,
      device=None,
      dtype=None,
      **kwargs,
  ) -> None:
      from einops import rearrange
      from flash_attn.flash_attention import FlashAttention
  
      assert batch_first
      factory_kwargs = {"device": device, "dtype": dtype}
      super().__init__()
      self.embed_dim = embed_dim
      self.num_heads = num_heads
      self.causal = causal
  
      assert (
          self.embed_dim % num_heads == 0
      ), "self.kdim must be divisible by num_heads"
      self.head_dim = self.embed_dim // num_heads
      assert self.head_dim in [
          16,
          32,
          64,
      ], "Only support head_dim == 16, 32, or 64"
  
      self.inner_attn = FlashAttention(
          attention_dropout=attention_dropout, **factory_kwargs
      )
      self.rearrange = rearrange
  
imports: from abc import abstractmethod
  import math
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from .fp16_util import convert_module_to_f16, convert_module_to_f32
  from .nn import checkpoint, conv_nd, linear, avg_pool_nd, zero_module, normalization, timestep_embedding
label_desc: Use image node_id 1 to create a new QKVFlashAttention object from inherited base classes: nn with example: obj = QKVFlashAttention(embed_dim, num_heads, batch_first, attention_dropout, causal, device, dtype)