code: def attention_fn(
      self,
      query_layer,
      key_layer,
      value_layer,
      attention_mask,
      hidden_size_per_partition,
      layer_id,
      layer_past=None,
      scaling_attention_score=True,
      use_cache=False,
  ):
      key_layer = key_layer.permute(1, 2, 0, 3).contiguous()
      value_layer = value_layer.permute(1, 2, 0, 3).contiguous()
      # query_layer = query_layer.permute(1, 2, 0, 3)
  
      cur_length, batch_size = (
          query_layer.shape[0],
          query_layer.shape[1],
      )
      device = query_layer.device
  
      if layer_past is not None:
          cache_k, cache_v = layer_past[0], layer_past[1]
          cache_k = cache_k.permute(1, 2, 0, 3)
          cache_v = cache_v.permute(1, 2, 0, 3)
          past_length = cache_k.size(2)
          if cache_k.stride()[1] <= cache_k.size(2) * cache_k.size(3):
              max_cache_length = (
                  past_length + cur_length + KV_CACHE_ALLOC_BLOCK_LENGTH
              )
              new_cache_k, new_cache_v = extend_kv_cache(
                  batch_size,
                  self.num_attention_heads_per_partition,
                  self.hidden_size_per_attention_head,
                  past_length,
                  max_cache_length,
                  dtype=query_layer.dtype,
                  device=device,
              )
              new_cache_k[:] = cache_k
              new_cache_v[:] = cache_v
              cache_k = new_cache_k
              cache_v = new_cache_v
          key_layer, value_layer = append_kv_cache(
              cache_k, cache_v, key_layer, value_layer
          )
  
      elif use_cache:
          max_cache_length = (
              max(KV_CACHE_ALLOC_MIN_LENGTH, cur_length)
              + KV_CACHE_ALLOC_BLOCK_LENGTH
          )
          key_cache, value_cache = init_kv_cache(
              batch_size,
              self.num_attention_heads_per_partition,
              self.hidden_size_per_attention_head,
              cur_length,
              max_cache_length,
              dtype=query_layer.dtype,
              device=device,
          )
          key_cache[:] = key_layer
          value_cache[:] = value_layer
          key_layer = key_cache
          value_layer = value_cache
  
      # seqlen, batch, num_attention_heads, hidden_size_per_attention_head
      b, nh, seq_len, hidden_size = key_layer.shape
  
      if use_cache:
          present = (
              key_layer.permute(2, 0, 1, 3),
              value_layer.permute(2, 0, 1, 3),
          )
      else:
          present = None
  
      pytorch_major_version = int(torch.__version__.split(".")[0])
      if query_layer.size(0) > 1 and pytorch_major_version >= 2:
          query_layer = query_layer.permute(1, 2, 0, 3)
          if (
              attention_mask is None
              and query_layer.shape[2] == key_layer.shape[2]
          ):
              if torch.is_autocast_cpu_enabled():
                  attention_mask = torch.ones(
                      query_layer.shape[2],
                      key_layer.shape[2],
                      dtype=torch.bool,
                  ).tril(diagonal=0)
                  attention_mask = attention_mask.masked_fill(
                      ~attention_mask,
                      -float("inf"),
                  )
                  attention_mask = attention_mask.to(
                      torch.get_autocast_cpu_dtype()
                  )
                  query_layer = query_layer.to(
                      torch.get_autocast_cpu_dtype()
                  )
                  key_layer = key_layer.to(
                      torch.get_autocast_cpu_dtype()
                  )
                  value_layer = value_layer.to(
                      torch.get_autocast_cpu_dtype()
                  )
                  context_layer = (
                      torch.nn.functional.scaled_dot_product_attention(
                          query_layer,
                          key_layer,
                          value_layer,
                          attention_mask,
                          is_causal=False,
                      )
                  )
              else:
                  context_layer = (
                      torch.nn.functional.scaled_dot_product_attention(
                          query_layer,
                          key_layer,
                          value_layer,
                          attention_mask,
                          is_causal=True,
                      )
                  )
          else:
              if attention_mask is not None:
                  attention_mask = ~attention_mask
              attention_mask = attention_mask.masked_fill(
                  ~attention_mask,
                  -float("inf"),
              )
              if torch.is_autocast_cpu_enabled():
                  query_layer = query_layer.to(
                      torch.get_autocast_cpu_dtype()
                  )
                  key_layer = key_layer.to(
                      torch.get_autocast_cpu_dtype()
                  )
                  value_layer = value_layer.to(
                      torch.get_autocast_cpu_dtype()
                  )
                  attention_mask = attention_mask.to(
                      torch.get_autocast_cpu_dtype()
                  )
              context_layer = (
                  torch.nn.functional.scaled_dot_product_attention(
                      query_layer,
                      key_layer,
                      value_layer,
                      attention_mask,
                  )
              )
          context_layer = context_layer.permute(2, 0, 1, 3)
          new_context_layer_shape = context_layer.size()[:-2] + (
              self.hidden_size_per_partition,
          )
          context_layer = context_layer.reshape(
              *new_context_layer_shape
          )
          attention_probs = None
  
      else:
          query_key_layer_scaling_coeff = float(layer_id + 1)
          if scaling_attention_score:
              query_layer = query_layer / (
                  math.sqrt(hidden_size) * query_key_layer_scaling_coeff
              )
  
          # ===================================
          # Raw attention scores. [b, np, s, s]
          # ===================================
  
          # [b, np, sq, sk]
          output_size = (
              query_layer.size(1),
              query_layer.size(2),
              query_layer.size(0),
              key_layer.size(2),
          )
  
          # [sq, b, np, hn] -> [sq, b * np, hn]
          query_layer = query_layer.view(
              output_size[2], output_size[0] * output_size[1], -1
          )
          # [sk, b, np, hn] -> [sk, b * np, hn]
          key_layer = key_layer.view(
              output_size[0] * output_size[1], output_size[3], -1
          )
  
          matmul_input_buffer = torch.empty(
              output_size[0] * output_size[1],
              output_size[2],
              output_size[3],
              dtype=query_layer.dtype,
              device=query_layer.device,
          )
  
          matmul_result = torch.empty(
              output_size[0] * output_size[1],
              output_size[2],
              output_size[3],
              dtype=query_layer.dtype,
              device=query_layer.device,
          )
  
          torch.baddbmm(
              matmul_input_buffer,
              query_layer.transpose(0, 1),  # [b * np, sq, hn]
              key_layer.transpose(1, 2),  # [b * np, hn, sk]
              beta=0.0,
              alpha=1.0,
              out=matmul_result,
          )
  
          # change view to [b, np, sq, sk]
          attention_scores = matmul_result.view(*output_size)
  
          if self.scale_mask_softmax:
              self.scale_mask_softmax.scale = (
                  query_key_layer_scaling_coeff
              )
              attention_probs = self.scale_mask_softmax(
                  attention_scores, attention_mask.contiguous()
              )
          else:
              if not (attention_mask == 0).all():
                  # if auto-regressive, skip
                  attention_scores.masked_fill_(
                      attention_mask, -10000.0
                  )
              dtype = attention_scores.dtype
              attention_scores = attention_scores.float()
              attention_scores = (
                  attention_scores * query_key_layer_scaling_coeff
              )
  
              attention_probs = F.softmax(attention_scores, dim=-1)
  
              attention_probs = attention_probs.type(dtype)
  
          # =========================
          # Context layer. [sq, b, hp]
          # =========================
  
          # value_layer -> context layer.
          # [sk, b, np, hn] --> [b, np, sq, hn]
  
          # context layer shape: [b, np, sq, hn]
          output_size = (
              value_layer.size(0),
              value_layer.size(1),
              query_layer.size(0),
              value_layer.size(3),
          )
  
          # change view [sk, b * np, hn]
          value_layer = value_layer.view(
              output_size[0] * output_size[1], value_layer.size(2), -1
          )
  
          # change view [b * np, sq, sk]
          attention_probs = attention_probs.view(
              output_size[0] * output_size[1], output_size[2], -1
          )
  
          # matmul: [b * np, sq, hn]
          context_layer = torch.empty(
              output_size[0] * output_size[1],
              output_size[2],
              value_layer.size(-1),
              dtype=value_layer.dtype,
              device=query_layer.device,
          )
          torch.bmm(attention_probs, value_layer, out=context_layer)
  
          # change view [b, np, sq, hn]
          context_layer = context_layer.view(*output_size)
  
          # [b, np, sq, hn] --> [sq, b, np, hn]
          context_layer = context_layer.permute(2, 0, 1, 3).contiguous()
  
          # [sq, b, np, hn] --> [sq, b, hp]
          new_context_layer_shape = context_layer.size()[:-2] + (
              hidden_size_per_partition,
          )
          context_layer = context_layer.view(*new_context_layer_shape)
  
      outputs = (context_layer, present, attention_probs)
  
      return outputs
  
file_path: BigDL/python/llm/src/bigdl/llm/transformers/models/chatglm.py
imports: import math
  import torch
  import torch.utils.checkpoint
  import torch.nn.functional
  from typing import Optional, Tuple
  from bigdl.llm.transformers.models.utils import init_kv_cache, extend_kv_cache, append_kv_cache
label_desc: Use image node_id 3 for calling a global function with example usage: attention_fn(query_layer, key_layer, value_layer, attention_mask, hidden_size_per_partition, layer_id, layer_past, scaling_attention_score, use_cache) and returns: outputs