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
docstr: Initialize tensors used to manage memory states
code: def _init_mem_states(self, size, device):
      """Initialize tensors used to manage memory states"""
      self.mem_state = torch.ones(
          (size,), dtype=torch.bool, device=device
      )
      self.mem_cum_sum = torch.empty(
          (size,), dtype=torch.int32, device=device
      )
      self.indexes = torch.arange(
          0, size, dtype=torch.long, device=device
      )
  
file_path: colossalai/colossalai/inference/kv_cache/kvcache_manager.py
imports: import torch
  from transformers.utils import logging
label_desc: Use image node_id 2 for calling the MemoryManager obj's underlying member method code with example usage: obj._init_mem_states(size, device) without return types