[Bugfix][KV Cache] Emit BlockRemoved only after the last physical copy is evicted

tests/v1/core/test_prefix_caching.py

@@ -2088,7 +2088,12 @@ def test_prefix_cache_stats_disabled():
 
 
 def test_maybe_evict_cached_block():
-    pool = BlockPool(num_gpu_blocks=4, enable_caching=True, hash_block_size=16)
+    pool = BlockPool(
+        num_gpu_blocks=4,
+        enable_caching=True,
+        hash_block_size=16,
+        enable_kv_cache_events=True,
+    )
     block_hash0 = make_block_hash_with_group_id(BlockHash(b"10"), 1000)
     block_hash1 = make_block_hash_with_group_id(BlockHash(b"20"), 2000)
     block_hash2 = make_block_hash_with_group_id(BlockHash(b"30"), 3000)
@@ -2120,19 +2125,82 @@ def test_maybe_evict_cached_block():
         block_hash0: {block0.block_id: block0, block3.block_id: block3},
         block_hash2: block2,
     }
+    [removed_event] = pool.take_events()
+    assert isinstance(removed_event, BlockRemoved)
+    assert removed_event.block_hashes == [
+        kv_cache_utils.maybe_convert_block_hash(get_block_hash(block_hash1))
+    ]
+    assert removed_event.group_idx == get_group_id(block_hash1)
+
     # Evict block0: block_hash0 entry should NOT be removed, as block3
     # also use the same hash
     pool._maybe_evict_cached_block(block0)
     assert pool.cached_block_hash_to_block._cache == {
         block_hash0: {block3.block_id: block3},
         block_hash2: block2,
     }
+    # BlockRemoved represents disappearance of the cache key, so reclaiming
+    # one physical copy must remain invisible while another copy is resident.
+    assert pool.take_events() == []
+
     # Evict block2
     pool._maybe_evict_cached_block(block2)
     assert pool.cached_block_hash_to_block._cache == {block_hash0: {3: block3}}
+    [removed_event] = pool.take_events()
+    assert isinstance(removed_event, BlockRemoved)
+    assert removed_event.block_hashes == [
+        kv_cache_utils.maybe_convert_block_hash(get_block_hash(block_hash2))
+    ]
+    assert removed_event.group_idx == get_group_id(block_hash2)
+
     # Evict block3
     pool._maybe_evict_cached_block(block3)
     assert pool.cached_block_hash_to_block._cache == {}
+    [removed_event] = pool.take_events()
+    assert isinstance(removed_event, BlockRemoved)
+    assert removed_event.block_hashes == [
+        kv_cache_utils.maybe_convert_block_hash(get_block_hash(block_hash0))
+    ]
+    assert removed_event.group_idx == get_group_id(block_hash0)
+
+
+def test_block_removed_last_copy_is_scoped_by_group():
+    pool = BlockPool(
+        num_gpu_blocks=4,
+        enable_caching=True,
+        hash_block_size=16,
+        enable_kv_cache_events=True,
+    )
+    external_hash = BlockHash(b"same-hash")
+    group0_key = make_block_hash_with_group_id(external_hash, 0)
+    group1_key = make_block_hash_with_group_id(external_hash, 1)
+    group0_block0, group0_block1, group1_block = pool.get_new_blocks(3)
+
+    for block, key in (
+        (group0_block0, group0_key),
+        (group0_block1, group0_key),
+        (group1_block, group1_key),
+    ):
+        block.set_block_hash(key)
+        pool.cached_block_hash_to_block.insert(key, block)
+
+    # One group-0 copy remains, so no logical removal is published.
+    pool._maybe_evict_cached_block(group0_block0)
+    assert pool.take_events() == []
+
+    # Removing group 1 does not affect the remaining group-0 copy.
+    pool._maybe_evict_cached_block(group1_block)
+    [removed_event] = pool.take_events()
+    assert isinstance(removed_event, BlockRemoved)
+    assert removed_event.group_idx == 1
+    assert pool.get_cached_block(external_hash, [0]) == [group0_block1]
+    assert pool.get_cached_block(external_hash, [1]) is None
+
+    # Group 0 publishes its own removal only when its final copy disappears.
+    pool._maybe_evict_cached_block(group0_block1)
+    [removed_event] = pool.take_events()
+    assert isinstance(removed_event, BlockRemoved)
+    assert removed_event.group_idx == 0
 
 
 @pytest.mark.parametrize("blocks_to_cache", [2, 3, 10])
@@ -3644,6 +3712,26 @@ def test_block_lookup_cache_multi_blocks_per_key():
     assert cache.pop(key1, 12) is None
 
 
