class_docstr: Base abstract class for managing blocked KV caches. Will probably have a single
  implementation for now.
docstr: The key
  
  Parameters:
      block_size (int): The number of tokens to allocate in each block.
docstr_tok: ["The","key","Parameters",":","block_size","(","int",")",":","The","number","of","tokens","to","allocate","in","each","block","."]
code: def __init__(
      self,
      config: DSStateManagerConfig,
      kv_configs: Tuple[KVCacheConfig, ...],
      base_mp_group: Optional[Any] = None,
  ) -> None:
      """
      The key
  
      Parameters:
          block_size (int): The number of tokens to allocate in each block.
      """
      self._config = config
      self._kv_configs = kv_configs
  
      # Load our helpers for host allocation.
      self._ragged_utils = RaggedUtilsBuilder().load()
  
      # Initialize the allocator for tracking sequences (so this doesn't need to be ad-hoc).
      self._tracking_allocator = BlockedAllocator(
          self._config.max_tracked_sequences
      )
  
      all_block_ids = []
      all_block_ids_shadow = []
  
      for cache_config in self._kv_configs:
          # Storage to back tracking the KV cache allocation.
          ids_shape = (
              self._config.max_tracked_sequences,
              cache_config.num_allocation_groups,
              cache_config.max_blocks_per_allocation_group,
          )
  
          all_block_ids.append(
              torch.zeros(
                  ids_shape,
                  dtype=torch.int32,
                  device=get_accelerator().current_device(),
              )
          )
          all_block_ids_shadow.append(
              self._ragged_utils.allocate_fast_host_buffer(
                  all_block_ids[-1]
              )
          )
  
      self._all_block_ids = tuple(all_block_ids)
      self._all_block_ids_shadow = tuple(all_block_ids_shadow)
  
      # Initialize the sequence container.
      self._seqs = {}
  
      # Finally initialize the KV cache.
      self._kv_cache = BlockedKVCache(
          self._kv_configs,
          self._config.memory_config,
          mp_group=base_mp_group,
          offload=self._config.offload,
      )
  
imports: import torch
  from typing import Any, Dict, Optional, Tuple
  from deepspeed.accelerator import get_accelerator
  from deepspeed.ops.op_builder import RaggedUtilsBuilder
  from deepspeed.utils.logging import logger
  from .blocked_allocator import BlockedAllocator
  from .kv_cache import BlockedKVCache
  from .manager_configs import DSStateManagerConfig, KVCacheConfig
  from .sequence_descriptor import DSSequenceDescriptor
label_desc: Use image node_id 1 to create a new DSStateManager object  with example: obj = DSStateManager(config, kv_configs, base_mp_group)