class_docstr: Multi-head self attention.
  
  Using torch or triton attention implementation enables user to also use
  additive bias.
class_docstr_tok: ["Multi-head","self","attention",".","Using","torch","or","triton","attention","implementation","enables","user","to","also","use","additive","bias","."]
code: def __init__(
      self,
      config,
      prefix,
      weights,
  ):
      super().__init__()
      attn_impl = config.attn_config["attn_impl"]
      self.attn_impl = config.attn_config["attn_impl"]
      self.clip_qkv = config.attn_config["clip_qkv"]
      self.qk_ln = config.attn_config["qk_ln"]
      self.d_model = config.d_model
      d_model = config.d_model
      self.n_heads = config.n_heads
      self.softmax_scale = config.attn_config["softmax_scale"]
      if self.softmax_scale is None:
          self.softmax_scale = 1 / math.sqrt(
              self.d_model / self.n_heads
          )
      self.attn_dropout_p = config.attn_config["attn_pdrop"]
  
      if self.n_heads % weights.process_group.size() != 0:
          raise ValueError(
              f"`n_heads` must be divisible by `num_shards` (got `n_heads`: {self.n_heads} "
              f"and `num_shards`: {weights.process_group.size()}"
          )
      self.n_heads = self.n_heads // weights.process_group.size()
      self.Wqkv = load_col(
          config,
          prefix=f"{prefix}.Wqkv",
          weights=weights,
          bias=not config.no_bias,
      )
      if self.qk_ln:
          raise NotImplementedError("qk_ln is not supported")
      if self.attn_impl == "flash":
          self.attn_fn = flash_attn_fn
      elif self.attn_impl == "triton":
          self.attn_fn = triton_flash_attn_fn
      elif self.attn_impl == "torch":
          self.attn_fn = scaled_multihead_dot_product_attention
      else:
          raise ValueError(
              f"attn_impl={attn_impl!r} is an invalid setting."
          )
      self.out_proj = TensorParallelRowLinear.load(
          config,
          prefix=f"{prefix}.out_proj",
          weights=weights,
          bias=not config.no_bias,
      )
  
imports: import math
  import os
  import warnings
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.nn
  import torch.nn.functional
  from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
  from einops import rearrange
  from packaging import version
  from text_generation_server.utils.layers import TensorParallelEmbedding, TensorParallelColumnLinear, TensorParallelRowLinear, TensorParallelHead, get_linear
label_desc: Use image node_id 1 to create a new MultiheadAttention object from inherited base classes: nn with example: obj = MultiheadAttention(config, prefix, weights)