code: def compute_attention(
      fused_qkv: torch.Tensor,
      layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]],
      alibi: torch.Tensor,
      attention_mask: torch.Tensor,
      head_mask: Optional[torch.Tensor],
      beta: float,
      inv_norm_factor: float,
      num_heads: int,
      use_cache: bool,
  ):
      batch_size, q_length, three_times_hidden_size = fused_qkv.shape
      head_dim = three_times_hidden_size // (3 * num_heads)
      batch_size * num_heads
  
      ### TODO @thomasw21: this takes quite a bit of time, how do I accelerate that?
      # 3 x [batch_size, seq_length, num_heads, head_dim]
      (query_layer, key_layer, value_layer) = _split_heads(
          fused_qkv, num_heads=num_heads, head_dim=head_dim
      )
  
      if layer_past is not None:
          past_key, past_value = layer_past
          # concatenate along seq_length dimension:
          #  - key: [batch_size * self.num_heads, head_dim, kv_length]
          #  - value: [batch_size * self.num_heads, kv_length, head_dim]
          past_key = past_key.view(-1, *past_key.shape[-2:])
          key_layer = torch.cat((past_key, key_layer), dim=2)
          past_value = past_value.view(-1, *past_value.shape[-2:])
          value_layer = torch.cat((past_value, value_layer), dim=1)
  
      _, _, kv_length = key_layer.shape
  
      if use_cache is True:
          present = (key_layer, value_layer)
      else:
          present = None
      ###
  
      # [batch_size * num_heads, q_length, kv_length]
      # we use `torch.Tensor.baddbmm` instead of `torch.baddbmm` as the latter isn't supported by TorchScript v1.11
      attention_scores = alibi.baddbmm(
          batch1=query_layer,
          batch2=key_layer,
          beta=beta,
          alpha=inv_norm_factor,
      )
  
      # cast attention scores to fp32, compute scaled softmax and cast back to initial dtype - [batch_size, num_heads, q_length, kv_length]
      input_dtype = attention_scores.dtype
      # `float16` has a minimum value of -65504.0, whereas `bfloat16` and `float32` have a minimum value of `-3.4e+38`
      if input_dtype == torch.float16:
          attention_scores = attention_scores.to(torch.float)
      # torch.finfo not supported by torch.jit, we temporarily remplace with `-1e34`
      attn_weights = attention_scores.masked_fill_(
          attention_mask, torch.finfo(attention_scores.dtype).min
      )
      attention_probs = F.softmax(
          attn_weights, dim=-1, dtype=torch.float32
      ).to(input_dtype)
  
      # # [batch_size, num_heads, q_length, kv_length]
      # attention_probs = self.attention_dropout(attention_probs)
  
      if head_mask is not None:
          attention_probs = attention_probs * head_mask
  
      # matmul: [batch_size * num_heads, q_length, head_dim]
      context_layer = torch.bmm(
          attention_probs, value_layer, out=query_layer
      )
  
      # change view [batch_size, num_heads, q_length, head_dim]
      context_layer = _merge_heads(
          context_layer, num_heads=num_heads, head_dim=head_dim
      )
  
      return context_layer, present, attention_probs
  
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 2 for calling the BloomAttention obj's underlying member method code with example usage: obj.compute_attention(fused_qkv, layer_past, alibi, attention_mask, head_mask, beta, inv_norm_factor, num_heads, use_cache) and returns: context_layer,  present,  attention_probs