code: def token_attn(
      q,
      k,
      v,
      attn_out,
      kv_cache_loc,
      kv_cache_start_loc,
      kv_cache_seq_len,
      max_len_in_batch,
      other_kv_index,
  ):
      total_token_num = k.shape[0]
      batch_size, head_num, head_dim = q.shape
      calcu_shape1 = (batch_size, head_num, head_dim)
      att_m_tensor = torch.empty(
          (head_num, total_token_num), dtype=q.dtype, device="cuda"
      )
  
      lightllm_llama_token_att_fwd(
          q,
          k,
          att_m_tensor,
          kv_cache_loc,
          kv_cache_start_loc,
          kv_cache_seq_len,
          max_len_in_batch,
      )
  
      if triton.__version__ == "2.0.0":
          prob = torch.empty_like(att_m_tensor)
          lightllm_llama_token_softmax_fwd(
              att_m_tensor,
              kv_cache_start_loc,
              kv_cache_seq_len,
              prob,
              max_len_in_batch,
          )
          att_m_tensor = None
  
          lightllm_llama_token_att_fwd2(
              prob,
              v,
              attn_out.view(calcu_shape1),
              kv_cache_loc,
              kv_cache_start_loc,
              kv_cache_seq_len,
              max_len_in_batch,
          )
  
          prob = None
          return
  
      elif triton.__version__ >= "2.1.0":
          Llama2TokenAttentionForwards.token_softmax_reducev_fwd(
              att_m_tensor,
              v,
              attn_out.view(calcu_shape1),
              kv_cache_loc,
              kv_cache_start_loc,
              kv_cache_seq_len,
              max_len_in_batch,
              other_kv_index,
          )
      else:
          raise Exception("not support triton version")
  
imports: import torch
label_desc: Use image node_id 3 for calling the Llama2TokenAttentionForwards obj's underlying member method code with example usage: obj.token_attn(q, k, v, attn_out, kv_cache_loc, kv_cache_start_loc, kv_cache_seq_len, max_len_in_batch, other_kv_index) without return types