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,
  ):
      hidden_states = hidden_states.transpose(0, 1).contiguous()
      if not kv_caches:
          kv_caches = [None for _ in range(self.num_layers)]
      presents = () if use_cache else None
      all_self_attentions = None
      all_hidden_states = () if output_hidden_states else None
  
      infer_state.decode_layer_id = 0
      for index in range(self.num_layers):
          layer = self.layers[index]
  
          layer_ret = layer(
              hidden_states,
              attention_mask,
              kv_cache=kv_caches[index],
              use_cache=use_cache,
              infer_state=infer_state,
          )
  
          infer_state.decode_layer_id += 1
  
          hidden_states, kv_cache = layer_ret
          if use_cache:
              presents = presents + (kv_cache,)
  
      if output_hidden_states:
          all_hidden_states = all_hidden_states + (hidden_states,)
  
      # Final layer norm.
      hidden_states = hidden_states.transpose(0, 1).contiguous()
  
      if self.post_layer_norm:
          hidden_states = self.final_layernorm(hidden_states)
  
      return (
          hidden_states,
          presents,
          all_hidden_states,
          all_self_attentions,
      )
  
file_path: colossalai/colossalai/legacy/inference/tensor_parallel/modeling/chatglm2.py
imports: import os
  from typing import Optional, Tuple
  import torch
  from torch.nn import CrossEntropyLoss
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
  from colossalai.inference.tensor_parallel.batch_infer_state import BatchInferState
  from colossalai.kernel.triton.token_attention_kernel import Llama2TokenAttentionForwards
  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) and returns: hidden_states,  presents,  all_hidden_states,  all_self_attentions