class_docstr: Manage token block indexes and allocate physical memory for key and value cache
  
  Args:
      size: maximum token number used as the size of key and value buffer
      dtype: data type of cached key and value
      head_num: number of heads the memory manager is responsible for
      head_dim: embedded size per head
      layer_num: the number of layers in the model
      device: device used to store the key and value cache
code: def __init__(
      self,
      size: int,
      dtype: torch.dtype,
      head_num: int,
      head_dim: int,
      layer_num: int,
      device: torch.device = torch.device("cuda"),
  ):
      self.logger = logging.get_logger(__name__)
      self.available_size = size
      self.max_len_in_batch = 0
      self._init_mem_states(size, device)
      self._init_kv_buffers(
          size, device, dtype, head_num, head_dim, layer_num
      )
  
file_path: ColossalAI/colossalai/legacy/inference/tensor_parallel/kvcache_manager.py
imports: import torch
  from transformers.utils import logging
label_desc: Use image node_id 1 to create a new MemoryManager object  with example: obj = MemoryManager(size, dtype, head_num, head_dim, layer_num, device)