fixes issues introduced in microsoft#4668

EleutherAI · Nov 16, 2023 · 5acef9e · 5acef9e
1 parent ce5e56a
commit 5acef9e
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/deepspeed/inference/v2/ragged/ragged_manager.py b/deepspeed/inference/v2/ragged/ragged_manager.py
@@ -116,7 +116,7 @@ def flush_sequence(self, uid: int) -> None:
  return
 
  seq = self._seqs[uid]
- for i in range(self.n_kv_caches):
+ for i in range(self.n_kv_cache_groups):
  self._kv_cache.free(seq.all_block_ids(cache_group=i), cache_group=i)
 
  self._tracking_allocator.free(seq.tracking_id)

diff --git a/deepspeed/inference/v2/ragged/sequence_descriptor.py b/deepspeed/inference/v2/ragged/sequence_descriptor.py
@@ -195,7 +195,7 @@ def kv_blocks_ptr(self, cache_group: int = 0) -> int:
  """
  return self._kv_cache_ids[cache_group].data_ptr()
 
- @property
+ #TODO: this was previously a property but causing issues with PR-4668 need to consult w. Connor
  def all_block_ids(self, cache_group: int = 0) -> torch.Tensor:
  """
  Return the Tensor containing all block IDs for this sequence in the specified cache group.