@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c6fa3895e9

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 12, 18:30 UTC

[KV Connector] Fix async connector prefix cache metrics (#28585)

Signed-off-by: Mark McLoughlin <markmc@redhat.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>

Mark McLoughlin committed Nov 21, 2025 at 22:45 UTC c6fa3895e90f6daef4d223188f6b4156311f40c9

3 files changed +24 -12

tests/v1/core/test_scheduler.py

+13 -4

@@ -1057,7 +1057,8 @@ def test_kv_connector_basic(is_async: bool):

1057 )

1058

1059

1060 -def test_external_prefix_cache_metrics():

1060 +@pytest.mark.parametrize("is_async", [False, True])

1061 +def test_external_prefix_cache_metrics(is_async: bool):

1062 """

1063 Verify connector prefix cache metrics are updated

1064 correctly when the scheduler processes requests with KV connector hits.

@@ -1067,7 +1068,9 @@ def test_external_prefix_cache_metrics():

1068 NUM_MATCHED_NEW_TOKENS = 4

1069 scheduler = create_scheduler(

1070 enable_prefix_caching=False,

1070 - use_kv_connector=mock_kv(matched_tokens=NUM_MATCHED_NEW_TOKENS, is_async=False),

1071 + use_kv_connector=mock_kv(

1072 + matched_tokens=NUM_MATCHED_NEW_TOKENS, is_async=is_async

1073 + ),

1074 )

1075

1076 # --- Prepare simple requests ---

@@ -1079,9 +1082,15 @@ def test_external_prefix_cache_metrics():

1082 num_tokens=NUM_TOKENS,

1083 max_tokens=MAX_TOKENS,

1084 )

1085 + req_ids = []

1086 + req_to_index = {}

1087 + for i, request in enumerate(requests):

1088 + scheduler.add_request(request)

1089 + req_ids.append(request.request_id)

1090 + req_to_index[request.request_id] = i

1091

1083 - for req in requests:

1084 - scheduler.add_request(req)

1092 + if is_async:

1093 + _step_until_kv_transfer_finished(scheduler, req_ids)

1094

1095 # --- Trigger scheduling and simulate model output ---

1096 output = scheduler.schedule()

vllm/v1/core/sched/scheduler.py

+8 -8

@@ -470,6 +470,7 @@ class Scheduler(SchedulerInterface):

470 skipped_waiting_requests.prepend_request(request)

471 continue

472

473 + request.num_external_computed_tokens = ext_tokens

474 num_external_computed_tokens = ext_tokens

475

476 # Total computed tokens (local + external).

@@ -576,9 +577,6 @@ class Scheduler(SchedulerInterface):

577 new_computed_blocks + new_blocks,

578 num_external_computed_tokens,

579 )

579 - self._update_connector_prefix_cache_stats(

580 - request, num_external_computed_tokens

581 - )

580

581 # Request was already popped from self.waiting

582 # unless it was re-added above due to new_blocks being None.

@@ -590,6 +588,8 @@ class Scheduler(SchedulerInterface):

588 request.status = RequestStatus.WAITING_FOR_REMOTE_KVS

589 continue

590

591 + self._update_connector_prefix_cache_stats(request)

592 +

593 req_index += 1

594 self.running.append(request)

595 if self.log_stats:

@@ -1380,15 +1380,13 @@ class Scheduler(SchedulerInterface):

1380 # KV Connector Related Methods

1381 ########################################################################

1382

1383 - def _update_connector_prefix_cache_stats(

1384 - self, request: Request, num_external_tokens: int

1385 - ) -> None:

1383 + def _update_connector_prefix_cache_stats(self, request: Request) -> None:

1384 if self.connector_prefix_cache_stats is None:

1385 return

1386

1387 self.connector_prefix_cache_stats.record(

1388 num_tokens=request.num_tokens,

1391 - num_hits=num_external_tokens,

1389 + num_hits=request.num_external_computed_tokens,

1390 preempted=request.num_preemptions > 0,

1391 )

1392

@@ -1571,9 +1569,11 @@ class Scheduler(SchedulerInterface):

1569 marked_invalid_block = True

1570 # Truncate the computed tokens at the first failed block

1571 request.num_computed_tokens = idx * self.block_size

1574 - total_affected_tokens += (

1572 + num_affected_tokens = (

1573 req_num_computed_tokens - request.num_computed_tokens

1574 )

1575 + total_affected_tokens += num_affected_tokens

1576 + request.num_external_computed_tokens -= num_affected_tokens

1577

1578 if is_affected:

1579 if not marked_invalid_block:

vllm/v1/request.py

@@ -121,6 +121,9 @@ class Request:

121 # The number of requests being preempted by the scheduler

122 self.num_preemptions = 0

123

124 + # The number of tokens that have been computed remotely.

125 + self.num_external_computed_tokens = 0

126 +

127 self.block_hashes: list[BlockHash] = []

128 self.get_hash_new_full_blocks: Callable[[], list[BlockHash]] | None = None

129 if block_hasher is not None: