class_docstr: This class serves a micro library for bloom inference forwards.
  We intend to replace the forward methods for BloomForCausalLM, BloomModel, BloomBlock, and BloomAttention,
  as well as prepare_inputs_for_generation method for BloomForCausalLM.
  For future improvement, we might want to skip replacing methods for BloomForCausalLM,
  and call BloomModel.forward iteratively in TpInferEngine
class_docstr_tok: ["This","class","serves","a","micro","library","for","bloom","inference","forwards",".","We","intend","to","replace","the","forward","methods","for","BloomForCausalLM",",","BloomModel",",","BloomBlock",",","and","BloomAttention",",","as","well","as","prepare_inputs_for_generation","method","for","BloomForCausalLM",".","For","future","improvement",",","we","might","want","to","skip","replacing","methods","for","BloomForCausalLM",",","and","call","BloomModel.forward","iteratively","in","TpInferEngine"]
code: def bloom_attention_forward(
      self: BloomAttention,
      hidden_states: torch.Tensor,
      residual: torch.Tensor,
      alibi: torch.Tensor,
      attention_mask: torch.Tensor,
      layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
      head_mask: Optional[torch.Tensor] = None,
      use_cache: bool = False,
      output_attentions: bool = False,
      infer_state: Optional[BatchInferState] = None,
  ):
      fused_qkv = self.query_key_value(
          hidden_states
      )  # [batch_size, seq_length, 3 x hidden_size]
  
      # 3 x [batch_size, seq_length, num_heads, head_dim]
      (query_layer, key_layer, value_layer) = self._split_heads(
          fused_qkv
      )
      batch_size, q_length, H, D_HEAD = query_layer.shape
      k = key_layer.reshape(
          -1, H, D_HEAD
      )  # batch_size * q_length, H, D_HEAD, q_lenth == 1
      v = value_layer.reshape(
          -1, H, D_HEAD
      )  # batch_size * q_length, H, D_HEAD, q_lenth == 1
  
      mem_manager = infer_state.cache_manager
      layer_id = infer_state.decode_layer_id
  
      if infer_state.is_context_stage:
          # context process
          max_input_len = q_length
          b_start_loc = infer_state.start_loc
          b_seq_len = infer_state.seq_len[:batch_size]
          q = query_layer.reshape(-1, H, D_HEAD)
  
          copy_kv_cache_to_dest(
              k,
              infer_state.context_mem_index,
              mem_manager.key_buffer[layer_id],
          )
          copy_kv_cache_to_dest(
              v,
              infer_state.context_mem_index,
              mem_manager.value_buffer[layer_id],
          )
  
          # output = self.output[:batch_size*q_length, :, :]
          output = torch.empty_like(q)
  
          if HAS_LIGHTLLM_KERNEL:
              lightllm_bloom_context_attention_fwd(
                  q,
                  k,
                  v,
                  output,
                  alibi,
                  b_start_loc,
                  b_seq_len,
                  max_input_len,
              )
          else:
              bloom_context_attn_fwd(
                  q,
                  k,
                  v,
                  output,
                  b_start_loc,
                  b_seq_len,
                  max_input_len,
                  alibi,
              )
  
          context_layer = output.view(batch_size, q_length, H * D_HEAD)
      else:
          # query_layer = query_layer.transpose(1, 2).reshape(batch_size * self.num_heads, q_length, self.head_dim)
          # need shape: batch_size, H, D_HEAD (q_length == 1), input q shape : (batch_size, q_length(1), H, D_HEAD)
          assert (
              q_length == 1
          ), "for non-context process, we only support q_length == 1"
          q = query_layer.reshape(-1, H, D_HEAD)
  
          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[layer_id][
                  infer_state.decode_mem_start : infer_state.decode_mem_end,
                  :,
                  :,
              ]
              cache_v = infer_state.cache_manager.value_buffer[
                  layer_id
              ][
                  infer_state.decode_mem_start : infer_state.decode_mem_end,
                  :,
                  :,
              ]
              cache_k.copy_(k)
              cache_v.copy_(v)
          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_cache_to_dest(
                  k,
                  infer_state.decode_mem_index,
                  mem_manager.key_buffer[layer_id],
              )
              copy_kv_cache_to_dest(
                  v,
                  infer_state.decode_mem_index,
                  mem_manager.value_buffer[layer_id],
              )
  
          b_start_loc = infer_state.start_loc
          b_loc = infer_state.block_loc
          b_seq_len = infer_state.seq_len
          output = torch.empty_like(q)
          token_attention_fwd(
              q,
              mem_manager.key_buffer[layer_id],
              mem_manager.value_buffer[layer_id],
              output,
              b_loc,
              b_start_loc,
              b_seq_len,
              infer_state.max_len_in_batch,
              alibi,
          )
  
          context_layer = output.view(batch_size, q_length, H * D_HEAD)
  
      # NOTE: always set present as none for now, instead of returning past key value to the next decoding,
      #       we create the past key value pair from the cache manager
      present = None
  
      # aggregate results across tp ranks. See here: https://github.com/pytorch/pytorch/issues/76232
      if self.pretraining_tp > 1 and self.slow_but_exact:
          slices = self.hidden_size / self.pretraining_tp
          output_tensor = torch.zeros_like(context_layer)
          for i in range(self.pretraining_tp):
              output_tensor = output_tensor + F.linear(
                  context_layer[
                      :, :, int(i * slices) : int((i + 1) * slices)
                  ],
                  self.dense.weight[
                      :, int(i * slices) : int((i + 1) * slices)
                  ],
              )
      else:
          output_tensor = self.dense(context_layer)
  
      # dropout is not required here during inference
      output_tensor = residual + output_tensor
  
      outputs = (output_tensor, present)
      assert (
          output_attentions is False
      ), "we do not support output_attentions at this time"
  
      return outputs
  
file_path: colossalai/colossalai/inference/engine/modeling/bloom.py
imports: import math
  import warnings
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.distributed
  from torch.nn import functional
  from transformers.models.bloom.modeling_bloom import BaseModelOutputWithPastAndCrossAttentions, BloomAttention, BloomBlock, BloomForCausalLM, BloomModel
  from transformers.utils import logging
  from colossalai.inference.kv_cache.batch_infer_state import BatchInferState
  from colossalai.kernel.triton import bloom_context_attn_fwd, copy_kv_cache_to_dest, token_attention_fwd
  from colossalai.pipeline.stage_manager import PipelineStageManager
label_desc: Use image node_id 4 for calling the BloomInferenceForwards obj's underlying member method code with example usage: obj.bloom_attention_forward(hidden_states, residual, alibi, attention_mask, layer_past, head_mask, use_cache, output_attentions, infer_state) and returns: outputs