code: def __init__(
      self,
      num_blocks: int,
      num_layers: int,
      num_heads: int,
      head_size: int,
      repeat_slots: bool,
      dtype: torch.dtype,
      device: torch.device,
  ):
      self.block_size = BLOCK_SIZE
      self.num_blocks = num_blocks
      self.repeat_slots = repeat_slots
  
      element_size = torch.tensor([], dtype=dtype).element_size()
      x = self.block_size // element_size
  
      self.kv_cache = [
          (
              torch.empty(
                  (
                      num_blocks,
                      num_heads,
                      head_size // x,
                      self.block_size,
                      x,
                  ),
                  dtype=dtype,
                  device=device,
              ),
              torch.empty(
                  (num_blocks, num_heads, head_size, self.block_size),
                  dtype=dtype,
                  device=device,
              ),
          )
          for _ in range(num_layers)
      ]
      self.free_block_mask = torch.ones(
          num_blocks, dtype=torch.int32, device="cpu"
      )
      self.slots = torch.arange(
          0, num_blocks * self.block_size, dtype=torch.int32
      ).view(num_blocks, self.block_size)
  
file_path: text-generation-inference/server/text_generation_server/models/cache_manager.py
imports: import math
  import torch
  from typing import Optional, List, Tuple
label_desc: Use image node_id 1 to create a new CacheManager object  with example: obj = CacheManager(num_blocks, num_layers, num_heads, head_size, repeat_slots, dtype, device)