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 = self.hidden_size // self.num_heads
      self.num_key_value_heads = config.num_key_value_heads
      self.num_key_value_groups = (
          self.num_heads // self.num_key_value_heads
      )
      self.max_position_embeddings = config.max_position_embeddings
      self.rope_theta = config.rope_theta
  
      if (self.head_dim * self.num_heads) != self.hidden_size:
          raise ValueError(
              f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
              f" and `num_heads`: {self.num_heads})."
          )
      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._init_rope()
  
      self.kv_cache = H2OKVCache_LayerWise(
          hh_size=config.hh_size,
          recent_size=config.recent_size,
          k_seq_dim=2,
          v_seq_dim=2,
      )
  
imports: import math
  from typing import Optional, Tuple
  import pdb
  import torch
  from torch import nn
  import torch.utils.checkpoint
  import torch.nn.functional
  from transformers.models.llama.modeling_llama import LlamaAttention, rotate_half, apply_rotary_pos_emb, LlamaRotaryEmbedding, apply_rotary_pos_emb, LlamaForCausalLM
  import types
  from transformers.configuration_utils import PretrainedConfig
label_desc: Use image node_id 1 to create a new H2OLlamaAttention_streaming object from inherited base classes: nn with example: obj = H2OLlamaAttention_streaming(config)