code: def __init__(
      self,
      embed_dim: int,
      num_heads: int,
      dropout: float = 0.0,
      scale=None,
  ):
      super().__init__()
      assert (
          embed_dim % num_heads == 0
      ), f"the embed dim ({embed_dim}) is not divisible by the number of attention heads ({num_heads})."
      if scale is not None:
          self.scale = scale
      else:
          self.scale = 1 / math.sqrt(embed_dim // num_heads)
      self.dropout = dropout
  
      if not HAS_MEM_EFF_ATTN and not HAS_FLASH_ATTN:
          raise Exception("flash attention can not support!")
  
file_path: colossalai/colossalai/kernel/cuda_native/mha/mha.py
imports: import math
  from typing import Optional
  import torch
  from einops import rearrange
  from ..scaled_softmax import AttnMaskType
  from .flash_attn_2 import HAS_FLASH_ATTN
  from .mem_eff_attn import HAS_MEM_EFF_ATTN
  from .utils import Repad, SeqLenInfo, Unpad
label_desc: Use image node_id 1 to create a new ColoAttention object from inherited base classes: torch.nn with example: obj = ColoAttention(embed_dim, num_heads, dropout, scale)