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","."]
docstr: Args:
      labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
          Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
          config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
          (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
docstr_tok: ["Args",":","labels","(","`","torch.LongTensor","`","of","shape","`","(","batch_size",",","sequence_length",")","`",",","*","optional","*",")",":","Labels","for","computing","the","masked","language","modeling","loss",".","Indices","should","either","be","in","`","[","0",",","...",",","config.vocab_size","]","`","or","-100","(","see","`","input_ids","`","docstring",")",".","Tokens","with","indices","set","to","`","-100","`","are","ignored","(","masked",")",",","the","loss","is","only","computed","for","the","tokens","with","labels","in","`","[","0",",","...",",","config.vocab_size","]","`","."]
code: def llama_causal_lm_forward(
      self: LlamaForCausalLM,
      input_ids: torch.LongTensor = None,
      attention_mask: Optional[torch.Tensor] = None,
      position_ids: Optional[torch.LongTensor] = None,
      past_key_values: Optional[List[torch.FloatTensor]] = None,
      inputs_embeds: Optional[torch.FloatTensor] = None,
      labels: Optional[torch.LongTensor] = None,
      use_cache: Optional[bool] = None,
      output_attentions: Optional[bool] = None,
      output_hidden_states: Optional[bool] = None,
      return_dict: Optional[bool] = None,
      infer_state: BatchInferState = None,
      stage_manager: Optional[PipelineStageManager] = None,
      hidden_states: Optional[torch.FloatTensor] = None,
      stage_index: Optional[List[int]] = None,
  ):
      r"""
      Args:
          labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
              Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
              config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
              (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
  
      """
      logger = logging.get_logger(__name__)
  
      return_dict = (
          return_dict
          if return_dict is not None
          else self.config.use_return_dict
      )
  
      if output_attentions:
          logger.warning_once(
              "output_attentions=True is not supported for pipeline models at the moment."
          )
          output_attentions = False
      if output_hidden_states:
          logger.warning_once(
              "output_hidden_states=True is not supported for pipeline models at the moment."
          )
          output_hidden_states = False
  
      # If is first stage and after warmup, go throught lm_head first
      if stage_manager.is_first_stage() and hidden_states is not None:
          lm_logits = self.lm_head(hidden_states)
          return {"logits": lm_logits}
  
      # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
      outputs = LlamaInferenceForwards.llama_model_forward(
          self.model,
          input_ids=input_ids,
          attention_mask=attention_mask,
          position_ids=position_ids,
          past_key_values=past_key_values,
          inputs_embeds=inputs_embeds,
          use_cache=use_cache,
          output_attentions=output_attentions,
          output_hidden_states=output_hidden_states,
          return_dict=return_dict,
          infer_state=infer_state,
          stage_manager=stage_manager,
          hidden_states=hidden_states,
          stage_index=stage_index,
      )
  
      return outputs
  
file_path: ColossalAI/colossalai/legacy/inference/hybridengine/modeling/llama.py
imports: import math
  from typing import List, Optional, Tuple
  import torch
  from transformers.models.llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaForCausalLM, LlamaModel
  from transformers.utils import logging
  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 colossalai.pipeline.stage_manager import PipelineStageManager
  from ._utils import copy_kv_to_mem_cache
label_desc: Use image node_id 1 for calling the LlamaInferenceForwards obj's underlying member method code with example usage: obj.llama_causal_lm_forward(input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict, infer_state, stage_manager, hidden_states, stage_index) and returns: outputs, dict