code: def llama_triton_context_attention(
      query_states,
      key_states,
      value_states,
      attn_output,
      infer_state,
      num_key_value_groups=1,
  ):
      # if num_key_value_groups == 1:
      if HAS_LIGHTLLM_KERNEL is False:
          llama_context_attn_fwd(
              query_states,
              key_states,
              value_states,
              attn_output,
              infer_state.start_loc,
              infer_state.seq_len,
              # infer_state.cache_manager.past_key_values_length,
              infer_state.max_len_in_batch,
          )
      else:
          lightllm_llama_context_attention_fwd(
              query_states,
              key_states,
              value_states,
              attn_output,
              infer_state.start_loc,
              infer_state.seq_len,
              # infer_state.cache_manager.past_key_values_length,
              infer_state.max_len_in_batch,
          )
  
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 a global function with example usage: llama_triton_context_attention(query_states, key_states, value_states, attn_output, infer_state, num_key_value_groups) without return types