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_encoder_forward(
      self: GLMTransformer,
      hidden_states,
      attention_mask,
      kv_caches=None,
      use_cache: Optional[bool] = True,
      output_hidden_states: Optional[bool] = False,
      infer_state: Optional[BatchInferState] = None,
      stage_manager: Optional[PipelineStageManager] = None,
      stage_index: Optional[List[int]] = None,
      shard_config: ShardConfig = None,
  ):
      hidden_states = hidden_states.transpose(0, 1).contiguous()
  
      infer_state.decode_layer_id = 0
      start_idx, end_idx = stage_index[0], stage_index[1]
      if kv_caches is None:
          kv_caches = tuple([None] * (end_idx - start_idx + 1))
  
      for idx, kv_cache in zip(range(start_idx, end_idx), kv_caches):
          layer = self.layers[idx]
          layer_ret = layer(
              hidden_states,
              attention_mask,
              kv_cache=kv_cache,
              use_cache=use_cache,
              infer_state=infer_state,
          )
          infer_state.decode_layer_id += 1
  
          hidden_states, _ = layer_ret
  
      hidden_states = hidden_states.transpose(0, 1).contiguous()
  
      if self.post_layer_norm and (
          stage_manager.is_last_stage() or stage_manager.num_stages == 1
      ):
          # Final layer norm.
          hidden_states = self.final_layernorm(hidden_states)
  
      return 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 3 for calling the ChatGLM2InferenceForwards obj's underlying member method code with example usage: obj.chatglm_encoder_forward(hidden_states, attention_mask, kv_caches, use_cache, output_hidden_states, infer_state, stage_manager, stage_index, shard_config) and returns: hidden_states