class_docstr: patched version of LlamaDecoderLayer to pass through the precalculated cu_seqlens
class_docstr_tok: ["patched","version","of","LlamaDecoderLayer","to","pass","through","the","precalculated","cu_seqlens"]
docstr: Args:
      hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
      attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
          `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
      output_attentions (`bool`, *optional*):
          Whether or not to return the attentions tensors of all attention layers. See `attentions` under
          returned tensors for more detail.
      use_cache (`bool`, *optional*):
          If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
          (see `past_key_values`).
      past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
      cu_seqlens (`torch.Tensor`, *optional*) cumulative sequence len when packing
docstr_tok: ["Args",":","hidden_states","(","`","torch.FloatTensor","`",")",":","input","to","the","layer","of","shape","`","(","batch",",","seq_len",",","embed_dim",")","`","attention_mask","(","`","torch.FloatTensor","`",",","*","optional","*",")",":","attention","mask","of","size","`","(","batch",",","1",",","tgt_len",",","src_len",")","`","where","padding","elements","are","indicated","by","very","large","negative","values",".","output_attentions","(","`","bool","`",",","*","optional","*",")",":","Whether","or","not","to","return","the","attentions","tensors","of","all","attention","layers",".","See","`","attentions","`","under","returned","tensors","for","more","detail",".","use_cache","(","`","bool","`",",","*","optional","*",")",":","If","set","to","`","True","`",",","`","past_key_values","`","key","value","states","are","returned","and","can","be","used","to","speed","up","decoding","(","see","`","past_key_values","`",")",".","past_key_value","(","`","Tuple","(","torch.FloatTensor",")","`",",","*","optional","*",")",":","cached","past","key","and","value","projection","states","cu_seqlens","(","`","torch.Tensor","`",",","*","optional","*",")","cumulative","sequence","len","when","packing"]
code: def forward(
      self,
      hidden_states: torch.Tensor,
      attention_mask: Optional[torch.Tensor] = None,
      position_ids: Optional[torch.LongTensor] = None,
      past_key_value: Optional[Tuple[torch.Tensor]] = None,
      output_attentions: Optional[bool] = False,
      use_cache: Optional[bool] = False,
      padding_mask: Optional[torch.LongTensor] = None,
      cu_seqlens: Optional[torch.Tensor] = None,
      max_seqlen: Optional[torch.Tensor] = None,
  ) -> Tuple[
      torch.FloatTensor,
      Optional[Tuple[torch.FloatTensor, torch.FloatTensor]],
  ]:
      """
      Args:
          hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
          attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
              `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
          output_attentions (`bool`, *optional*):
              Whether or not to return the attentions tensors of all attention layers. See `attentions` under
              returned tensors for more detail.
          use_cache (`bool`, *optional*):
              If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
              (see `past_key_values`).
          past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
          cu_seqlens (`torch.Tensor`, *optional*) cumulative sequence len when packing
      """
  
      residual = hidden_states
  
      hidden_states = self.input_layernorm(hidden_states)
  
      # Self Attention
      (
          hidden_states,
          self_attn_weights,
          present_key_value,
      ) = self.self_attn(
          hidden_states=hidden_states,
          attention_mask=attention_mask,
          position_ids=position_ids,
          past_key_value=past_key_value,
          output_attentions=output_attentions,
          use_cache=use_cache,
          padding_mask=padding_mask,
          cu_seqlens=cu_seqlens,
          max_seqlen=max_seqlen,
      )
      hidden_states = residual + hidden_states
  
      # Fully Connected
      residual = hidden_states
      hidden_states = self.post_attention_layernorm(hidden_states)
      hidden_states = self.mlp(hidden_states)
      hidden_states = residual + hidden_states
  
      outputs = (hidden_states,)
  
      if output_attentions:
          outputs += (self_attn_weights,)
  
      if use_cache:
          outputs += (present_key_value,)
  
      return outputs
  
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 for calling the LlamaDecoderLayer obj's underlying member method code with example usage: obj.forward(hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache, padding_mask, cu_seqlens, max_seqlen) and returns: outputs