+def test_block_lookup_cache_pop_with_remaining_count():
+    cache = BlockHashToBlockMap()
+    key = BlockHashWithGroupId(b"hash")
+    blocks = [KVCacheBlock(block_id) for block_id in range(3)]
+    for block in blocks:
+        cache.insert(key, block)
+
+    assert cache.pop_with_remaining_count(key, 100) == (None, 3)
+    assert cache.pop_with_remaining_count(key, 0) == (blocks[0], 2)
+    assert cache.pop_with_remaining_count(key, 1) == (blocks[1], 1)
+    assert cache.pop_with_remaining_count(key, 2) == (blocks[2], 0)
+    assert cache.pop_with_remaining_count(key, 2) == (None, 0)
+
+    single_key = BlockHashWithGroupId(b"single-hash")
+    single_block = KVCacheBlock(3)
+    cache.insert(single_key, single_block)
+    assert cache.pop_with_remaining_count(single_key, 100) == (None, 1)
+    assert cache.pop_with_remaining_count(single_key, 3) == (single_block, 0)
+
+
 def test_can_fit_full_sequence_swa_cap_admits_long_prompt():
     """Hybrid full+SWA model with a pool sized at the startup minimum should
     admit a prompt longer than the SWA cap, because SlidingWindowManager

vllm/v1/core/block_pool.py

@@ -107,31 +107,39 @@ def pop(self, key: BlockHashWithGroupId, block_id: int) -> KVCacheBlock | None:
         """
         Checks if block_hash exists and pop block_id from the cache
         """
+        block, _ = self.pop_with_remaining_count(key, block_id)
+        return block
+
+    def pop_with_remaining_count(
+        self, key: BlockHashWithGroupId, block_id: int
+    ) -> tuple[KVCacheBlock | None, int]:
+        """Pop a block and return the number of same-hash copies remaining."""
         blocks = self._cache.pop(key, None)
         if blocks is None:
             # block_hash not found in the cache
-            return None
+            return None, 0
         # TODO(Jialin): If key is found, block_id should always present
         # in blocks. We currently keep the original behaviour for safety.
         #
         # Will add block_id == blocks.block_id assertion and
         # use del blocks[block_id] instead as followup.
         if isinstance(blocks, KVCacheBlock):
             if blocks.block_id == block_id:
-                return blocks
+                return blocks, 0
             # If the single block ID doesn't match, we should put the
             # block back (it should happen rarely)
             self._cache[key] = blocks
-            return None
+            return None, 1
         if isinstance(blocks, dict):
             # Try to pop block_id from the block dict, and if dict still
             # contain blocks, put back to the cache.
             block = blocks.pop(block_id, None)
-            if len(blocks) > 0:
+            remaining_count = len(blocks)
+            if remaining_count > 0:
                 self._cache[key] = blocks
-            return block
+            return block, remaining_count
         self._unexpected_blocks_type(blocks)
-        return None
+        return None, 0
 
     def __len__(self) -> int:
         return len(self._cache)
@@ -571,31 +579,39 @@ def _get_partial_block_parent_hash_and_start(
     def _remove_cached_block_hashes(
         self,
         block: KVCacheBlock,
-    ) -> list[BlockHashWithGroupId]:
+    ) -> list[tuple[BlockHashWithGroupId, int]]:
         block_hashes: list[BlockHashWithGroupId] = []
         if block.block_hash is not None:
             block_hashes.append(block.block_hash)
         block_hashes.extend(self.cached_block_hashes_by_block.pop(block.block_id, ()))
         if not block_hashes:
             return []
 
-        removed_hashes: list[BlockHashWithGroupId] = []
+        removed_hashes: list[tuple[BlockHashWithGroupId, int]] = []
         for block_hash in block_hashes:
-            if (
-                self.cached_block_hash_to_block.pop(block_hash, block.block_id)
-                is not None
-            ):
-                removed_hashes.append(block_hash)
+            removed_block, remaining_count = (
+                self.cached_block_hash_to_block.pop_with_remaining_count(
+                    block_hash, block.block_id
+                )
+            )
+            if removed_block is not None:
+                removed_hashes.append((block_hash, remaining_count))
         block.reset_hash()
         return removed_hashes
 
     def _emit_block_removed_events(
         self,
-        block_hashes: list[BlockHashWithGroupId],
+        block_hashes: list[tuple[BlockHashWithGroupId, int]],
     ) -> None:
         if not self.enable_kv_cache_events:
             return
-        for block_hash in block_hashes:
+        for block_hash, remaining_count in block_hashes:
+            # BlockRemoved means the cache key is no longer resident, not that
+            # one of potentially several physical copies was reclaimed. Keep
+            # per-copy accounting internal and notify consumers only after the
+            # final copy in this KV cache group has disappeared.
+            if remaining_count > 0:
+                continue
             self.kv_event_queue.append(
                 BlockRemoved(
                     block_hashes=[maybe_convert_block_hash(get_block_hash(block_hash))],
@@ -640,7 +656,7 @@ def move_block_hashes(
         assert dst_block.block_hash is None
         assert dst_block.block_id not in self.cached_block_hashes_by_block
         num_tokens = src_block.block_hash_num_tokens
-        for block_hash in self._remove_cached_block_hashes(src_block):
+        for block_hash, _ in self._remove_cached_block_hashes(src_block):
             # `num_tokens` only applies to the first (primary) insertion.
             self._insert_block_hash(block_hash, dst_block, num_tokens=num_tokens)