Skip to content

Commit af376ca

Browse files
authored
[Core] Minimize number of dict lookup in _maybe_evict_cached_block (#21281)
Signed-off-by: Jialin Ouyang <[email protected]>
1 parent e7b2042 commit af376ca

File tree

1 file changed

+21
-16
lines changed

1 file changed

+21
-16
lines changed

vllm/v1/core/block_pool.py

Lines changed: 21 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -243,22 +243,27 @@ def _maybe_evict_cached_block(self, block: KVCacheBlock) -> bool:
243243
True if the block is evicted, False otherwise.
244244
"""
245245
block_hash = block.block_hash
246-
if block_hash and block_hash in self.cached_block_hash_to_block:
247-
block.reset_hash()
248-
del self.cached_block_hash_to_block[block_hash][block.block_id]
249-
250-
if len(self.cached_block_hash_to_block[block_hash]) == 0:
251-
del self.cached_block_hash_to_block[block_hash]
252-
253-
if self.enable_kv_cache_events:
254-
# FIXME (Chen): Not sure whether we should return `hash_value`
255-
# or `(hash_value, group_id)` here. But it's fine now because
256-
# we disable hybrid kv cache manager when kv cache event is
257-
# enabled, so there is only one group.
258-
self.kv_event_queue.append(
259-
BlockRemoved(block_hashes=[block_hash.get_hash_value()]))
260-
return True
261-
return False
246+
if block_hash is None:
247+
# The block doesn't have hash, eviction is not needed
248+
return False
249+
blocks_by_id = self.cached_block_hash_to_block.get(block_hash)
250+
if blocks_by_id is None:
251+
# block_hash not found in cached_block_hash_to_block,
252+
# eviction is not needed
253+
return False
254+
block.reset_hash()
255+
blocks_by_id.pop(block.block_id, None)
256+
if blocks_by_id:
257+
del self.cached_block_hash_to_block[block_hash]
258+
259+
if self.enable_kv_cache_events:
260+
# FIXME (Chen): Not sure whether we should return `hash_value`
261+
# or `(hash_value, group_id)` here. But it's fine now because
262+
# we disable hybrid kv cache manager when kv cache event is
263+
# enabled, so there is only one group.
264+
self.kv_event_queue.append(
265+
BlockRemoved(block_hashes=[block_hash.get_hash_value()]))
266+
return True
262267

263268
def touch(self, blocks: tuple[list[KVCacheBlock], ...]) -> None:
264269
"""Touch a block increases its reference count by 1, and may remove

0 commit comments

Comments
 (0)