@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e5588e49bc

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

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

GongLei-HW committed Jul 10, 2026 at 13:46 UTC e5588e49bc2642670116664a7fc4096e27adb179
5 files changed +270 -25
tests/v1/core/test_prefix_caching.py
+118 -1
@@ -12,7 +12,12 @@ import torch
12
13 import vllm.v1.core.kv_cache_manager as kv_cache_manager
14 import vllm.v1.core.kv_cache_utils as kv_cache_utils
15 -from vllm.distributed.kv_events import AllBlocksCleared, BlockRemoved, BlockStored
15 +from vllm.distributed.kv_events import (
16 + MEDIUM_GPU,
17 + AllBlocksCleared,
18 + BlockRemoved,
19 + BlockStored,
20 +)
21 from vllm.lora.request import LoRARequest
22 from vllm.multimodal.inputs import (
23 MultiModalFeatureSpec,
@@ -2454,6 +2459,118 @@ def test_block_removed_event_group_idx(group_id: int):
2459 assert event.group_idx == group_id
2460
2461
2462 +def test_emit_cached_block_events():
2463 + """emit_cached_block_events emits one BlockStored for already-cached
2464 + (reused) prefix blocks, carrying the correct group_idx /
2465 + parent_block_hash / token_ids, and without mutating block state."""
2466 + block_size = 4
2467 + num_cached_blocks = 3
2468 + kv_cache_group_id = 1
2469 + num_tokens = block_size * 4 # 4 full blocks; reuse the first 3
2470 +
2471 + pool = BlockPool(
2472 + num_gpu_blocks=8,
2473 + enable_caching=True,
2474 + hash_block_size=block_size,
2475 + enable_kv_cache_events=True,
2476 + )
2477 +
2478 + req = make_request(
2479 + "req_emit_cached",
2480 + prompt_token_ids=list(range(num_tokens)),
2481 + block_size=block_size,
2482 + hash_fn=sha256,
2483 + )
2484 + assert len(req.block_hashes) >= num_cached_blocks
2485 +
2486 + # Snapshot block state to prove emit_cached_block_events does not mutate it.
2487 + free_before = pool.get_num_free_blocks()
2488 + assert len(pool.cached_block_hash_to_block) == 0
2489 +
2490 + pool.emit_cached_block_events(
2491 + request=req,
2492 + num_cached_blocks=num_cached_blocks,
2493 + block_size=block_size,
2494 + kv_cache_group_id=kv_cache_group_id,
2495 + )
2496 +
2497 + # No block-state mutation: nothing allocated, nothing inserted into the
2498 + # prefix-cache map.
2499 + assert pool.get_num_free_blocks() == free_before
2500 + assert len(pool.cached_block_hash_to_block) == 0
2501 +
2502 + events = pool.take_events()
2503 + assert len(events) == 1
2504 + event = events[0]
2505 + assert isinstance(event, BlockStored)
2506 +
2507 + expected_hashes = [
2508 + kv_cache_utils.maybe_convert_block_hash(req.block_hashes[i])
2509 + for i in range(num_cached_blocks)
2510 + ]
2511 + assert event.block_hashes == expected_hashes
2512 + # Reused blocks start from block 0, so there is no parent block hash.
2513 + assert event.parent_block_hash is None
2514 + assert event.token_ids == list(req.all_token_ids[: num_cached_blocks * block_size])
2515 + assert event.group_idx == kv_cache_group_id
2516 + assert event.block_size == block_size
2517 + assert event.medium == MEDIUM_GPU
2518 + assert event.lora_id is None
2519 + assert event.lora_name is None
2520 +
2521 +
2522 +def test_emit_cached_block_events_disabled():
2523 + """No events are emitted when enable_kv_cache_events is False."""
2524 + block_size = 4
2525 + pool = BlockPool(
2526 + num_gpu_blocks=8,
2527 + enable_caching=True,
2528 + hash_block_size=block_size,
2529 + enable_kv_cache_events=False,
2530 + )
2531 + req = make_request(
2532 + "req_emit_disabled",
2533 + prompt_token_ids=list(range(block_size * 4)),
2534 + block_size=block_size,
2535 + hash_fn=sha256,
2536 + )
2537 +
2538 + pool.emit_cached_block_events(
2539 + request=req,
2540 + num_cached_blocks=3,
2541 + block_size=block_size,
2542 + kv_cache_group_id=0,
2543 + )
2544 +
2545 + assert pool.take_events() == []
2546 +
2547 +
2548 +def test_emit_cached_block_events_zero_cached():
2549 + """No events are emitted when num_cached_blocks == 0."""
2550 + block_size = 4
2551 + pool = BlockPool(
2552 + num_gpu_blocks=8,
2553 + enable_caching=True,
2554 + hash_block_size=block_size,
2555 + enable_kv_cache_events=True,
2556 + )
2557 + req = make_request(
2558 + "req_emit_zero",
2559 + prompt_token_ids=list(range(block_size * 4)),
2560 + block_size=block_size,
2561 + hash_fn=sha256,
2562 + )
2563 +
2564 + pool.emit_cached_block_events(
2565 + request=req,
2566 + num_cached_blocks=0,
2567 + block_size=block_size,
2568 + kv_cache_group_id=0,
2569 + )
2570 +
2571 + assert pool.take_events() == []
2572 +
2573 +
2574 def test_eagle_enabled_removes_last_block():
2575 """Verify Eagle does NOT remove blocks when request
2576 length is divisible by block size."""
vllm/v1/core/block_pool.py
+107 -24
@@ -14,8 +14,6 @@ from vllm.logger import init_logger
14 from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
15 from vllm.v1.core.kv_cache_utils import (
16 BlockHash,
17 - BlockHashList,
18 - BlockHashListWithBlockSize,
17 BlockHashWithGroupId,
18 ExternalBlockHash,
19 FreeKVCacheBlockQueue,
@@ -25,6 +23,7 @@ from vllm.v1.core.kv_cache_utils import (
23 get_group_id,
24 make_block_hash_with_group_id,
25 maybe_convert_block_hash,
26 + resolve_block_hashes,
27 )
28 from vllm.v1.request import Request
29
@@ -261,17 +260,7 @@ class BlockPool:
260 return
261 new_full_blocks = blocks[num_cached_blocks:num_full_blocks]
262 assert block_mask is None or len(block_mask) == len(new_full_blocks)
264 - if block_size == self.hash_block_size:
265 - # Common case.
266 - block_hashes: BlockHashList = request.block_hashes
267 - else:
268 - # block_size is a multiple of hash_block_size. This happens when
269 - # different KV cache groups have different block sizes.
270 - assert block_size % self.hash_block_size == 0
271 - block_hashes = BlockHashListWithBlockSize(
272 - request.block_hashes, self.hash_block_size, block_size
273 - )
274 - assert len(block_hashes) >= num_full_blocks
263 + block_hashes = resolve_block_hashes(request, self.hash_block_size, block_size)
264
265 new_block_hashes = block_hashes[num_cached_blocks:]
266 new_hashes: list[ExternalBlockHash] | None = (
@@ -338,23 +327,117 @@ class BlockPool:
327 extra_keys_list.append(extra_keys)
328
329 self.kv_event_queue.append(
341 - BlockStored(
330 + self._build_block_stored_event(
331 + request,
332 block_hashes=new_hashes,
333 parent_block_hash=parent_block_hash,
344 - token_ids=request.all_token_ids[start_token_idx:end_token_idx],
334 + start_token_idx=start_token_idx,
335 + end_token_idx=end_token_idx,
336 block_size=block_size,
346 - lora_id=request.lora_request.adapter_id
347 - if request.lora_request
348 - else None,
349 - medium=MEDIUM_GPU,
350 - lora_name=request.lora_request.name
351 - if request.lora_request
352 - else None,
353 - extra_keys=extra_keys_list if extra_keys_list else None,
354 - group_idx=kv_cache_group_id,
337 + kv_cache_group_id=kv_cache_group_id,
338 + extra_keys_list=extra_keys_list,
339 )
340 )
341
342 + def _build_block_stored_event(
343 + self,
344 + request: Request,
345 + block_hashes: list[ExternalBlockHash] | None,
346 + parent_block_hash: ExternalBlockHash | None,
347 + start_token_idx: int,
348 + end_token_idx: int,
349 + block_size: int,
350 + kv_cache_group_id: int,
351 + extra_keys_list: list[tuple[Any, ...] | None],
352 + ) -> BlockStored:
353 + """Build a ``BlockStored`` KV event for ``request``.
354 +
355 + Shared by ``cache_full_blocks`` (newly cached blocks) and
356 + ``emit_cached_block_events`` (prefix-cache-reused blocks) so both emit
357 + identical event shapes for downstream consumers.
358 + """
359 + return BlockStored(
360 + block_hashes=block_hashes,
361 + parent_block_hash=parent_block_hash,
362 + token_ids=request.all_token_ids[start_token_idx:end_token_idx],
363 + block_size=block_size,
364 + lora_id=request.lora_request.adapter_id if request.lora_request else None,
365 + medium=MEDIUM_GPU,
366 + lora_name=request.lora_request.name if request.lora_request else None,
367 + extra_keys=extra_keys_list if extra_keys_list else None,
368 + group_idx=kv_cache_group_id,
369 + )
370 +
371 + def emit_cached_block_events(
372 + self,
373 + request: Request,
374 + num_cached_blocks: int,
375 + block_size: int,
376 + kv_cache_group_id: int,
377 + ) -> None:
378 + """Generate BlockStored events for blocks reused from prefix cache.
379 +
380 + Unlike cache_full_blocks(), this does NOT modify block state —
381 + the blocks are already cached. It only generates events so that
382 + external consumers (e.g. gateway) can learn about reused blocks.
383 +
384 + Args:
385 + request: The request whose prefix cache blocks were reused.
386 + num_cached_blocks: Number of blocks that were cache hits.
387 + block_size: Number of tokens per block.
388 + kv_cache_group_id: The KV cache group ID.
389 + """
390 + if not self.enable_kv_cache_events or num_cached_blocks == 0:
391 + return
392 +
393 + block_hashes = resolve_block_hashes(request, self.hash_block_size, block_size)
394 +
395 + # Collect external hashes and extra_keys for cached blocks.
396 + cached_hashes: list[ExternalBlockHash] = []
397 + extra_keys_list: list[tuple[Any, ...] | None] = []
398 + curr_mm_idx = 0
399 + for i in range(num_cached_blocks):
400 + block_start = i * block_size
401 + block_end = block_start + block_size
402 + cached_hashes.append(maybe_convert_block_hash(block_hashes[i]))
403 + extra_keys, curr_mm_idx = generate_block_hash_extra_keys(
404 + request, block_start, block_end, curr_mm_idx
405 + )
406 + extra_keys_list.append(extra_keys)
407 +
408 + if not cached_hashes:
409 + return
410 +
411 + # Prefix-cache hits always form a contiguous prefix starting at block 0,
412 + # so the first (and thus the whole group's) parent block hash is None.
413 + parent_block_hash: ExternalBlockHash | None = None
414 + start_token_idx = 0
415 + end_token_idx = num_cached_blocks * block_size
416 +
417 + logger.debug(
418 + "EmitCachedBlock event: block_size=%d, "
419 + "num_cached_blocks=%d, parent_block_hash=%s, "
420 + "token_ids_len=%d, group_idx=%s",
421 + block_size,
422 + num_cached_blocks,
423 + parent_block_hash,
424 + len(request.all_token_ids[start_token_idx:end_token_idx]),
425 + kv_cache_group_id,
426 + )
427 +
428 + self.kv_event_queue.append(
429 + self._build_block_stored_event(
430 + request,
431 + block_hashes=cached_hashes,
432 + parent_block_hash=parent_block_hash,
433 + start_token_idx=start_token_idx,
434 + end_token_idx=end_token_idx,
435 + block_size=block_size,
436 + kv_cache_group_id=kv_cache_group_id,
437 + extra_keys_list=extra_keys_list,
438 + )
439 + )
440 +
441 def cache_partial_block(
442 self,
443 request: Request,
vllm/v1/core/kv_cache_manager.py
+21
@@ -136,6 +136,7 @@ class KVCacheManager:
136 max_in_flight_tokens = max_model_len
137
138 self.enable_caching = enable_caching
139 + self.enable_kv_cache_events = enable_kv_cache_events
140 self.use_eagle = use_eagle
141 self.log_stats = log_stats
142 self.metrics_collector = metrics_collector
@@ -235,6 +236,26 @@ class KVCacheManager:
236 )
237 )
238
239 + # When kv_cache_report_mode is "full", emit BlockStored events
240 + # for the reused prefix cache blocks so that external consumers
241 + # (e.g. gateway) can learn about them.
242 + if (
243 + num_new_computed_tokens > 0
244 + and self.enable_kv_cache_events
245 + and getattr(request, "kv_cache_report_mode", "incremental") == "full"
246 + ):
247 + for group_idx, group_blocks in enumerate(computed_blocks):
248 + num_blocks = len(group_blocks)
249 + if num_blocks > 0:
250 + group = self.kv_cache_config.kv_cache_groups[group_idx]
251 + block_size = group.kv_cache_spec.block_size
252 + self.block_pool.emit_cached_block_events(
253 + request,
254 + num_blocks,
255 + block_size,
256 + group_idx,
257 + )
258 +
259 if self.log_stats:
260 assert self.prefix_cache_stats is not None
261 self.prefix_cache_stats.record(
vllm/v1/core/kv_cache_utils.py
+19
@@ -2249,3 +2249,22 @@ class BlockHashListWithBlockSize:
2249
2250
2251 BlockHashList = list[BlockHash] | BlockHashListWithBlockSize
2252 +
2253 +
2254 +def resolve_block_hashes(
2255 + request: Request,
2256 + hash_block_size: int,
2257 + block_size: int,
2258 +) -> BlockHashList:
2259 + """Resolve the block-hash view for ``request`` at ``block_size``.
2260 +
2261 + When ``block_size`` equals ``hash_block_size``, reuse the request's
2262 + precomputed ``block_hashes`` directly; otherwise recalculate at
2263 + ``block_size`` granularity (``block_size`` must be a multiple of
2264 + ``hash_block_size``, which happens when KV cache groups differ in
2265 + block size).
2266 + """
2267 + if block_size == hash_block_size:
2268 + return request.block_hashes
2269 + assert block_size % hash_block_size == 0
2270 + return BlockHashListWithBlockSize(request.block_hashes, hash_block_size, block_size)
vllm/v1/request.py
+5
@@ -115,6 +115,11 @@ class Request:
115 self.kv_transfer_params = sampling_params.extra_args.get(
116 "kv_transfer_params"
117 )
118 + self.kv_cache_report_mode = sampling_params.extra_args.get(
119 + "kv_cache_report_mode", "incremental"
120 + )
121 + else:
122 + self.kv_cache_report_mode = "incremental"
123 else:
124 raise ValueError("sampling_params and pooling_params can't both be unset")
125