code: def mistral_attention_forward(
      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,
      padding_mask: Optional[torch.Tensor] = None,
  ) -> Tuple[
      torch.Tensor,
      Optional[torch.Tensor],
      Optional[Tuple[torch.Tensor]],
  ]:
      bsz, q_len, hidden_size = hidden_states.size()
      device = hidden_states.device
      # for flash attention
      original_dtype = hidden_states.dtype
  
      use_fuse_rope = should_use_fuse_rope(
          self, hidden_states, position_ids
      )
      enough_kv_room = is_enough_kv_cache_room_4_31(past_key_value)
      decoding_fast_path = use_decoding_fast_path(
          self.q_proj.qtype, use_fuse_rope, enough_kv_room, bsz * q_len
      )
  
      if decoding_fast_path:
          hidden_states = hidden_states.view(1, -1)
          kv_seq_len = past_key_value[0].shape[-2]
          cache_k = past_key_value[0]
          cache_v = past_key_value[1]
          import linear_q4_0
  
          (
              query_states,
              key_states,
              value_states,
          ) = linear_q4_0.forward_qkv(
              hidden_states,
              self.q_proj.weight,
              self.k_proj.weight,
              self.v_proj.weight,
              position_ids,
              cache_k,
              cache_v,
              self.q_proj.weight.qtype,
              kv_seq_len,
              self.head_dim,
          )
          kv_seq_len += 1
      else:
          query_states = self.q_proj(hidden_states)
          key_states = self.k_proj(hidden_states)
          value_states = self.v_proj(hidden_states)
  
          query_states = query_states.view(
              bsz, q_len, self.num_heads, self.head_dim
          ).transpose(1, 2)
          key_states = key_states.view(
              bsz, q_len, self.num_key_value_heads, self.head_dim
          ).transpose(1, 2)
          value_states = value_states.view(
              bsz, q_len, self.num_key_value_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 use_fuse_rope:
              (
                  query_states,
                  key_states,
              ) = apply_rotary_pos_emb_no_cache_xpu(
                  query_states, key_states, position_ids, "mistral"
              )
          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,
                  "mistral",
              )
  
          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 not enough_kv_room:
                  # allocate new
                  new_cache_k, new_cache_v = extend_kv_cache(
                      bsz,
                      self.num_key_value_heads,  # Support GQA
                      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_key_value_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 not self.training and not hidden_states.requires_grad:
          fsdp_flag = use_flash_attention(query_states, key_states)
      else:
          fsdp_flag = False
      if fsdp_flag:
          attention_dtype = (
              torch.float16
          )  # use fp16 for flash attention
      else:
          attention_dtype = original_dtype
  
      # repeat k/v heads if n_kv_heads < n_heads
      key_states = repeat_kv(key_states, self.num_key_value_groups).to(
          device, dtype=attention_dtype
      )
      value_states = repeat_kv(
          value_states, self.num_key_value_groups
      ).to(device, dtype=attention_dtype)
  
      if fsdp_flag:
          attn_output = F.scaled_dot_product_attention(
              query_states.to(dtype=attention_dtype),
              key_states,
              value_states,
              is_causal=True,
          )
          attn_weights = None
          attn_output = attn_output.transpose(1, 2).contiguous()
          attn_output = attn_output.reshape(bsz, q_len, hidden_size)
      else:
          attn_output, attn_weights = compute_attn_outputs_weights(
              query_states,
              key_states,
              value_states,
              bsz,
              q_len,
              kv_seq_len,
              self.num_heads,
              self.head_dim,
              self.hidden_size,
              attention_mask,
          )
  
      attn_output = self.o_proj(attn_output)
  
      if not output_attentions:
          attn_weights = None
  
      return (
          attn_output.to(original_dtype),
          attn_weights,
          past_key_value,
      )
  
file_path: BigDL/python/llm/src/bigdl/llm/transformers/models/mistral.py
imports: import math
  from typing import Optional, Tuple
  import torch
  from torch import nn
  import torch.nn.functional
  from bigdl.llm.utils.common import invalidInputError
  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, apply_rotary_pos_emb_no_cache_xpu
  from bigdl.llm.transformers.models.utils import is_enough_kv_cache_room_4_31, is_enough_kv_cache_room_4_36
  from bigdl.llm.transformers.low_bit_linear import SYM_INT4, FP8E5
  from bigdl.llm.transformers.models.utils import use_flash_attention
label_desc: Use image node_id 5 for calling a global function with example usage: mistral_attention_forward(hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache, padding_mask) and returns: attn_output,  attn_weights,  past_key_value