class_docstr: This class holds forwards for Chatglm2 inference.
  We intend to replace the forward methods for ChatGLMModel, ChatGLMEecoderLayer, and ChatGLMAttention.
class_docstr_tok: ["This","class","holds","forwards","for","Chatglm2","inference",".","We","intend","to","replace","the","forward","methods","for","ChatGLMModel",",","ChatGLMEecoderLayer",",","and","ChatGLMAttention","."]
code: def chatglm_model_forward(
      self: ChatGLMModel,
      input_ids: Optional[torch.Tensor] = None,
      position_ids: Optional[torch.Tensor] = None,
      attention_mask: Optional[torch.BoolTensor] = None,
      full_attention_mask: Optional[torch.BoolTensor] = None,
      past_key_values: Optional[
          Tuple[Tuple[torch.Tensor, torch.Tensor], ...]
      ] = None,
      inputs_embeds: Optional[torch.Tensor] = None,
      use_cache: 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,
      shard_config: ShardConfig = None,
  ):
      output_hidden_states = (
          output_hidden_states
          if output_hidden_states is not None
          else self.config.output_hidden_states
      )
  
      if stage_manager.is_first_stage():
          if input_ids is not None and inputs_embeds is not None:
              raise ValueError(
                  "You cannot specify both input_ids and 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 input_ids or inputs_embeds"
              )
          if inputs_embeds is None:
              inputs_embeds = self.embedding(input_ids)
          if position_ids is None:
              position_ids = get_position_ids(
                  batch_size, seq_length, input_ids.device
              )
          hidden_states = inputs_embeds
      else:
          assert (
              hidden_states is not None
          ), "hidden_states should not be None in non-first stage"
          seq_length, batch_size, _ = hidden_states.shape
          if position_ids is None:
              position_ids = get_position_ids(
                  batch_size, seq_length, 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
  
      seq_length_with_past = seq_length + past_key_values_length
  
      # prefill stage at first
      if seq_length != 1:
          infer_state.is_context_stage = True
          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[
                  :, seq_length_with_past - 1
              ] = infer_state.decode_mem_index
          else:
              print(f" *** Encountered allocation non-contiguous")
              print(
                  f"    infer_state.cache_manager.past_key_values_length: {infer_state.cache_manager.past_key_values_length}"
              )
              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[
                  :, seq_length_with_past - 1
              ] = infer_state.decode_mem_index
  
      # related to rotary embedding
      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()
  
      if self.pre_seq_len is not None:
          if past_key_values is None:
              past_key_values = self.get_prompt(
                  batch_size=batch_size,
                  device=input_ids.device,
                  dtype=inputs_embeds.dtype,
              )
          if attention_mask is not None:
              attention_mask = torch.cat(
                  [
                      attention_mask.new_ones(
                          (batch_size, self.pre_seq_len)
                      ),
                      attention_mask,
                  ],
                  dim=-1,
              )
      if full_attention_mask is None:
          if (
              attention_mask is not None and not attention_mask.all()
          ) or (past_key_values and seq_length != 1):
              full_attention_mask = get_masks(
                  self,
                  input_ids,
                  infer_state.cache_manager.past_key_values_length,
                  padding_mask=attention_mask,
              )
  
      # Run encoder.
      hidden_states = self.encoder(
          hidden_states,
          full_attention_mask,
          kv_caches=past_key_values,
          use_cache=use_cache,
          output_hidden_states=output_hidden_states,
          infer_state=infer_state,
          stage_manager=stage_manager,
          stage_index=stage_index,
          shard_config=shard_config,
      )
  
      # update indices
      infer_state.start_loc = 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
  
      return {"hidden_states": hidden_states}
  
file_path: ColossalAI/colossalai/inference/engine/modeling/chatglm2.py
imports: from typing import List, Optional, Tuple
  import torch
  from transformers.utils import logging
  from colossalai.inference.kv_cache import BatchInferState
  from colossalai.kernel.triton.token_attention_kernel import Llama2TokenAttentionForwards
  from colossalai.pipeline.stage_manager import PipelineStageManager
  from colossalai.shardformer import ShardConfig
  from colossalai.shardformer.modeling.chatglm2_6b.modeling_chatglm import ChatGLMForConditionalGeneration, ChatGLMModel, GLMBlock, GLMTransformer, SelfAttention, split_tensor_along_last_dim
  from ._utils import copy_kv_to_mem_cache
label_desc: Use image node_id 2 for calling the ChatGLM2InferenceForwards obj's underlying member method code with example usage: obj.chatglm_model_forward(input_ids, position_ids, attention_mask, full_attention_mask, past_key_values, inputs_embeds, use_cache, output_hidden_states, return_dict, infer_state, stage_manager, hidden_states, stage_index, shard_config) and returns: dict