class_docstr: Multi-headed attention from 'Attention Is All You Need' paper
class_docstr_tok: ["Multi-headed","attention","from","'Attention","Is","All","You","Need","'","paper"]
code: def __init__(self, config: LlamaConfig):
      super().__init__()
      self.config = config
      self.hidden_size = config.hidden_size
      self.num_heads = config.num_attention_heads
      self.head_dim = config.head_dim
      self.num_key_value_heads = config.num_key_value_heads
      self.num_key_value_groups = (
          self.num_heads // self.num_key_value_heads
      )
      self.pretraining_tp = config.pretraining_tp
      self.max_position_embeddings = config.max_position_embeddings
  
      self.q_proj = nn.Linear(
          self.hidden_size, self.num_heads * self.head_dim, bias=False
      )
      self.k_proj = nn.Linear(
          self.hidden_size,
          self.num_key_value_heads * self.head_dim,
          bias=False,
      )
      self.v_proj = nn.Linear(
          self.hidden_size,
          self.num_key_value_heads * self.head_dim,
          bias=False,
      )
      self.o_proj = nn.Linear(
          self.num_heads * self.head_dim, self.hidden_size, bias=False
      )
      self.sin, self.cos = generate_fixed_pos_embedding(
          self.head_dim, self.max_position_embeddings, 1.0, 1e4
      )
  
imports: import math
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.nn.functional
  import torch.utils.checkpoint
  from torch import nn
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
  from transformers.modeling_utils import PreTrainedModel
  from transformers.models.llama.modeling_llama import LlamaConfig, LlamaRMSNorm
  from transformers.utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
  from colossalai.kernel.cuda_native.mha.flash_attn_2 import HAS_FLASH_ATTN
  from colossalai.kernel.triton.llama_act_combine_kernel import HAS_TRITON
  from colossalai.moe.layers import SparseMLP
  from colossalai.moe.manager import MOE_MANAGER
  from colossalai.moe.utils import get_activation, set_moe_args
label_desc: Use image node_id 1 to create a new OpenMoeAttention object from inherited base classes: nn with example: obj = OpenMoeAttention(config)