class_docstr: Engine class for tensor parallel inference.
  
  Args:
      model (Module): original model, e.g. huggingface CausalLM
      shard_config (ShardConfig): The config for sharding original model
      max_batch_size (int): maximum batch size
      max_input_len (int): maximum input length of sequence
      max_output_len (int): maximum output length of output tokens
      dtype (torch.dtype): datatype used to init KV cache space
      device (str): device the KV cache of engine to be initialized on
  
  Examples:
      >>> # define model and shard config for your inference
      >>> model = ...
      >>> generate_kwargs = ...
      >>> shard_config = ShardConfig(enable_tensor_parallelism=True, extra_kwargs={"inference_only": True})
      >>> infer_engine = TPInferEngine(model, shard_config, MAX_BATCH_SIZE, MAX_INPUT_LEN, MAX_OUTPUT_LEN)
      >>> outputs = infer_engine.generate(input_ids, **generate_kwargs)
class_docstr_tok: ["Engine","class","for","tensor","parallel","inference",".","Args",":","model","(","Module",")",":","original","model",",","e.g",".","huggingface","CausalLM","shard_config","(","ShardConfig",")",":","The","config","for","sharding","original","model","max_batch_size","(","int",")",":","maximum","batch","size","max_input_len","(","int",")",":","maximum","input","length","of","sequence","max_output_len","(","int",")",":","maximum","output","length","of","output","tokens","dtype","(","torch.dtype",")",":","datatype","used","to","init","KV","cache","space","device","(","str",")",":","device","the","KV","cache","of","engine","to","be","initialized","on","Examples",":",">",">",">","#","define","model","and","shard","config","for","your","inference",">",">",">","model","=","...",">",">",">","generate_kwargs","=","...",">",">",">","shard_config","=","ShardConfig","(","enable_tensor_parallelism=True",",","extra_kwargs=","{","``","inference_only","''",":","True","}",")",">",">",">","infer_engine","=","TPInferEngine","(","model",",","shard_config",",","MAX_BATCH_SIZE",",","MAX_INPUT_LEN",",","MAX_OUTPUT_LEN",")",">",">",">","outputs","=","infer_engine.generate","(","input_ids",",","*","*","generate_kwargs",")"]
code: def __init__(
      self,
      model: nn.Module,
      shard_config: ShardConfig,
      max_batch_size: int,
      max_input_len: int,
      max_output_len: int,
      dtype: torch.dtype = torch.float16,
      device: str = "cuda",
  ) -> None:
      self.max_batch_size = max_batch_size
      self.max_input_len = max_input_len
      self.max_output_len = max_output_len
      self.max_total_token_num = self.max_batch_size * (
          self.max_input_len + self.max_output_len
      )
      # Constraints relatable with specs of devices and model
      # This may change into an optional arg in the future
      assert (
          self.max_batch_size <= 64
      ), "Max batch size exceeds the constraint"
      assert (
          self.max_input_len + self.max_output_len <= 4096
      ), "Max length exceeds the constraint"
  
      self.dtype = dtype
  
      self.head_dim = (
          model.config.hidden_size // model.config.num_attention_heads
      )
      self.head_num = model.config.num_attention_heads
      num_hidden_layers = (
          model.config.num_hidden_layers
          if hasattr(model.config, "num_hidden_layers")
          else model.config.num_layers
      )
      self.layer_num = num_hidden_layers
  
      self.multi_query_group_num = model.config.num_attention_heads
      # default to attention_heads
      if hasattr(model.config, "multi_query_attention"):
          self.multi_query_attention = getattr(
              model.config, "multi_query_attention"
          )
  
      if hasattr(model.config, "multi_query_group_num"):
          self.multi_query_group_num = getattr(
              model.config, "multi_query_group_num"
          )
  
      if hasattr(model.config, "num_key_value_heads"):
          self.multi_query_group_num = getattr(
              model.config, "num_key_value_heads"
          )
  
      self.tp_size = (
          -1
      )  # to be set with given shard config in self.prepare_shard_config
      self.cache_manager = None
  
      self.max_dq_buffer_size = 1
      self.max_inner_outer_dim = 1
      self.gptq_temp_state_buffer = None
      self.gptq_temp_dq_buffer = None
      self.bits = -1
      self.use_act_order = False
  
      self.shard_config = shard_config
      self.model = None
      self.cache = {}
  
      # optimize the original model by sharding with ShardFormer
      self._optimize_model(model=model.to(device))
  
imports: from typing import Any, Callable, List, Optional, Union
  import torch
  import torch.nn
  from transformers import BloomForCausalLM, LlamaForCausalLM
  from transformers.generation import GenerationConfig
  from transformers.generation.stopping_criteria import StoppingCriteriaList
  from transformers.tokenization_utils_base import BatchEncoding
  from colossalai.shardformer import ShardConfig, ShardFormer
  from colossalai.shardformer.policies.auto_policy import get_autopolicy
  from .batch_infer_state import BatchInferState
  from .kvcache_manager import MemoryManager
label_desc: Use image node_id 1 to create a new TPInferEngine object  with example: obj = TPInferEngine(model, shard_config, max_batch_size, max_input_len, max_output_len, dtype, device)