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_glmblock_forward(
      self: GLMBlock,
      hidden_states,
      attention_mask,
      kv_cache=None,
      use_cache=True,
      infer_state: Optional[BatchInferState] = None,
  ):
      # hidden_states: [s, b, h]
  
      # Layer norm at the beginning of the transformer layer.
      layernorm_output = self.input_layernorm(hidden_states)
      # Self attention.
      attention_output, kv_cache = self.self_attention(
          layernorm_output,
          attention_mask,
          kv_cache=kv_cache,
          use_cache=use_cache,
          infer_state=infer_state,
      )
      # Residual connection.
      if self.apply_residual_connection_post_layernorm:
          residual = layernorm_output
      else:
          residual = hidden_states
      layernorm_input = torch.nn.functional.dropout(
          attention_output,
          p=self.hidden_dropout,
          training=self.training,
      )
      layernorm_input = residual + layernorm_input
      # Layer norm post the self attention.
      layernorm_output = self.post_attention_layernorm(layernorm_input)
      # MLP.
      mlp_output = self.mlp(layernorm_output)
  
      # Second residual connection.
      if self.apply_residual_connection_post_layernorm:
          residual = layernorm_output
      else:
          residual = layernorm_input
  
      output = torch.nn.functional.dropout(
          mlp_output, p=self.hidden_dropout, training=self.training
      )
      output = residual + output
      return output, kv_cache
  
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 4 for calling the ChatGLM2InferenceForwards obj's underlying member method code with example usage: obj.chatglm_glmblock_forward(hidden_states, attention_mask, kv_cache, use_cache, infer_state) and returns: output,  kv_cache