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_flash_attn_kvcache_forward(
      self: SelfAttention,
      hidden_states,
      attention_mask,
      kv_cache=None,
      use_cache=True,
      infer_state: Optional[BatchInferState] = None,
  ):
      assert (
          use_cache is True
      ), "use_cache should be set to True using this chatglm attention"
      # hidden_states: original :[sq, b, h] --> this [b, sq, h]
      batch_size = hidden_states.shape[0]
      hidden_size = hidden_states.shape[-1]
      # Attention heads [sq, b, h] --> [sq, b, (np * 3 * hn)]
      mixed_x_layer = self.query_key_value(hidden_states)
      if self.multi_query_attention:
          (query_layer, key_layer, value_layer) = mixed_x_layer.split(
              [
                  self.num_attention_heads_per_partition
                  * self.hidden_size_per_attention_head,
                  self.num_multi_query_groups_per_partition
                  * self.hidden_size_per_attention_head,
                  self.num_multi_query_groups_per_partition
                  * self.hidden_size_per_attention_head,
              ],
              dim=-1,
          )
          query_layer = query_layer.view(
              query_layer.size()[:-1]
              + (
                  self.num_attention_heads_per_partition,
                  self.hidden_size_per_attention_head,
              )
          )
          key_layer = key_layer.view(
              key_layer.size()[:-1]
              + (
                  self.num_multi_query_groups_per_partition,
                  self.hidden_size_per_attention_head,
              )
          )
          value_layer = value_layer.view(
              value_layer.size()[:-1]
              + (
                  self.num_multi_query_groups_per_partition,
                  self.hidden_size_per_attention_head,
              )
          )
  
      else:
          new_tensor_shape = mixed_x_layer.size()[:-1] + (
              self.num_attention_heads_per_partition,
              3 * self.hidden_size_per_attention_head,
          )
          mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)
          # [sq, b, np, 3 * hn] --> 3 [sq, b, np, hn]
          (
              query_layer,
              key_layer,
              value_layer,
          ) = split_tensor_along_last_dim(mixed_x_layer, 3)
      cos, sin = infer_state.position_cos, infer_state.position_sin
  
      chatglm2_rotary_emb_fwd(
          query_layer.view(
              -1,
              self.num_attention_heads_per_partition,
              self.hidden_size_per_attention_head,
          ),
          cos,
          sin,
      )
      if self.multi_query_attention:
          chatglm2_rotary_emb_fwd(
              key_layer.view(
                  -1,
                  self.num_multi_query_groups_per_partition,
                  self.hidden_size_per_attention_head,
              ),
              cos,
              sin,
          )
      else:
          chatglm2_rotary_emb_fwd(
              key_layer.view(
                  -1,
                  self.num_attention_heads_per_partition,
                  self.hidden_size_per_attention_head,
              ),
              cos,
              sin,
          )
  
      # reshape q k v  to [bsz*sql, num_heads, head_dim]   2*1 ,32/2 ,128
      query_layer = query_layer.reshape(
          -1,
          self.num_attention_heads_per_partition,
          self.hidden_size_per_attention_head,
      )
      key_layer = key_layer.reshape(
          -1,
          self.num_multi_query_groups_per_partition,
          self.hidden_size_per_attention_head,
      )
      value_layer = value_layer.reshape(
          -1,
          self.num_multi_query_groups_per_partition,
          self.hidden_size_per_attention_head,
      )
  
      if infer_state.is_context_stage:
          # first token generation:
          # copy key and value calculated in current step to memory manager
          copy_kv_to_mem_cache(
              infer_state.decode_layer_id,
              key_layer,
              value_layer,
              infer_state.context_mem_index,
              infer_state.cache_manager,
          )
          attn_output = torch.empty_like(
              query_layer.contiguous().view(-1, self.projection_size)
          )
  
          # NOTE: no bug in context attn fwd (del it )
          lightllm_llama2_context_attention_fwd(
              query_layer,
              key_layer,
              value_layer,
              attn_output.view(
                  -1,
                  self.num_attention_heads_per_partition,
                  self.hidden_size_per_attention_head,
              ),
              infer_state.start_loc,
              infer_state.seq_len,
              infer_state.max_len_in_batch,
          )
  
      else:
          if infer_state.decode_is_contiguous:
              # if decode is contiguous, then we copy to key cache and value cache in cache manager directly
              cache_k = infer_state.cache_manager.key_buffer[
                  infer_state.decode_layer_id
              ][
                  infer_state.decode_mem_start : infer_state.decode_mem_end,
                  :,
                  :,
              ]
              cache_v = infer_state.cache_manager.value_buffer[
                  infer_state.decode_layer_id
              ][
                  infer_state.decode_mem_start : infer_state.decode_mem_end,
                  :,
                  :,
              ]
              cache_k.copy_(key_layer)
              cache_v.copy_(value_layer)
          else:
              # if decode is not contiguous, use triton kernel to copy key and value cache
              # k, v shape: [batch_size, num_heads, head_dim/embed_size_per_head
              copy_kv_to_mem_cache(
                  infer_state.decode_layer_id,
                  key_layer,
                  value_layer,
                  infer_state.decode_mem_index,
                  infer_state.cache_manager,
              )
  
          # second token and follows
          attn_output = torch.empty_like(
              query_layer.contiguous().view(-1, self.projection_size)
          )
          cache_k = infer_state.cache_manager.key_buffer[
              infer_state.decode_layer_id
          ][: infer_state.decode_mem_end, :, :]
          cache_v = infer_state.cache_manager.value_buffer[
              infer_state.decode_layer_id
          ][: infer_state.decode_mem_end, :, :]
  
          # ==================================
          # core attention computation is replaced by triton kernel
          # ==================================
          Llama2TokenAttentionForwards.token_attn(
              query_layer,
              cache_k,
              cache_v,
              attn_output,
              infer_state.block_loc,
              infer_state.start_loc,
              infer_state.seq_len,
              infer_state.max_len_in_batch,
              infer_state.other_kv_index,
          )
  
          # print('after attention',torch.isnan(attn_output).any())
  
      # =================
      # Output:[b,sq, h]
      # =================
      output = self.dense(attn_output).reshape(
          batch_size, -1, hidden_size
      )
  
      return output, kv_cache
  
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 5 for calling the ChatGLM2InferenceForwards obj's underlying member method code with example usage: obj.chatglm_flash_attn_kvcache_forward(hidden_states, attention_mask, kv_cache, use_cache, infer_state) and returns: output,  kv_cache