class_docstr: CaiInferEngine is a class that handles the pipeline parallel inference.
  
  Args:
      tp_size (int): the size of tensor parallelism.
      pp_size (int): the size of pipeline parallelism.
      model (`nn.Module`): the model not in pipeline style, and will be modified with `ShardFormer`.
      model_policy (`colossalai.shardformer.policies.base_policy.Policy`): the policy to shardformer model.
      micro_batch_size (int): the micro batch size.
      micro_batch_buffer_size (int): the buffer size for micro batch. Normally, it should be the same as the number of pipeline stages.
      max_batch_size (int): the maximum batch size.
      max_input_len (int): the maximum input length.
      max_output_len (int): the maximum output length.
  
  Example:
  
  ```python
  from colossalai.inference import InferEngine
  from colossalai.inference.pipeline.policies import LlamaModelInferPolicy
  import colossalai
  from transformers import LlamaForCausalLM, LlamaTokenizer
  
  colossalai.launch_from_torch(config={})
  
  model = LlamaForCausalLM.from_pretrained("your_path_to_model")
  tokenizer = LlamaTokenizer.from_pretrained("/home/lczyh/share/models/llama-7b-hf")
  # assume the model is infered with 2 pipeline stages
  inferengine = CaiInferEngine(pp_size=2, model=model, model_policy=LlamaModelInferPolicy())
  
  input = ["Introduce a landmark in China ","Introduce a landmark in China "]
  data = tokenizer(input, return_tensors='pt')
  output = inferengine.inference([data.to('cuda').data])
  
  ```
code: def _init_manager(
      self,
      model,
      max_batch_size: int,
      max_input_len: int,
      max_output_len: int,
  ) -> None:
      max_total_token_num = max_batch_size * (
          max_input_len + max_output_len
      )
      head_dim = (
          model.config.hidden_size // model.config.num_attention_heads
      )
      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
      )
      layer_num = num_hidden_layers // self.pp_size
  
      cache_manager = MemoryManager(
          max_total_token_num, self.dtype, head_num, head_dim, layer_num
      )
      return cache_manager
  
file_path: ColossalAI/colossalai/legacy/inference/hybridengine/engine.py
imports: import torch
  import torch.distributed
  import torch.nn
  from transformers.tokenization_utils_base import BatchEncoding
  from colossalai.cluster import ProcessGroupMesh
  from colossalai.pipeline.schedule.generate import GenerateSchedule
  from colossalai.pipeline.stage_manager import PipelineStageManager
  from colossalai.shardformer import ShardConfig, ShardFormer
  from colossalai.shardformer.policies.base_policy import Policy
  from ..pipeline.microbatch_manager import MicroBatchManager
  from ..tensor_parallel.kvcache_manager import MemoryManager
label_desc: Use image node_id 4 for calling the CaiInferEngine obj's underlying member method code with example usage: obj._init_manager(model, max_batch_size, max_input_len, max_output_len) and returns: cache_manager