class_docstr: Fused QKV Attention layer for incrementally improved training efficiency
class_docstr_tok: ["Fused","QKV","Attention","layer","for","incrementally","improved","training","efficiency"]
code: def __init__(
      self,
      config,
      q: torch.nn.Linear,  # pylint: disable=invalid-name
      k: torch.nn.Linear,  # pylint: disable=invalid-name
      v: torch.nn.Linear,  # pylint: disable=invalid-name
      o: torch.nn.Linear,  # pylint: disable=invalid-name
  ):
      super().__init__(config)
      self.config = config
      self.init_device = next(iter(q.state_dict().values())).device
  
      # define equivalent fused qkv projection
      self.out_features: List[int] = [
          q.out_features,
          k.out_features,
          v.out_features,
      ]
      self.qkv_proj = torch.nn.Linear(
          q.in_features,
          sum(self.out_features),
          device=self.init_device,
          bias=False,
      )
      self.o_proj = o
  
      # overwrite initialized weights with pretrained weights
      self.qkv_proj.weight.data = torch.cat(
          (q.weight.data, k.weight.data, v.weight.data), dim=0
      )
  
imports: import logging
  import warnings
  from functools import partial
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.nn.functional
  import transformers
  from einops import rearrange
  from flash_attn.bert_padding import pad_input, unpad_input
  from transformers.modeling_outputs import BaseModelOutputWithPast
  from transformers.models.llama.modeling_llama import LlamaAttention
  from transformers.models.llama.modeling_llama import LlamaDecoderLayer
  from transformers.models.llama.modeling_llama import LlamaMLP, apply_rotary_pos_emb, repeat_kv
  from xformers.ops import SwiGLU
  from axolotl.monkeypatch.utils import get_cu_seqlens_from_pos_ids, set_module_name
label_desc: Use image node_id 1 to create a new FusedAttention object from inherited base classes: LlamaAttention with example: obj = FusedAttention(config, q, k, v, o)