skip looking up cache

microsoft · May 31, 2024 · ea50fb5 · ea50fb5
1 parent 0c8e0e6
commit ea50fb5
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/deepspeed/inference/v2/engine_v2.py b/deepspeed/inference/v2/engine_v2.py
@@ -260,6 +260,7 @@ def setup_cached_sequence(self, uid: int, cached_length: int, block_ids: torch.T
             seq.pre_forward(cached_length)
             seq.post_forward()
             seq.extend_kv_cache(block_ids)
+            seq.num_prefix_cache_blocks = len(block_ids)
             self._state_manager.increment_ref_count(block_ids)
             self._state_manager._kv_cache.allocate_blocks(block_ids)