code: def forward(
      self,
      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,
  ):
      fused_qkv = self.query_key_value(
          hidden_states
      )  # [batch_size, seq_length, 3 x hidden_size]
      batch_size, q_length, _ = fused_qkv.shape
  
      if layer_past is not None:
          past_key, past_value = layer_past
          layer_past = (
              past_key.view(-1, *past_key.shape[-2:]),
              past_value.view(-1, *past_value.shape[-2:]),
          )
  
      if CUSTOM_KERNELS_ENABLED:
          assert (
              self.training is False
          ), "Only foward pass was implemented"
          assert (
              attention_mask.shape[-1] < 4096
          ), "Custom kernel support only up to 4096 tokens"
          (
              context_layer,
              present,
              attention_probs,
          ) = fused_bloom_attention_cuda.forward(
              fused_qkv,
              layer_past,
              alibi,
              attention_mask,
              head_mask,
              self.beta,
              self.inv_norm_factor,
              self.num_heads,
              use_cache,
          )
      else:
          (
              context_layer,
              present,
              attention_probs,
          ) = self.compute_attention(
              fused_qkv=fused_qkv,
              layer_past=layer_past,
              alibi=alibi,
              attention_mask=attention_mask,
              head_mask=head_mask,
              beta=self.beta,
              inv_norm_factor=self.inv_norm_factor,
              num_heads=self.num_heads,
              use_cache=use_cache,
          )
  
      # 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)
  
      # output_tensor = dropout_add(output_tensor, residual, self.hidden_dropout, self.training)
      output_tensor += residual
  
      outputs = (output_tensor, present)
      if output_attentions:
          outputs += (attention_probs,)
  
      return outputs
  
imports: import math
  import os
  import warnings
  from typing import Optional, Tuple, Union
  import torch
  import torch.distributed
  import torch.utils.checkpoint
  from torch import nn
  from torch.nn import LayerNorm
  from torch.nn import functional
  from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions
  from transformers import BloomConfig, PreTrainedModel
  from text_generation_server.utils.layers import TensorParallelColumnLinear, TensorParallelEmbedding, TensorParallelRowLinear, TensorParallelHead
label_desc: Use image node_id 3 for calling the BloomAttention obj's underlying member method code with example usage: obj.forward(hidden_states, residual, alibi, attention_mask, layer_past, head_mask, use_cache, output_attentions) and returns: outputs