code: def llama_model_forward(
      self,
      input_ids: torch.LongTensor = None,
      attention_mask: Optional[torch.Tensor] = None,
      position_ids: Optional[torch.LongTensor] = None,
      past_key_values: Optional[List[torch.FloatTensor]] = None,
      inputs_embeds: Optional[torch.FloatTensor] = None,
      use_cache: Optional[bool] = None,
      output_attentions: Optional[bool] = None,
      output_hidden_states: Optional[bool] = None,
      return_dict: Optional[bool] = None,
  ) -> Union[Tuple, BaseModelOutputWithPast]:
      output_attentions = (
          output_attentions
          if output_attentions is not None
          else self.config.output_attentions
      )
      output_hidden_states = (
          output_hidden_states
          if output_hidden_states is not None
          else self.config.output_hidden_states
      )
      use_cache = (
          use_cache if use_cache is not None else self.config.use_cache
      )
  
      return_dict = (
          return_dict
          if return_dict is not None
          else self.config.use_return_dict
      )
  
      # retrieve input_ids and inputs_embeds
      if input_ids is not None and inputs_embeds is not None:
          raise ValueError(
              "You cannot specify both input_ids and inputs_embeds at the same time"
          )
      elif input_ids is not None:
          batch_size, seq_length = input_ids.shape
      elif inputs_embeds is not None:
          batch_size, seq_length, _ = inputs_embeds.shape
      else:
          raise ValueError(
              "You have to specify either input_ids or inputs_embeds"
          )
  
      infer_state = self.infer_state
      if infer_state.is_context_stage:
          past_key_values_length = 0
      else:
          past_key_values_length = infer_state.max_len_in_batch - 1
  
      seq_length_with_past = seq_length + past_key_values_length
  
      # NOTE: differentiate with prefill stage
      #       block_loc require different value-assigning method for two different stage
      # NOTE: differentiate with prefill stage
      #       block_loc require different value-assigning method for two different stage
      if infer_state.is_context_stage:
          infer_state.context_mem_index = (
              infer_state.cache_manager.alloc(
                  infer_state.total_token_num
              )
          )
          infer_state.init_block_loc(
              infer_state.block_loc,
              infer_state.seq_len,
              seq_length,
              infer_state.context_mem_index,
          )
      else:
          alloc_mem = infer_state.cache_manager.alloc_contiguous(
              batch_size
          )
          if alloc_mem is not None:
              infer_state.decode_is_contiguous = True
              infer_state.decode_mem_index = alloc_mem[0]
              infer_state.decode_mem_start = alloc_mem[1]
              infer_state.decode_mem_end = alloc_mem[2]
              infer_state.block_loc[
                  :, seq_length_with_past - 1
              ] = infer_state.decode_mem_index
          else:
              print(f" *** Encountered allocation non-contiguous")
              print(
                  f"    infer_state.cache_manager.max_len_in_batch: {infer_state.max_len_in_batch}"
              )
              infer_state.decode_is_contiguous = False
              alloc_mem = infer_state.cache_manager.alloc(batch_size)
              infer_state.decode_mem_index = alloc_mem
              infer_state.block_loc[
                  :, seq_length_with_past - 1
              ] = infer_state.decode_mem_index
  
      if position_ids is None:
          device = (
              input_ids.device
              if input_ids is not None
              else inputs_embeds.device
          )
          position_ids = torch.arange(
              past_key_values_length,
              seq_length + past_key_values_length,
              dtype=torch.long,
              device=device,
          )
          position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
      else:
          position_ids = position_ids.view(-1, seq_length).long()
  
      if inputs_embeds is None:
          inputs_embeds = self.embed_tokens(input_ids)
      # embed positions
      if attention_mask is None:
          attention_mask = torch.ones(
              (batch_size, seq_length_with_past),
              dtype=torch.bool,
              device=inputs_embeds.device,
          )
          padding_mask = None
      else:
          if 0 in attention_mask:
              padding_mask = attention_mask
          else:
              padding_mask = None
  
      attention_mask = self._prepare_decoder_attention_mask(
          attention_mask,
          (batch_size, seq_length),
          inputs_embeds,
          past_key_values_length,
      )
  
      hidden_states = inputs_embeds
  
      if self.gradient_checkpointing and self.training:
          raise NotImplementedError(
              "not implement gradient_checkpointing and training options "
          )
  
      if past_key_values_length == 0:
          position_cos = torch.index_select(
              self._cos_cached, 0, position_ids.view(-1)
          ).view(position_ids.view(-1).shape[0], -1)
          position_sin = torch.index_select(
              self._sin_cached, 0, position_ids.view(-1)
          ).view(position_ids.view(-1).shape[0], -1)
      else:
          position_cos = torch.index_select(
              self._cos_cached, 0, position_ids.view(-1)
          ).view(batch_size, -1)
          position_sin = torch.index_select(
              self._sin_cached, 0, position_ids.view(-1)
          ).view(batch_size, -1)
  
      # decoder layers
      all_hidden_states = () if output_hidden_states else None
      all_self_attns = () if output_attentions else None
      next_decoder_cache = () if use_cache else None
      infer_state.decode_layer_id = 0
      for idx, decoder_layer in enumerate(self.layers):
          if output_hidden_states:
              all_hidden_states += (hidden_states,)
  
          past_key_value = (
              past_key_values[idx]
              if past_key_values is not None
              else None
          )
  
          layer_outputs = decoder_layer(
              hidden_states,
              rotary_emb=(position_cos, position_sin),
              attention_mask=attention_mask,
              position_ids=position_ids,
              past_key_value=past_key_value,
              output_attentions=output_attentions,
              use_cache=use_cache,
              padding_mask=padding_mask,
              infer_state=infer_state,
          )
  
          hidden_states = layer_outputs[0]
          infer_state.decode_layer_id += 1
  
          if use_cache:
              next_decoder_cache += (
                  layer_outputs[2 if output_attentions else 1],
              )
  
          if output_attentions:
              all_self_attns += (layer_outputs[1],)
  
      hidden_states = self.norm(hidden_states)
  
      # add hidden states from the last decoder layer
      if output_hidden_states:
          all_hidden_states += (hidden_states,)
  
      infer_state.is_context_stage = False
      infer_state.start_loc = infer_state.start_loc + torch.arange(
          0, batch_size, dtype=torch.int32, device="cuda"
      )
      infer_state.seq_len += 1
      infer_state.max_len_in_batch += 1
  
      next_cache = next_decoder_cache if use_cache else None
      if not return_dict:
          return tuple(
              v
              for v in [
                  hidden_states,
                  next_cache,
                  all_hidden_states,
                  all_self_attns,
              ]
              if v is not None
          )
      return BaseModelOutputWithPast(
          last_hidden_state=hidden_states,
          past_key_values=next_cache,
          hidden_states=all_hidden_states,
          attentions=all_self_attns,
      )
  
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 3 for calling a global function with example usage: llama_model_forward(input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict) and returns: BaseModelOutputWithPast, tuple