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_model_forward(
      self: LlamaModel,
      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,
      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,
  ):
      return_dict = (
          return_dict
          if return_dict is not None
          else self.config.use_return_dict
      )
  
      use_cache = (
          use_cache if use_cache is not None else self.config.use_cache
      )
      # retrieve input_ids and inputs_embeds
      if stage_manager is None or stage_manager.is_first_stage():
          if input_ids is not None and inputs_embeds is not None:
              raise ValueError(
                  "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
              )
          elif input_ids is not None:
              batch_size, seq_length = input_ids.shape
          elif inputs_embeds is not None:
              batch_size, seq_length, _ = inputs_embeds.shape
          else:
              raise ValueError(
                  "You have to specify either decoder_input_ids or decoder_inputs_embeds"
              )
          device = (
              input_ids.device
              if input_ids is not None
              else inputs_embeds.device
          )
          if inputs_embeds is None:
              inputs_embeds = self.embed_tokens(input_ids)
          hidden_states = inputs_embeds
      else:
          assert stage_manager is not None
          assert (
              hidden_states is not None
          ), f"hidden_state should not be none in stage {stage_manager.stage}"
          input_shape = hidden_states.shape[:-1]
          batch_size, seq_length = input_shape
          device = hidden_states.device
  
      if infer_state.is_context_stage:
          past_key_values_length = 0
      else:
          past_key_values_length = infer_state.max_len_in_batch - 1
  
      # NOTE: differentiate with prefill stage
      #       block_loc require different value-assigning method for two different stage
      if use_cache and seq_length != 1:
          # NOTE assume prefill stage
          # allocate memory block
          infer_state.is_context_stage = (
              True  # set prefill stage, notify attention layer
          )
          infer_state.context_mem_index = (
              infer_state.cache_manager.alloc(
                  infer_state.total_token_num
              )
          )
          infer_state.init_block_loc(
              infer_state.block_loc,
              infer_state.seq_len,
              seq_length,
              infer_state.context_mem_index,
          )
      else:
          infer_state.is_context_stage = False
          alloc_mem = infer_state.cache_manager.alloc_contiguous(
              batch_size
          )
          if alloc_mem is not None:
              infer_state.decode_is_contiguous = True
              infer_state.decode_mem_index = alloc_mem[0]
              infer_state.decode_mem_start = alloc_mem[1]
              infer_state.decode_mem_end = alloc_mem[2]
              infer_state.block_loc[
                  :, infer_state.max_len_in_batch - 1
              ] = infer_state.decode_mem_index
          else:
              infer_state.decode_is_contiguous = False
              alloc_mem = infer_state.cache_manager.alloc(batch_size)
              infer_state.decode_mem_index = alloc_mem
              infer_state.block_loc[
                  :, infer_state.max_len_in_batch - 1
              ] = infer_state.decode_mem_index
  
      if position_ids is None:
          position_ids = torch.arange(
              past_key_values_length,
              seq_length + past_key_values_length,
              dtype=torch.long,
              device=device,
          )
          position_ids = position_ids.repeat(batch_size, 1)
          position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
      else:
          position_ids = position_ids.view(-1, seq_length).long()
  
      if infer_state.is_context_stage:
          infer_state.position_cos = torch.index_select(
              self._cos_cached, 0, position_ids.view(-1)
          ).view(position_ids.view(-1).shape[0], -1)
          infer_state.position_sin = torch.index_select(
              self._sin_cached, 0, position_ids.view(-1)
          ).view(position_ids.view(-1).shape[0], -1)
  
      else:
          seq_len = infer_state.seq_len
          infer_state.position_cos = torch.index_select(
              self._cos_cached, 0, seq_len - 1
          ).view(seq_len.shape[0], -1)
          infer_state.position_sin = torch.index_select(
              self._sin_cached, 0, seq_len - 1
          ).view(seq_len.shape[0], -1)
          infer_state.other_kv_index = infer_state.block_loc[
              0, infer_state.max_len_in_batch - 1
          ].item()
  
      # embed positions
      if attention_mask is None:
          attention_mask = torch.ones(
              (batch_size, infer_state.max_len_in_batch),
              dtype=torch.bool,
              device=hidden_states.device,
          )
  
      attention_mask = self._prepare_decoder_attention_mask(
          attention_mask,
          (batch_size, seq_length),
          hidden_states,
          past_key_values_length,
      )
  
      # decoder layers
      infer_state.decode_layer_id = 0
  
      start_idx, end_idx = stage_index[0], stage_index[1]
      if past_key_values is None:
          past_key_values = tuple([None] * (end_idx - start_idx + 1))
  
      for idx, past_key_value in zip(
          range(start_idx, end_idx), past_key_values
      ):
          decoder_layer = self.layers[idx]
          # NOTE: modify here for passing args to decoder layer
          layer_outputs = decoder_layer(
              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,
          )
          infer_state.decode_layer_id += 1
          hidden_states = layer_outputs[0]
  
      if stage_manager.is_last_stage() or stage_manager.num_stages == 1:
          hidden_states = self.norm(hidden_states)
  
      # update indices
      # infer_state.block_loc[:, infer_state.max_len_in_batch-1] = infer_state.total_token_num + torch.arange(0, batch_size, dtype=torch.int32, device="cuda")
      infer_state.start_loc += torch.arange(
          0, batch_size, dtype=torch.int32, device="cuda"
      )
      infer_state.seq_len += 1
      infer_state.max_len_in_batch += 1
  
      # if not return_dict:
      #     return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
  
      # return BaseModelOutputWithPast(
      #     last_hidden_state=hidden_states,
      #     past_key_values=next_cache,
      #     hidden_states=all_hidden_states,
      #     attentions=all_self_attns,
      # )
      return {"hidden_states": hidden_states}
  
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 2 for calling the LlamaInferenceForwards obj's underlying member method code with example usage: obj.llama_model_forward(input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict, infer_state, stage_manager, hidden_states, stage_index) and returns: dict