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_model_forward(
      self: BloomModel,
      input_ids: Optional[torch.LongTensor] = None,
      past_key_values: Optional[
          Tuple[Tuple[torch.Tensor, torch.Tensor], ...]
      ] = None,
      attention_mask: Optional[torch.Tensor] = None,
      head_mask: Optional[torch.LongTensor] = None,
      inputs_embeds: Optional[torch.LongTensor] = None,
      use_cache: Optional[bool] = None,
      output_attentions: Optional[bool] = None,
      output_hidden_states: Optional[bool] = None,
      return_dict: Optional[bool] = None,
      infer_state: Optional[BatchInferState] = None,
      **deprecated_arguments,
  ) -> Union[
      Tuple[torch.Tensor, ...],
      BaseModelOutputWithPastAndCrossAttentions,
  ]:
      logger = logging.get_logger(__name__)
  
      if deprecated_arguments.pop("position_ids", False) is not False:
          # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None`
          warnings.warn(
              "`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. You can safely ignore"
              " passing `position_ids`.",
              FutureWarning,
          )
      if len(deprecated_arguments) > 0:
          raise ValueError(
              f"Got unexpected arguments: {deprecated_arguments}"
          )
  
      output_attentions = (
          output_attentions
          if output_attentions is not None
          else self.config.output_attentions
      )
      output_hidden_states = (
          output_hidden_states
          if output_hidden_states is not None
          else self.config.output_hidden_states
      )
      use_cache = (
          use_cache if use_cache is not None else self.config.use_cache
      )
      return_dict = (
          return_dict
          if return_dict is not None
          else self.config.use_return_dict
      )
  
      if input_ids is not None and inputs_embeds is not None:
          raise ValueError(
              "You cannot specify both input_ids and inputs_embeds at the same time"
          )
      elif input_ids is not None:
          batch_size, seq_length = input_ids.shape
      elif inputs_embeds is not None:
          batch_size, seq_length, _ = inputs_embeds.shape
      else:
          raise ValueError(
              "You have to specify either input_ids or inputs_embeds"
          )
  
      # still need to keep past_key_values to fit original forward flow
      if past_key_values is None:
          past_key_values = tuple([None] * len(self.h))
  
      # Prepare head mask if needed
      # 1.0 in head_mask indicate we keep the head
      # attention_probs has shape batch_size x num_heads x N x N
      # head_mask has shape n_layer x batch x num_heads x N x N
      head_mask = self.get_head_mask(head_mask, self.config.n_layer)
  
      if inputs_embeds is None:
          inputs_embeds = self.word_embeddings(input_ids)
  
      hidden_states = self.word_embeddings_layernorm(inputs_embeds)
  
      presents = () if use_cache else None
      all_self_attentions = () if output_attentions else None
      all_hidden_states = () if output_hidden_states else None
  
      if self.gradient_checkpointing and self.training:
          if use_cache:
              logger.warning_once(
                  "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
              )
              use_cache = False
  
      # NOTE determine if BatchInferState is passed in via arg
      #      if not, get the attr binded to the model
      # We might wantto remove setattr later
      if infer_state is None:
          assert hasattr(self, "infer_state")
          infer_state = self.infer_state
  
      # infer_state.cache_manager = self.cache_manager
      if infer_state.is_context_stage:
          past_key_values_length = 0
      else:
          past_key_values_length = infer_state.max_len_in_batch - 1
  
      if use_cache and seq_length != 1:
          # prefill stage
          infer_state.is_context_stage = (
              True  # set prefill stage, notify attention layer
          )
          infer_state.context_mem_index = (
              infer_state.cache_manager.alloc(
                  infer_state.total_token_num
              )
          )
          BatchInferState.init_block_loc(
              infer_state.block_loc,
              infer_state.seq_len,
              seq_length,
              infer_state.context_mem_index,
          )
      else:
          infer_state.is_context_stage = False
          alloc_mem = infer_state.cache_manager.alloc_contiguous(
              batch_size
          )
          if alloc_mem is not None:
              infer_state.decode_is_contiguous = True
              infer_state.decode_mem_index = alloc_mem[0]
              infer_state.decode_mem_start = alloc_mem[1]
              infer_state.decode_mem_end = alloc_mem[2]
              infer_state.block_loc[
                  :, infer_state.max_len_in_batch - 1
              ] = infer_state.decode_mem_index
          else:
              print(f" *** Encountered allocation non-contiguous")
              print(
                  f"    infer_state.max_len_in_batch : {infer_state.max_len_in_batch}"
              )
              infer_state.decode_is_contiguous = False
              alloc_mem = infer_state.cache_manager.alloc(batch_size)
              infer_state.decode_mem_index = alloc_mem
              # infer_state.decode_key_buffer = torch.empty((batch_size, self.tp_head_num_, self.head_dim_), dtype=torch.float16, device="cuda")
              # infer_state.decode_value_buffer = torch.empty((batch_size, self.tp_head_num_, self.head_dim_), dtype=torch.float16, device="cuda")
              infer_state.block_loc[
                  :, infer_state.max_len_in_batch - 1
              ] = infer_state.decode_mem_index
  
      if attention_mask is None:
          attention_mask = torch.ones(
              (batch_size, infer_state.max_len_in_batch),
              device=hidden_states.device,
          )
      else:
          attention_mask = attention_mask.to(hidden_states.device)
  
      # NOTE revise: we might want to store a single 1D alibi(length is #heads) in model,
      #      or store to BatchInferState to prevent re-calculating
      #      When we have multiple process group (e.g. dp together with tp), we need to pass the pg to here
      # alibi = generate_alibi(self.num_heads).contiguous().cuda()
      tp_size = dist.get_world_size()
      curr_tp_rank = dist.get_rank()
      alibi = (
          generate_alibi(self.num_heads * tp_size)
          .contiguous()[
              curr_tp_rank
              * self.num_heads : (curr_tp_rank + 1)
              * self.num_heads
          ]
          .cuda()
      )
      causal_mask = self._prepare_attn_mask(
          attention_mask,
          input_shape=(batch_size, seq_length),
          past_key_values_length=past_key_values_length,
      )
  
      infer_state.decode_layer_id = 0
      for i, (block, layer_past) in enumerate(
          zip(self.h, past_key_values)
      ):
          if output_hidden_states:
              all_hidden_states = all_hidden_states + (hidden_states,)
  
          if self.gradient_checkpointing and self.training:
              # NOTE: currently our KV cache manager does not handle this condition
              def create_custom_forward(module):
                  def custom_forward(*inputs):
                      # None for past_key_value
                      return module(
                          *inputs,
                          use_cache=use_cache,
                          output_attentions=output_attentions,
                      )
  
                  return custom_forward
  
              outputs = torch.utils.checkpoint.checkpoint(
                  create_custom_forward(block),
                  hidden_states,
                  alibi,
                  causal_mask,
                  layer_past,
                  head_mask[i],
              )
          else:
              outputs = block(
                  hidden_states,
                  layer_past=layer_past,
                  attention_mask=causal_mask,
                  head_mask=head_mask[i],
                  use_cache=use_cache,
                  output_attentions=output_attentions,
                  alibi=alibi,
                  infer_state=infer_state,
              )
  
          infer_state.decode_layer_id += 1
          hidden_states = outputs[0]
          if use_cache is True:
              presents = presents + (outputs[1],)
  
          if output_attentions:
              all_self_attentions = all_self_attentions + (
                  outputs[2 if use_cache else 1],
              )
  
      # Add last hidden state
      hidden_states = self.ln_f(hidden_states)
  
      if output_hidden_states:
          all_hidden_states = all_hidden_states + (hidden_states,)
  
      # update indices of kv cache block
      # NOT READY FOR PRIME TIME
      # might want to remove this part, instead, better to pass the BatchInferState from model forward,
      #       and update these information in engine.generate after model foward called
      infer_state.start_loc = infer_state.start_loc + torch.arange(
          0, batch_size, dtype=torch.int32, device="cuda"
      )
      infer_state.seq_len += 1
      infer_state.max_len_in_batch += 1
  
      if not return_dict:
          return tuple(
              v
              for v in [
                  hidden_states,
                  presents,
                  all_hidden_states,
                  all_self_attentions,
              ]
              if v is not None
          )
  
      return BaseModelOutputWithPastAndCrossAttentions(
          last_hidden_state=hidden_states,
          past_key_values=presents,  # should always be (None, None, ..., None)
          hidden_states=all_hidden_states,
          attentions=all_self_attentions,
      )
  
file_path: ColossalAI/colossalai/legacy/inference/tensor_parallel/modeling/bloom.py
imports: import math
  import warnings
  from typing import Optional, Tuple, Union
  import torch
  import torch.distributed
  from torch.nn import CrossEntropyLoss
  from torch.nn import functional
  from transformers.models.bloom.modeling_bloom import BaseModelOutputWithPastAndCrossAttentions, BloomAttention, BloomBlock, BloomForCausalLM, BloomModel, CausalLMOutputWithCrossAttentions
  from transformers.utils import logging
  from colossalai.inference.tensor_parallel.batch_infer_state import BatchInferState
  from colossalai.kernel.triton import bloom_context_attn_fwd, copy_kv_cache_to_dest, token_attention_fwd
label_desc: Use image node_id 1 for calling the BloomInferenceForwards obj's underlying member method code with example usage: obj.bloom_model_forward(input_ids, past_key_values, attention_mask, head_mask, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict, infer_state) and returns: BaseModelOutputWithPastAndCrossAttentions, tuple, custom_forward, module