code: def forward(
      self,
      hidden_states: torch.Tensor,
      rotary_emb: Tuple[torch.Tensor] = None,
      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.LongTensor] = None,
      infer_state: Optional[BatchInferState] = None,
  ) -> Tuple[
      torch.Tensor,
      Optional[torch.Tensor],
      Optional[Tuple[torch.Tensor]],
  ]:
      bsz, q_len, _ = hidden_states.size()
  
      query_states = self.q_proj(hidden_states)
      key_states = self.k_proj(hidden_states)
      value_states = self.v_proj(hidden_states)
  
      cos = rotary_emb[0]
      sin = rotary_emb[1]
  
      int8_rotary_embedding_fwd(
          query_states.view(-1, self.num_heads, self.head_dim),
          cos,
          sin,
          self.q_output_scale.item(),
          self.q_rotary_output_scale.item(),
      )
      int8_rotary_embedding_fwd(
          key_states.view(-1, self.num_heads, self.head_dim),
          cos,
          sin,
          self.k_output_scale.item(),
          self.k_rotary_output_scale.item(),
      )
  
      def _copy_kv_to_mem_cache(
          layer_id,
          key_buffer,
          value_buffer,
          context_mem_index,
          mem_manager,
      ):
          copy_kv_cache_to_dest(
              key_buffer,
              context_mem_index,
              mem_manager.key_buffer[layer_id],
          )
          copy_kv_cache_to_dest(
              value_buffer,
              context_mem_index,
              mem_manager.value_buffer[layer_id],
          )
          return
  
      query_states = query_states.view(
          -1, self.num_heads, self.head_dim
      )
      key_states = key_states.view(-1, self.num_heads, self.head_dim)
      value_states = value_states.view(
          -1, self.num_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)
  
          smooth_llama_context_attn_fwd(
              query_states,
              key_states,
              value_states,
              attn_output,
              self.q_rotary_output_scale.item(),
              self.k_rotary_output_scale.item(),
              self.v_output_scale.item(),
              self.out_input_scale.item(),
              infer_state.start_loc,
              infer_state.seq_len,
              q_len,
          )
  
      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,
              )
  
          # (batch_size, seqlen, nheads, headdim)
          attn_output = torch.empty_like(query_states)
  
          smooth_token_attention_fwd(
              query_states,
              infer_state.cache_manager.key_buffer[
                  infer_state.decode_layer_id
              ],
              infer_state.cache_manager.value_buffer[
                  infer_state.decode_layer_id
              ],
              attn_output,
              self.q_rotary_output_scale.item(),
              self.k_rotary_output_scale.item(),
              self.v_output_scale.item(),
              self.out_input_scale.item(),
              infer_state.block_loc,
              infer_state.start_loc,
              infer_state.seq_len,
              infer_state.max_len_in_batch,
          )
  
      attn_output = attn_output.view(
          bsz, q_len, self.num_heads * self.head_dim
      )
      attn_output = self.o_proj(attn_output)
  
      return attn_output, None, None
  
imports: import math
  import os
  import types
  from collections import defaultdict
  from functools import partial
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.nn
  import torch.nn.functional
  from torch_int.nn.bmm import BMM_S8T_S8N_F32T, BMM_S8T_S8N_S8T
  from transformers import PreTrainedModel
  from transformers.modeling_outputs import BaseModelOutputWithPast
  from transformers.models.llama.configuration_llama import LlamaConfig
  from transformers.models.llama.modeling_llama import LLAMA_INPUTS_DOCSTRING, LlamaAttention, LlamaDecoderLayer, LlamaMLP, LlamaRotaryEmbedding, repeat_kv, rotate_half
  from transformers.utils import add_start_docstrings_to_model_forward
  from colossalai.inference.tensor_parallel.batch_infer_state import BatchInferState
  from colossalai.kernel.triton import copy_kv_cache_to_dest, int8_rotary_embedding_fwd, smooth_llama_context_attn_fwd, smooth_token_attention_fwd
  from .base_model import BaseSmoothForCausalLM
  from .linear import W8A8B8O8Linear, W8A8BFP32O32LinearSiLU, W8A8BFP32OFP32Linear
label_desc: Use image node_id 5 for calling the LLamaSmoothquantAttention obj's underlying member method code with example usage: obj.forward(hidden_states, rotary_emb, attention_mask, position_ids, past_key_value, output_attentions, use_cache, padding_mask, infer_state) and returns: attn_output,  None,  None