class_docstr: This class holds forwards for llama inference.
  We intend to replace the forward methods for LlamaModel, LlamaDecoderLayer, and LlamaAttention for LlamaForCausalLM.
class_docstr_tok: ["This","class","holds","forwards","for","llama","inference",".","We","intend","to","replace","the","forward","methods","for","LlamaModel",",","LlamaDecoderLayer",",","and","LlamaAttention","for","LlamaForCausalLM","."]
code: def llama_decoder_layer_forward(
      self: LlamaDecoderLayer,
      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,
      infer_state: Optional[BatchInferState] = None,
  ) -> Tuple[
      torch.FloatTensor,
      Optional[Tuple[torch.FloatTensor, torch.FloatTensor]],
  ]:
      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,
          infer_state=infer_state,
      )
  
      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
  
file_path: ColossalAI/colossalai/legacy/inference/tensor_parallel/modeling/llama.py
imports: import math
  from typing import List, Optional, Tuple
  import torch
  from transformers.modeling_outputs import BaseModelOutputWithPast
  from transformers.models.llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaModel
  from colossalai.inference.tensor_parallel.batch_infer_state import BatchInferState
  from colossalai.kernel.triton import llama_context_attn_fwd, token_attention_fwd
  from colossalai.kernel.triton.token_attention_kernel import Llama2TokenAttentionForwards
  from ._utils import copy_kv_to_mem_cache
label_desc: Use image node_id 2 for calling the LlamaInferenceForwards obj's underlying member method code with example usage: obj.llama_decoder_layer_forward(hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache, infer_state) and returns: outputs