class_docstr: Multi-Head Attention module from "Attention is All You Need"
  :cite:`DBLP:journals/corr/VaswaniSPUJGKP17`.
  
  Similar to standard `dot` attention but uses
  multiple attention distributions simulataneously
  to select relevant items.
  
  .. mermaid::
  
     graph BT
        A[key]
        B[value]
        C[query]
        O[output]
        subgraph Attn
          D[Attn 1]
          E[Attn 2]
          F[Attn N]
        end
        A --> D
        C --> D
        A --> E
        C --> E
        A --> F
        C --> F
        D --> O
        E --> O
        F --> O
        B --> O
  
  Also includes several additional tricks.
  
  Args:
     head_count (int): number of parallel heads
     model_dim (int): the dimension of keys/values/queries,
         must be divisible by head_count
     dropout (float): dropout parameter
     max_relative_positions (int): max relative positions
     attn_type: "self" or "context"
docstr: Compute the context vector and the attention vectors.
  
  Args:
     key (Tensor): set of `key_len`
         key vectors ``(batch, key_len, dim)``
     value (Tensor): set of `key_len`
         value vectors ``(batch, key_len, dim)``
     query (Tensor): set of `query_len`
         query vectors  ``(batch, query_len, dim)``
     mask: binary mask 1/0 indicating which keys have
         zero / non-zero attention ``(batch, query_len, key_len)``
     step (int): decoding step (used for Rotary embedding)
  Returns:
     (Tensor, Tensor):
  
     * output context vectors ``(batch, query_len, dim)``
     * Attention vector in heads ``(batch, head, query_len, key_len)``.
