code: def allocate(
      self,
      needed_blocks_slots: List[Tuple[int, int]],
      blocks: int,
      max_blocks: int,
      device: torch.device,
  ):
      # Get free blocks indices by finding values in mask that are not set to 0
      free_block_indices = self.free_block_mask.nonzero()
      assert (
          len(free_block_indices) >= blocks
      ), f"Out of available cache blocks: asked {blocks}, only {len(free_block_indices)} free blocks"
  
      # Slice by the number of required blocks
      block_indices = free_block_indices[:blocks]
      block_indices = block_indices.flatten()
  
      # Padded block tables
      block_tables_tensor = torch.zeros(
          (len(needed_blocks_slots), max_blocks), dtype=torch.int32
      )
  
      # Allocate paged attention blocks
      cumulative_blocks = 0
      slots = []
      block_tables = []
      for i, (needed_blocks, needed_slots) in enumerate(
          needed_blocks_slots
      ):
          # Get allocated blocks for this sequence
          allocated_blocks = block_indices[
              cumulative_blocks : cumulative_blocks + needed_blocks
          ]
          # Get slots for the allocated blocks
          all_slots = self.slots[allocated_blocks].flatten()
  
          # Repeat slots in the case of context sliding window
          if needed_slots > len(all_slots) and self.repeat_slots:
              repeats = math.ceil(needed_slots / len(all_slots))
              all_slots = all_slots.repeat(repeats)
  
          allocated_slots = all_slots[:needed_slots]
  
          slots.append(allocated_slots)
          block_tables.append(allocated_blocks.tolist())
          block_tables_tensor[i, :needed_blocks] = allocated_blocks
          cumulative_blocks += needed_blocks
  
      block_tables = block_tables
      block_tables_tensor = block_tables_tensor.to(device)
      slots = torch.concat(slots).to(device)
  
      # Allocate the required number of blocks by setting the mask to 0
      self.free_block_mask[block_indices] = 0
  
      return block_tables, block_tables_tensor, slots
  
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 2 for calling the CacheManager obj's underlying member method code with example usage: obj.allocate(needed_blocks_slots, blocks, max_blocks, device) and returns: block_tables,  block_tables_tensor,  slots