code: def __init__(
      self,
      hidden_size: int,
      num_heads: int,
  ):
      super().__init__()
      self.hidden_size = hidden_size
      self.num_heads = num_heads
      self.head_dim = hidden_size // num_heads
  
      if (self.head_dim * num_heads) != self.hidden_size:
          raise ValueError(
              f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
              f" and `num_heads`: {num_heads})."
          )
  
      self.qk_bmm = BMM_S8T_S8N_F32T(1.0)
      self.pv_bmm = BMM_S8T_S8N_S8T(1.0)
  
      self.k_proj = W8A8B8O8Linear(hidden_size, hidden_size)
      self.v_proj = W8A8B8O8Linear(hidden_size, hidden_size)
      self.q_proj = W8A8B8O8Linear(hidden_size, hidden_size)
      self.o_proj = W8A8BFP32OFP32Linear(hidden_size, hidden_size)
  
      self.register_buffer("q_output_scale", torch.tensor([1.0]))
      self.register_buffer("k_output_scale", torch.tensor([1.0]))
      self.register_buffer("v_output_scale", torch.tensor([1.0]))
      self.register_buffer("q_rotary_output_scale", torch.tensor([1.0]))
      self.register_buffer("k_rotary_output_scale", torch.tensor([1.0]))
      self.register_buffer("out_input_scale", torch.tensor([1.0]))
      self.register_buffer("attn_input_scale", torch.tensor([1.0]))
  
      self._init_rope()
      self.num_key_value_heads = num_heads
  
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 1 to create a new LLamaSmoothquantAttention object from inherited base classes: nn with example: obj = LLamaSmoothquantAttention(hidden_size, num_heads)