[Core][KV events] Report prefix-cache-reused blocks in full report mode (vllm-project#45261)

Signed-off-by: Lei Gong <gonglei25@huawei.com>
Co-authored-by: Lei Gong <gonglei25@huawei.com>
Co-authored-by: Claude <noreply@anthropic.com>

Author: 3 people

tests/v1/core/test_prefix_caching.py

@@ -12,7 +12,12 @@
 
 import vllm.v1.core.kv_cache_manager as kv_cache_manager
 import vllm.v1.core.kv_cache_utils as kv_cache_utils
-from vllm.distributed.kv_events import AllBlocksCleared, BlockRemoved, BlockStored
+from vllm.distributed.kv_events import (
+    MEDIUM_GPU,
+    AllBlocksCleared,
+    BlockRemoved,
+    BlockStored,
+)
 from vllm.lora.request import LoRARequest
 from vllm.multimodal.inputs import (
     MultiModalFeatureSpec,
@@ -2454,6 +2459,118 @@ def test_block_removed_event_group_idx(group_id: int):
         assert event.group_idx == group_id
 
 
+def test_emit_cached_block_events():
+    """emit_cached_block_events emits one BlockStored for already-cached
+    (reused) prefix blocks, carrying the correct group_idx /
+    parent_block_hash / token_ids, and without mutating block state."""
+    block_size = 4
+    num_cached_blocks = 3
+    kv_cache_group_id = 1
+    num_tokens = block_size * 4  # 4 full blocks; reuse the first 3
+
+    pool = BlockPool(
+        num_gpu_blocks=8,
+        enable_caching=True,
+        hash_block_size=block_size,
+        enable_kv_cache_events=True,
+    )
+
+    req = make_request(
+        "req_emit_cached",
+        prompt_token_ids=list(range(num_tokens)),
+        block_size=block_size,
+        hash_fn=sha256,
+    )
+    assert len(req.block_hashes) >= num_cached_blocks
+
+    # Snapshot block state to prove emit_cached_block_events does not mutate it.
+    free_before = pool.get_num_free_blocks()
+    assert len(pool.cached_block_hash_to_block) == 0
+
+    pool.emit_cached_block_events(
+        request=req,
+        num_cached_blocks=num_cached_blocks,
+        block_size=block_size,
+        kv_cache_group_id=kv_cache_group_id,
+    )
+
+    # No block-state mutation: nothing allocated, nothing inserted into the
+    # prefix-cache map.
+    assert pool.get_num_free_blocks() == free_before
+    assert len(pool.cached_block_hash_to_block) == 0
+
+    events = pool.take_events()
+    assert len(events) == 1
+    event = events[0]
+    assert isinstance(event, BlockStored)
+
+    expected_hashes = [
+        kv_cache_utils.maybe_convert_block_hash(req.block_hashes[i])
+        for i in range(num_cached_blocks)
+    ]
+    assert event.block_hashes == expected_hashes
+    # Reused blocks start from block 0, so there is no parent block hash.
+    assert event.parent_block_hash is None
+    assert event.token_ids == list(req.all_token_ids[: num_cached_blocks * block_size])
+    assert event.group_idx == kv_cache_group_id
+    assert event.block_size == block_size
+    assert event.medium == MEDIUM_GPU
+    assert event.lora_id is None
+    assert event.lora_name is None
+
+
+def test_emit_cached_block_events_disabled():
+    """No events are emitted when enable_kv_cache_events is False."""
+    block_size = 4
+    pool = BlockPool(
+        num_gpu_blocks=8,
+        enable_caching=True,
+        hash_block_size=block_size,
+        enable_kv_cache_events=False,
+    )
+    req = make_request(
+        "req_emit_disabled",
+        prompt_token_ids=list(range(block_size * 4)),
+        block_size=block_size,
+        hash_fn=sha256,
+    )
+
+    pool.emit_cached_block_events(
+        request=req,
+        num_cached_blocks=3,
+        block_size=block_size,
+        kv_cache_group_id=0,
+    )
+
+    assert pool.take_events() == []
+
+
+def test_emit_cached_block_events_zero_cached():
+    """No events are emitted when num_cached_blocks == 0."""
+    block_size = 4
+    pool = BlockPool(
+        num_gpu_blocks=8,
+        enable_caching=True,
+        hash_block_size=block_size,
+        enable_kv_cache_events=True,
+    )
+    req = make_request(
+        "req_emit_zero",
+        prompt_token_ids=list(range(block_size * 4)),
+        block_size=block_size,
+        hash_fn=sha256,
+    )
+
+    pool.emit_cached_block_events(
+        request=req,
+        num_cached_blocks=0,
+        block_size=block_size,
+        kv_cache_group_id=0,
+    )
+
+    assert pool.take_events() == []
+
+
 def test_eagle_enabled_removes_last_block():
     """Verify Eagle does NOT remove blocks when request
     length is divisible by block size."""

vllm/v1/core/block_pool.py

@@ -14,8 +14,6 @@
 from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
 from vllm.v1.core.kv_cache_utils import (
     BlockHash,
-    BlockHashList,
-    BlockHashListWithBlockSize,
     BlockHashWithGroupId,
     ExternalBlockHash,
     FreeKVCacheBlockQueue,
@@ -25,6 +23,7 @@
     get_group_id,
     make_block_hash_with_group_id,
     maybe_convert_block_hash,
+    resolve_block_hashes,
 )
 from vllm.v1.request import Request
 
@@ -261,17 +260,7 @@ def cache_full_blocks(
             return
         new_full_blocks = blocks[num_cached_blocks:num_full_blocks]
         assert block_mask is None or len(block_mask) == len(new_full_blocks)
-        if block_size == self.hash_block_size:
-            # Common case.
-            block_hashes: BlockHashList = request.block_hashes
-        else:
-            # block_size is a multiple of hash_block_size. This happens when
-            # different KV cache groups have different block sizes.
-            assert block_size % self.hash_block_size == 0
-            block_hashes = BlockHashListWithBlockSize(
-                request.block_hashes, self.hash_block_size, block_size
-            )
-        assert len(block_hashes) >= num_full_blocks
+        block_hashes = resolve_block_hashes(request, self.hash_block_size, block_size)
 
         new_block_hashes = block_hashes[num_cached_blocks:]
         new_hashes: list[ExternalBlockHash] | None = (
@@ -338,23 +327,117 @@ def cache_full_blocks(
                 extra_keys_list.append(extra_keys)
 
             self.kv_event_queue.append(
-                BlockStored(
+                self._build_block_stored_event(
+                    request,
                     block_hashes=new_hashes,
                     parent_block_hash=parent_block_hash,
-                    token_ids=request.all_token_ids[start_token_idx:end_token_idx],
+                    start_token_idx=start_token_idx,
+                    end_token_idx=end_token_idx,
                     block_size=block_size,
-                    lora_id=request.lora_request.adapter_id
-                    if request.lora_request
-                    else None,
-                    medium=MEDIUM_GPU,
-                    lora_name=request.lora_request.name
-                    if request.lora_request
-                    else None,
-                    extra_keys=extra_keys_list if extra_keys_list else None,
-                    group_idx=kv_cache_group_id,
+                    kv_cache_group_id=kv_cache_group_id,
+                    extra_keys_list=extra_keys_list,
                 )
             )
 
+    def _build_block_stored_event(
+        self,
+        request: Request,
+        block_hashes: list[ExternalBlockHash] | None,
+        parent_block_hash: ExternalBlockHash | None,
+        start_token_idx: int,
+        end_token_idx: int,
+        block_size: int,
+        kv_cache_group_id: int,
+        extra_keys_list: list[tuple[Any, ...] | None],
+    ) -> BlockStored:
+        """Build a ``BlockStored`` KV event for ``request``.
+
+        Shared by ``cache_full_blocks`` (newly cached blocks) and
+        ``emit_cached_block_events`` (prefix-cache-reused blocks) so both emit
+        identical event shapes for downstream consumers.
+        """
+        return BlockStored(
+            block_hashes=block_hashes,
+            parent_block_hash=parent_block_hash,
+            token_ids=request.all_token_ids[start_token_idx:end_token_idx],
+            block_size=block_size,
+            lora_id=request.lora_request.adapter_id if request.lora_request else None,
+            medium=MEDIUM_GPU,
+            lora_name=request.lora_request.name if request.lora_request else None,
+            extra_keys=extra_keys_list if extra_keys_list else None,
+            group_idx=kv_cache_group_id,
+        )
+
+    def emit_cached_block_events(
+        self,
+        request: Request,
+        num_cached_blocks: int,
+        block_size: int,
+        kv_cache_group_id: int,
+    ) -> None:
+        """Generate BlockStored events for blocks reused from prefix cache.
+
+        Unlike cache_full_blocks(), this does NOT modify block state —
+        the blocks are already cached. It only generates events so that
+        external consumers (e.g. gateway) can learn about reused blocks.
+
+        Args:
+            request: The request whose prefix cache blocks were reused.
+            num_cached_blocks: Number of blocks that were cache hits.
+            block_size: Number of tokens per block.
+            kv_cache_group_id: The KV cache group ID.
+        """
+        if not self.enable_kv_cache_events or num_cached_blocks == 0:
+            return
+
+        block_hashes = resolve_block_hashes(request, self.hash_block_size, block_size)
+
+        # Collect external hashes and extra_keys for cached blocks.
+        cached_hashes: list[ExternalBlockHash] = []
+        extra_keys_list: list[tuple[Any, ...] | None] = []
+        curr_mm_idx = 0
+        for i in range(num_cached_blocks):
+            block_start = i * block_size
+            block_end = block_start + block_size
+            cached_hashes.append(maybe_convert_block_hash(block_hashes[i]))
+            extra_keys, curr_mm_idx = generate_block_hash_extra_keys(
+                request, block_start, block_end, curr_mm_idx
+            )
+            extra_keys_list.append(extra_keys)
+
+        if not cached_hashes:
+            return
+
+        # Prefix-cache hits always form a contiguous prefix starting at block 0,
+        # so the first (and thus the whole group's) parent block hash is None.
+        parent_block_hash: ExternalBlockHash | None = None
+        start_token_idx = 0
+        end_token_idx = num_cached_blocks * block_size
+
+        logger.debug(
+            "EmitCachedBlock event: block_size=%d, "
+            "num_cached_blocks=%d, parent_block_hash=%s, "
+            "token_ids_len=%d, group_idx=%s",
+            block_size,
+            num_cached_blocks,
+            parent_block_hash,
+            len(request.all_token_ids[start_token_idx:end_token_idx]),
+            kv_cache_group_id,
+        )
+
+        self.kv_event_queue.append(
+            self._build_block_stored_event(
+                request,
+                block_hashes=cached_hashes,
+                parent_block_hash=parent_block_hash,
+                start_token_idx=start_token_idx,
+                end_token_idx=end_token_idx,
+                block_size=block_size,
+                kv_cache_group_id=kv_cache_group_id,
+                extra_keys_list=extra_keys_list,
+            )
+        )
+
     def cache_partial_block(
         self,
         request: Request,

vllm/v1/core/kv_cache_manager.py

@@ -136,6 +136,7 @@ def __init__(
             max_in_flight_tokens = max_model_len
 
         self.enable_caching = enable_caching
+        self.enable_kv_cache_events = enable_kv_cache_events
         self.use_eagle = use_eagle
         self.log_stats = log_stats
         self.metrics_collector = metrics_collector
@@ -235,6 +236,26 @@ def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int]:
             )
         )
 
+        # When kv_cache_report_mode is "full", emit BlockStored events
+        # for the reused prefix cache blocks so that external consumers
+        # (e.g. gateway) can learn about them.
+        if (
+            num_new_computed_tokens > 0
+            and self.enable_kv_cache_events
+            and getattr(request, "kv_cache_report_mode", "incremental") == "full"
+        ):
+            for group_idx, group_blocks in enumerate(computed_blocks):
+                num_blocks = len(group_blocks)
+                if num_blocks > 0:
+                    group = self.kv_cache_config.kv_cache_groups[group_idx]
+                    block_size = group.kv_cache_spec.block_size
+                    self.block_pool.emit_cached_block_events(
+                        request,
+                        num_blocks,
+                        block_size,
+                        group_idx,
+                    )
+
         if self.log_stats:
             assert self.prefix_cache_stats is not None
             self.prefix_cache_stats.record(

vllm/v1/core/kv_cache_utils.py

@@ -2249,3 +2249,22 @@ def _get_value_at(self, idx: int) -> BlockHash:
 
 
 BlockHashList = list[BlockHash] | BlockHashListWithBlockSize
+
+
+def resolve_block_hashes(
+    request: Request,
+    hash_block_size: int,
+    block_size: int,
+) -> BlockHashList:
+    """Resolve the block-hash view for ``request`` at ``block_size``.
+
+    When ``block_size`` equals ``hash_block_size``, reuse the request's
+    precomputed ``block_hashes`` directly; otherwise recalculate at
+    ``block_size`` granularity (``block_size`` must be a multiple of
+    ``hash_block_size``, which happens when KV cache groups differ in
+    block size).
+    """
+    if block_size == hash_block_size:
+        return request.block_hashes
+    assert block_size % hash_block_size == 0
+    return BlockHashListWithBlockSize(request.block_hashes, hash_block_size, block_size)

vllm/v1/request.py

@@ -115,6 +115,11 @@ def __init__(
                 self.kv_transfer_params = sampling_params.extra_args.get(
                     "kv_transfer_params"
                 )
+                self.kv_cache_report_mode = sampling_params.extra_args.get(
+                    "kv_cache_report_mode", "incremental"
+                )
+            else:
+                self.kv_cache_report_mode = "incremental"
         else:
             raise ValueError("sampling_params and pooling_params can't both be unset")