class_docstr: This class holds forwards for llama inference.
  We intend to replace the forward methods for LlamaModel, LlamaDecoderLayer, and LlamaAttention for LlamaForCausalLM.
class_docstr_tok: ["This","class","holds","forwards","for","llama","inference",".","We","intend","to","replace","the","forward","methods","for","LlamaModel",",","LlamaDecoderLayer",",","and","LlamaAttention","for","LlamaForCausalLM","."]
code: def llama_flash_attn_kvcache_forward(
      self: LlamaAttention,
      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,
      infer_state: Optional[BatchInferState] = None,
  ) -> Tuple[
      torch.Tensor,
      Optional[torch.Tensor],
      Optional[Tuple[torch.Tensor]],
  ]:
      assert (
          use_cache is True
      ), "use_cache should be set to True using this llama attention"
  
      bsz, q_len, _ = hidden_states.size()
  
      # NOTE might think about better way to handle transposed k and v
      # key_states            [bs, seq_len, num_heads, head_dim/embed_size_per_head]
      # key_states_transposed [bs, num_heads, seq_len, head_dim/embed_size_per_head]
  
      query_states = self.q_proj(hidden_states).view(
          bsz, q_len, self.num_heads, self.head_dim
      )
      key_states = self.k_proj(hidden_states).view(
          bsz, q_len, self.num_key_value_heads, self.head_dim
      )
      value_states = self.v_proj(hidden_states).view(
          bsz, q_len, self.num_key_value_heads, self.head_dim
      )
  
      # NOTE might want to revise
      #   need some way to record the length of past key values cache
      #   since we won't return past_key_value_cache right now
  
      cos, sin = infer_state.position_cos, infer_state.position_sin
  
      llama_rotary_embedding_fwd(
          query_states.view(-1, self.num_heads, self.head_dim), cos, sin
      )
      llama_rotary_embedding_fwd(
          key_states.view(-1, self.num_key_value_heads, self.head_dim),
          cos,
          sin,
      )
  
      query_states = query_states.reshape(
          -1, self.num_heads, self.head_dim
      )
      key_states = key_states.reshape(
          -1, self.num_key_value_heads, self.head_dim
      )
      value_states = value_states.reshape(
          -1, self.num_key_value_heads, self.head_dim
      )
  
      if infer_state.is_context_stage:
          # first token generation
          # copy key and value calculated in current step to memory manager
          copy_kv_to_mem_cache(
              infer_state.decode_layer_id,
              key_states,
              value_states,
              infer_state.context_mem_index,
              infer_state.cache_manager,
          )
          attn_output = torch.empty_like(query_states)
  
          llama_triton_context_attention(
              query_states,
              key_states,
              value_states,
              attn_output,
              infer_state,
              num_key_value_groups=self.num_key_value_groups,
          )
      else:
          if infer_state.decode_is_contiguous:
              # if decode is contiguous, then we copy to key cache and value cache in cache manager directly
              cache_k = infer_state.cache_manager.key_buffer[
                  infer_state.decode_layer_id
              ][
                  infer_state.decode_mem_start : infer_state.decode_mem_end,
                  :,
                  :,
              ]
              cache_v = infer_state.cache_manager.value_buffer[
                  infer_state.decode_layer_id
              ][
                  infer_state.decode_mem_start : infer_state.decode_mem_end,
                  :,
                  :,
              ]
              cache_k.copy_(key_states)
              cache_v.copy_(value_states)
          else:
              # if decode is not contiguous, use triton kernel to copy key and value cache
              # k, v shape: [batch_size, num_heads, head_dim/embed_size_per_head
              copy_kv_to_mem_cache(
                  infer_state.decode_layer_id,
                  key_states,
                  value_states,
                  infer_state.decode_mem_index,
                  infer_state.cache_manager,
              )
  
          if HAS_LIGHTLLM_KERNEL:
              attn_output = torch.empty_like(query_states)
              llama_triton_token_attention(
                  query_states,
                  attn_output,
                  infer_state,
                  num_key_value_groups=self.num_key_value_groups,
              )
          else:
              self.num_heads // self.num_key_value_heads
              cache_k = infer_state.cache_manager.key_buffer[
                  infer_state.decode_layer_id
              ]
              cache_v = infer_state.cache_manager.value_buffer[
                  infer_state.decode_layer_id
              ]
  
              query_states = query_states.view(
                  bsz, -1, self.num_heads, self.head_dim
              )
              copy_cache_k = cache_k.view(
                  bsz, -1, self.num_key_value_heads, self.head_dim
              )
              copy_cache_v = cache_v.view(
                  bsz, -1, self.num_key_value_heads, self.head_dim
              )
  
              attn_output = flash_attn_with_kvcache(
                  q=query_states,
                  k_cache=copy_cache_k,
                  v_cache=copy_cache_v,
                  softmax_scale=1 / math.sqrt(self.head_dim),
                  causal=True,
              )
  
      attn_output = attn_output.view(bsz, q_len, self.hidden_size)
  
      attn_output = self.o_proj(attn_output)
  
      # return past_key_value as None
      return attn_output, None, None
  
file_path: ColossalAI/colossalai/legacy/inference/tensor_parallel/modeling/llama.py
imports: import math
  from typing import List, Optional, Tuple
  import torch
  from transformers.modeling_outputs import BaseModelOutputWithPast
  from transformers.models.llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaModel
  from colossalai.inference.tensor_parallel.batch_infer_state import BatchInferState
  from colossalai.kernel.triton import llama_context_attn_fwd, token_attention_fwd
  from colossalai.kernel.triton.token_attention_kernel import Llama2TokenAttentionForwards
  from ._utils import copy_kv_to_mem_cache
label_desc: Use image node_id 3 for calling the LlamaInferenceForwards obj's underlying member method code with example usage: obj.llama_flash_attn_kvcache_forward(hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache, infer_state) and returns: attn_output,  None,  None