code: def LlamaAttention_fast_forward_inference(
      self,
      hidden_states: torch.Tensor,
      past_key_value: Optional[Tuple[torch.Tensor]],
      position_ids,
  ):
      """
      https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L406
      Fast inference using KV cache.
      QK^T can be computed in 4 chunks
  
      [Q, q] @ [K, k].T where q, k are the new tokens.
      [QK^T, Qk^T]
      [qK^T, qk^T]
  
      Since the attention mask wipes Qk^T, we just get
      [QK^T,    0]
      [qK^T, qk^T]
  
      Since softmax is row-wise, we get
      softmax([QK^T,    0])
      softmax([qK^T, qk^T])
  
      We then multiply by   [V]
                            [v]
      softmax([QK^T,    0]) [softmax(QK^T)V] *
      softmax([qK^T, qk^T]) [softmax([qK^T, qk^T]) @ [V, v]]
  
      But notice * [softmax(QK^T)V] is just the last attention.
      We just need to compute the last final row.
  
      This means we can pass in a row of Q, but we need to
      remember K and V, which are called the KV cache.
      """
      Xn = hidden_states
      bsz, _, _ = hidden_states.size()
      K1, V1 = past_key_value
  
      Wq = self.q_proj.weight
      Wk = self.k_proj.weight
      Wv = self.v_proj.weight
      Wo = self.o_proj.weight
  
      n_heads = self.num_heads
      n_groups = self.num_key_value_groups
      n_kv_heads = self.num_key_value_heads
      head_dim = self.head_dim
      assert n_kv_heads * n_groups == n_heads
  
      Qn, Kn, Vn = original_apply_qkv(self, Xn)
      Qn = Qn.view(bsz, 1, n_heads, head_dim).transpose(1, 2)
      Kn = Kn.view(bsz, 1, n_kv_heads, head_dim).transpose(1, 2)
      Vn = Vn.view(bsz, 1, n_kv_heads, head_dim).transpose(1, 2)
  
      kv_seq_len = K1.shape[-2] + 1
      cos, sin = self.rotary_emb(Vn, seq_len=kv_seq_len)
      Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids)
  
      # New KV cache
      Kn = torch.cat([K1, Kn], dim=2)
      Vn = torch.cat([V1, Vn], dim=2)
  
      # Grouped query attention
      if n_groups != 1:
          _, _, cached_len, _ = Kn.shape
          Knn = Kn[:, :, None, :, :].expand(
              bsz, n_kv_heads, n_groups, cached_len, head_dim
          )
          Vnn = Vn[:, :, None, :, :].expand(
              bsz, n_kv_heads, n_groups, cached_len, head_dim
          )
          Knn = Knn.reshape(bsz, n_heads, cached_len, head_dim)
          Vnn = Vnn.reshape(bsz, n_heads, cached_len, head_dim)
      else:
          Knn, Vnn = Kn, Vn
  
      # Attention
      A = torch.matmul(Qn, Knn.transpose(2, 3))
      A *= 1.0 / (self.head_dim**0.5)
      A = torch.nn.functional.softmax(
          A, dim=-1, dtype=torch.float32
      ).to(A.dtype)
      A = torch.matmul(A, Vnn)
      A = A.transpose(1, 2)
      A = A.reshape(bsz, 1, self.hidden_size)
      A = original_apply_o(self, A)
      return A, (Kn, Vn)
  
imports: import torch
  from typing import Optional, Tuple, List, Union
  from torch.nn.functional import scaled_dot_product_attention
  from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
  from transformers.models.llama.modeling_llama import logger, BaseModelOutputWithPast, CausalLMOutputWithPast
  from ..kernels import *
  from ._utils import *
  from ._utils import __version__
  from transformers.models.llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaModel, LlamaForCausalLM
  from peft import PeftModelForCausalLM
  import gc
  import peft
  import bitsandbytes
  import numpy
  import types
  from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, AutoConfig
  from transformers import set_seed
  from peft import LoraConfig, TaskType, get_peft_model
label_desc: Use image node_id 3 for calling a global function with example usage: LlamaAttention_fast_forward_inference(hidden_states, past_key_value, position_ids) and returns: A,  Kn,  Vn