docstr_tok: ["Compute","the","context","vector","and","the","attention","vectors",".","Args",":","key","(","Tensor",")",":","set","of","`","key_len","`","key","vectors","``","(","batch",",","key_len",",","dim",")","``","value","(","Tensor",")",":","set","of","`","key_len","`","value","vectors","``","(","batch",",","key_len",",","dim",")","``","query","(","Tensor",")",":","set","of","`","query_len","`","query","vectors","``","(","batch",",","query_len",",","dim",")","``","mask",":","binary","mask","1\/0","indicating","which","keys","have","zero","\/","non-zero","attention","``","(","batch",",","query_len",",","key_len",")","``","step","(","int",")",":","decoding","step","(","used","for","Rotary","embedding",")","Returns",":","(","Tensor",",","Tensor",")",":","*","output","context","vectors","``","(","batch",",","query_len",",","dim",")","``","*","Attention","vector","in","heads","``","(","batch",",","head",",","query_len",",","key_len",")","``","."]
code: def forward(
      self,
      key: Tensor,
      value: Tensor,
      query: Tensor,
      mask: Optional[Tensor] = None,
      sliding_window: Optional[int] = 0,
      step: Optional[int] = 0,
      return_attn: Optional[bool] = False,
      self_attn_type: str = None,
  ) -> Tuple[Tensor, Tensor]:
      """
      Compute the context vector and the attention vectors.
  
      Args:
         key (Tensor): set of `key_len`
             key vectors ``(batch, key_len, dim)``
         value (Tensor): set of `key_len`
             value vectors ``(batch, key_len, dim)``
         query (Tensor): set of `query_len`
             query vectors  ``(batch, query_len, dim)``
         mask: binary mask 1/0 indicating which keys have
             zero / non-zero attention ``(batch, query_len, key_len)``
         step (int): decoding step (used for Rotary embedding)
      Returns:
         (Tensor, Tensor):
  
         * output context vectors ``(batch, query_len, dim)``
         * Attention vector in heads ``(batch, head, query_len, key_len)``.
      """
      # 1) Project key, value, and query.
      # as a reminder at training layer_cache[0] remains False
      key_pad_mask = self.layer_cache[1].get("key_pad_mask", None)
      if self.layer_cache[0]:
          # Retrieve keys and values from the KV cache (decoding mode only).
          if self.attn_type == "self":
              query, key, value = (
                  self.linear_query(query),
                  self.linear_keys(query),
                  self.linear_values(query),
              )
  
              query = shape(query, self.dim_per_head)
              key = shape(key, self.dim_per_head)
              value = shape(value, self.dim_per_head)
              start_pos = step
              seqlen = query.size(2)
  
              if (
                  step == 0
                  or not self.flash2
                  or self.self_attn_type != "scaled-dot-flash"
                  or self.max_relative_positions not in [0, -1]
                  or query.size(0) > 128
                  or query.dtype != torch.float16
              ):
                  if (
                      self.max_relative_positions == -1
                  ):  # Rotary Embeddings
                      if seqlen > self.rope.size(0):
                          self.rope = rotaryembeddings(
                              self.rotary_dim,
                              maxseqlen=(seqlen + 2048),
                              base=self.rotary_theta,
                          ).to(self.rope.device)
                      rope = self.rope[start_pos : start_pos + seqlen]
                      query, key = apply_rotary_emb(
                          query,
                          key,
                          rope,
                          interleave=self.rotary_interleave,
                      )
  
                  if self.layer_cache[1]["keys"].numel() != 0:
                      key = torch.cat(
                          (self.layer_cache[1]["keys"], key), dim=2
                      )
                      value = torch.cat(
                          (self.layer_cache[1]["values"], value), dim=2
                      )
                      if (
                          sliding_window > 0
                          and key.size(2) > sliding_window
                      ):
                          key = key[:, :, 1:, :]
                          value = value[:, :, 1:, :]
  
                  self.layer_cache[1]["keys"] = key
                  self.layer_cache[1]["values"] = value
  
              else:
                  if (
                      self.max_relative_positions == -1
                  ):  # Rotary Embeddings
                      if seqlen > self.rope.size(0):
                          self.rope = rotaryembeddings(
                              self.rotary_dim,
                              maxseqlen=(seqlen + 2048),
                              base=self.rotary_theta,
                          ).to(self.rope.device)
                          self.cos = (
                              self.rope[:, : self.rope.size(1) // 2]
                              .real.contiguous()
                              .half()
                          )
                          self.sin = (
                              self.rope[:, : self.rope.size(1) // 2]
                              .imag.contiguous()
                              .half()
                          )
                  if start_pos >= self.layer_cache[1]["keys"].size(2):
                      self.layer_cache[1]["keys"] = torch.cat(
                          [
                              self.layer_cache[1]["keys"],
                              torch.zeros(
                                  self.layer_cache[1]["keys"].shape[:-2]
                                  + (32,)
                                  + self.layer_cache[1]["keys"].shape[
                                      -1:
                                  ],
                                  device=query.device,
                              ).half(),
                          ],
                          dim=-2,
                      )
                      self.layer_cache[1]["values"] = torch.cat(
                          [
                              self.layer_cache[1]["values"],
                              torch.zeros(
                                  self.layer_cache[1]["values"].shape[
                                      :-2
                                  ]
                                  + (32,)
                                  + self.layer_cache[1]["values"].shape[
                                      -1:
                                  ],
                                  device=query.device,
                              ).half(),
                          ],
                          dim=-2,
                      )
                  if (
                      sliding_window > 0
                      and key.size(2) > sliding_window
                  ):
                      self.layer_cache[1]["keys"] = self.layer_cache[1][
                          "keys"
                      ][:, :, 1:, :]
                      self.layer_cache[1]["values"] = self.layer_cache[
                          1
                      ]["values"][:, :, 1:, :]
                  context = self.flash_attn_with_kvcache(
                      query.transpose(1, 2),
                      self.layer_cache[1]["keys"].transpose(1, 2),
                      self.layer_cache[1]["values"].transpose(1, 2),
                      key.transpose(1, 2),
                      value.transpose(1, 2),
                      rotary_cos=self.cos,
                      rotary_sin=self.sin,
                      cache_seqlens=step,
                      rotary_interleaved=self.rotary_interleave,
                  ).transpose(1, 2)
                  attn_output = self.final_linear(unshape(context))
                  if self.parallel_gpu > 1:
                      all_reduce(attn_output)
                  return attn_output, None
  
          elif self.attn_type == "context":
              query = self.linear_query(query)
              query = shape(query, self.dim_per_head)
              if self.layer_cache[1]["keys"].numel() == 0:
                  key, value = self.linear_keys(
                      key
                  ), self.linear_values(value)
                  key = shape(key, self.dim_per_head)
                  value = shape(value, self.dim_per_head)
              else:
                  key, value = (
                      self.layer_cache[1]["keys"],
                      self.layer_cache[1]["values"],
                  )
              self.layer_cache[1]["keys"] = key
              self.layer_cache[1]["values"] = value
  
          if key_pad_mask is not None:
              # Increase the cached key pad mask by concatenation.
              # For decoding only.
              if step > 0:
                  y = torch.zeros(
                      (key_pad_mask.size(0), key_pad_mask.size(1), 1),
                      dtype=torch.bool,
                      device=key_pad_mask.device,
                  )
                  self.layer_cache[1]["key_pad_mask"] = torch.cat(
                      (key_pad_mask, y), 2
                  )
                  key_pad_mask = self.layer_cache[1]["key_pad_mask"]
      else:
          # Retrieve keys and values from linear layers (training mode).
          key = self.maybe_ckpt(self.linear_keys, key)
          value = self.maybe_ckpt(self.linear_values, value)
          query = self.maybe_ckpt(self.linear_query, query)
  
          key = shape(key, self.dim_per_head)
          value = shape(value, self.dim_per_head)
          query = shape(query, self.dim_per_head)
  
          if self.max_relative_positions == -1:  # Rotary Embeddings
              start_pos = 0
              seqlen = query.size(2)
              if seqlen > self.rope.size(0):
                  self.rope = rotaryembeddings(
                      self.rotary_dim,
                      maxseqlen=(seqlen + 2048),
                      base=self.rotary_theta,
                  ).to(self.rope.device)
              rope = self.rope[start_pos : start_pos + seqlen].to(
                  query.device
              )
              query, key = apply_rotary_emb(
                  query, key, rope, interleave=self.rotary_interleave
              )
  
      b, h, l, d = key.size()
      if self.num_kv > 0:
          qh = query.size(1)
          # expand key on heads dimension when it's less than query heads (multi-query variant)
          key = key.view(b, -1, 1, l, d).repeat(1, 1, qh // h, 1, 1)
          key = key.view(b, qh, l, d)
  
          # expand value on heads dimension when it's less than query heads (multi-query variant)
          value = value.view(b, -1, 1, l, d).repeat(1, 1, qh // h, 1, 1)
          value = value.view(b, qh, l, d)
  
      # 2) When standard pos. enc. or rotary, use flash attention
  
      # Ultimately flashv2 will be part of pytorch https://github.com/pytorch/pytorch/pull/105602
      # In the meantime: if vanilla tranformer or Rotary embeddings (not rel_pos, not alibi)
      # then use flash2 if seq len > 256 otherwise use xtransformer from pt2 uptream
      flash2 = (
          self.flash2
          and l
          > 256  # https://github.com/Dao-AILab/flash-attention/issues/591
      )
  
      if (
          self.max_relative_positions in [-1, 0]
          and not return_attn
          and query.device != torch.device("cpu")
          and self.self_attn_type == "scaled-dot-flash"
      ):
          # Apply flash2 attention.
          causal = (
              self.is_decoder
              and self.attn_type == "self"
              and mask is not None
          )
          if self.is_decoder and self.attn_type == "self" and flash2:
              if causal:
                  window_size = (
                      (-1, -1)
                      if sliding_window == 0
                      else (sliding_window, 0)
                  )
              else:
                  window_size = (-1, -1)
              attn_output = self.flash_attn_func(
                  query.transpose(1, 2),
                  key.transpose(1, 2),
                  value.transpose(1, 2),
                  dropout_p=self.dropout_p,
                  causal=causal,
                  window_size=window_size,
              ).transpose(1, 2)
          else:
              # Apply scaled dot product attention.
              with torch.backends.cuda.sdp_kernel(
                  enable_flash=False,
                  enable_math=True,
                  enable_mem_efficient=True,
              ):
                  attn_output = scaled_dot_product_attention(
                      query,
                      key,
                      value,
                      ~mask if mask is not None else None,
                      self.dropout_p,
                      is_causal=causal,
                  )
          attn = None
  
      else:
          query /= sqrt(self.dim_per_head)
          # batch x num_heads x query_len x key_len
          scores = torch.matmul(query, key.transpose(2, 3))
  
          if self.relative_attention_bias is not None:
              q_len = (
                  key.size(2) if self.layer_cache[0] else query.size(2)
              )
              relative_position_bucket = compute_bias(
                  q_len,
                  key.size(2),
                  self.is_decoder,
                  self.max_relative_positions,
                  self.relative_positions_buckets,
                  device=key.device,
              )
              values = self.relative_attention_bias(
                  relative_position_bucket
              )  # shape (query_length, key_length, num_heads)
              position_bias = values.permute([2, 0, 1]).unsqueeze(
                  0
              )  # shape (1, num_heads, query_length, key_length)
              if self.layer_cache[0]:
                  position_bias = position_bias[
                      :, :, -query.size(2) :, :
                  ]
              scores.add_(position_bias)
  
          elif self.relative_positions_embeddings is not None:
              key_len = key.size(2)
              # 1 or key_len x key_len
              relative_positions_matrix = gen_relative_positions(
                  key_len,
                  self.max_relative_positions,
                  cache=self.layer_cache[0],
                  device=key.device,
              )
              #  1 or key_len x key_len x dim_per_head
              relations_keys = self.relative_positions_embeddings(
                  relative_positions_matrix
              )
              scores.add_(relative_matmul(query, relations_keys, True))
          elif self.max_relative_positions == -2:  # Alibi
              scores = self.alibi(scores)
  
          scores = scores.float()
          if key_pad_mask is not None and mask is None:
              mask = key_pad_mask.unsqueeze(1)
  
          if mask is not None:
              # not 100% necessary but expand to nb of heads
              mask = mask.expand(
                  -1, self.head_count // self.parallel_gpu, -1, -1
              )
              # now mask and scores have the same shape
              scores = scores.masked_fill(mask, -1e18)
  
          # 3) Apply attention dropout and compute context vectors.
          attn = self.softmax(scores).to(query.dtype)
          drop_attn = self.dropout(attn) if self.dropout_p > 0 else attn
  
          attn_output = torch.matmul(drop_attn, value)
  
          if self.relative_positions_embeddings is not None:
              # We use the same embeddings for key and value
              relations_values = relations_keys
              attn_output.add_(
                  relative_matmul(drop_attn, relations_values, False)
              )
  
      context = unshape(attn_output)
      if key_pad_mask is not None:
          if key_pad_mask.size(0) > 1 and context.size(1) > 1:
              x = (
                  key_pad_mask.squeeze(1)
                  .unsqueeze(2)
                  .expand(-1, -1, context.size(2))
              )
              context = context.masked_fill(x, 0)
  
      if self.layer_cache[0]:
          attn_output = self.final_linear(context)
      else:
          attn_output = self.maybe_ckpt(self.final_linear, context)
  
      if self.parallel_gpu > 1:
          all_reduce(attn_output)
  
      return attn_output, attn
  
imports: import torch
  import torch.nn
  from math import log, sqrt
  from torch import Tensor
  from typing import Optional, Tuple
  from torch.nn.functional import scaled_dot_product_attention
  from torch.utils.checkpoint import checkpoint
  from torch.nn.utils import skip_init
  from .alibi_position_bias import AlibiPositionalBias
  from torch.distributed import all_reduce
  from importlib import import_module
label_desc: Use image node_id 3 for calling the MultiHeadedAttention obj's underlying member method code with example usage: obj.forward(key, value, query, mask, sliding_window, step, return_attn, _attn_type) and returns: attn_output,  attn, attn_output,  None