code: def baichuan_attention_forward_7b(
      self,
      hidden_states: torch.Tensor,
      attention_mask: Optional[torch.Tensor] = None,
      position_ids: Optional[torch.LongTensor] = None,
      past_key_value: Optional[Tuple[torch.Tensor]] = None,
      output_attentions: bool = False,
      use_cache: bool = False,
  ) -> Tuple[
      torch.Tensor,
      Optional[torch.Tensor],
      Optional[Tuple[torch.Tensor]],
  ]:
      bsz, q_len, _ = hidden_states.size()
      device = hidden_states.device
  
      proj = self.W_pack(hidden_states)
      proj = (
          proj.unflatten(-1, (3, self.hidden_size))
          .unsqueeze(0)
          .transpose(0, -2)
          .squeeze(-2)
      )
      # batch_size x source_len x hidden_size
      query_states = (
          proj[0]
          .view(bsz, q_len, self.num_heads, self.head_dim)
          .transpose(1, 2)
      )
      # batch_size x target_len x head_size
      key_states = (
          proj[1]
          .view(bsz, q_len, self.num_heads, self.head_dim)
          .transpose(1, 2)
      )
      # batch_size x source_len x hidden_size
      value_states = (
          proj[2]
          .view(bsz, q_len, self.num_heads, self.head_dim)
          .transpose(1, 2)
      )
  
      kv_seq_len = key_states.shape[-2]
      if past_key_value is not None:
          kv_seq_len += past_key_value[0].shape[-2]
      if query_states.device.type == "xpu" and not (
          self.training and query_states.requires_grad
      ):
          query_states, key_states = apply_rotary_pos_emb_no_cache_xpu(
              query_states, key_states, position_ids, "baichuan"
          )
      else:
          cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
          query_states, key_states = apply_rotary_pos_emb(
              query_states,
              key_states,
              cos,
              sin,
              position_ids,
              "baichuan",
          )
      # [bsz, nh, t, hd]
  
      # if past_key_value is not None:
      #     # reuse k, v, self_attention
      #     key_states = torch.cat([past_key_value[0], key_states], dim=2)
      #     value_states = torch.cat([past_key_value[1], value_states], dim=2)
      if past_key_value is not None:
          # reuse k, v, self_attention
          cache_k = past_key_value[0]
          cache_v = past_key_value[1]
          if cache_k.stride()[1] <= cache_k.size(2) * cache_k.size(3):
              # allocate new
              new_cache_k, new_cache_v = extend_kv_cache(
                  bsz,
                  self.num_heads,
                  self.head_dim,
                  cache_k.size(2),
                  kv_seq_len + KV_CACHE_ALLOC_BLOCK_LENGTH,
                  dtype=cache_k.dtype,
                  device=device,
              )
              new_cache_k[:] = cache_k
              new_cache_v[:] = cache_v
              cache_k = new_cache_k
              cache_v = new_cache_v
  
          key_states, value_states = append_kv_cache(
              cache_k, cache_v, key_states, value_states
          )
  
      elif use_cache:
          max_cache_length = kv_seq_len + KV_CACHE_ALLOC_BLOCK_LENGTH
          new_key_states, new_value_states = init_kv_cache(
              bsz,
              self.num_heads,
              self.head_dim,
              kv_seq_len,
              max_cache_length,
              dtype=key_states.dtype,
              device=device,
          )
          new_key_states[:] = key_states
          new_value_states[:] = value_states
          key_states = new_key_states
          value_states = new_value_states
  
      past_key_value = (key_states, value_states) if use_cache else None
  
      if xops is not None and self.training:
          attn_weights = None
          query_states = query_states.transpose(1, 2)
          key_states = key_states.transpose(1, 2)
          value_states = value_states.transpose(1, 2)
          attn_output = xops.memory_efficient_attention(
              query_states,
              key_states,
              value_states,
              attn_bias=xops.LowerTriangularMask(),
          )
      else:
          if attention_mask is not None:
              if attention_mask.dtype == torch.bool:
                  attention_mask.masked_fill_(
                      attention_mask.logical_not(), float("-inf")
                  )
  
          scaling_factor = 1 / math.sqrt(query_states.size(-1))
          attn_output = torch.matmul(
              query_states * scaling_factor,
              key_states.transpose(-2, -1),
          )
          if attention_mask is not None:
              attn_output += attention_mask
          attn_output = torch.softmax(attn_output, -1)
          attn_output = torch.matmul(attn_output, value_states)
  
          attn_output = attn_output.transpose(1, 2)
      attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
      attn_output = self.o_proj(attn_output)
  
      if not output_attentions:
          attn_weights = None
  
      return attn_output, attn_weights, past_key_value
  
file_path: BigDL/python/llm/src/bigdl/llm/transformers/models/baichuan2.py
imports: import math
  from typing import Optional, Tuple
  import torch
  import torch.utils.checkpoint
  from torch.nn import functional
  from bigdl.llm.ggml.quantize import ggml_tensor_qtype
  from bigdl.llm.transformers.models.utils import init_kv_cache, extend_kv_cache, append_kv_cache
  from bigdl.llm.transformers.models.utils import apply_rotary_pos_emb
  from bigdl.llm.transformers.models.utils import apply_rotary_pos_emb_no_cache_xpu
  from bigdl.llm.transformers.models.utils import mlp_fusion_check
  from transformers.utils import logging
label_desc: Use image node_id 3 for calling a global function with example usage: baichuan_attention_forward_7b(hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache) and returns: attn_output,  attn_weights,  past_key_value