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

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 8, 21:30 UTC

[KV Connector][Metrics] Do not count local prefix cache hits in connector queries (#30522)

Signed-off-by: Mark McLoughlin <markmc@redhat.com>

Mark McLoughlin committed Feb 5, 2026 at 07:57 UTC 2abd97592f947c041ba70329532f0cf62dd8971f

3 files changed +115 -31

tests/v1/core/test_scheduler.py

+69 -10

@@ -1136,7 +1136,7 @@ def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]):

1136 prompt_logprobs_dict={},

1137 pooler_output=[],

1138 )

1139 - scheduler.update_from_output(output, EMPTY_OUTPUT)

1139 + initial_ecos = scheduler.update_from_output(output, EMPTY_OUTPUT)

1140

1141 # Simulate KV transfer completion using KVConnectorOutput.finished_recving

1142 output = scheduler.schedule()

@@ -1156,6 +1156,8 @@ def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]):

1156 for req_id in req_ids:

1157 assert req_id in scheduler.finished_recving_kv_req_ids

1158

1159 + return initial_ecos

1160 +

1161

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

1163 def test_kv_connector_basic(is_async: bool):

@@ -1286,29 +1288,72 @@ def test_kv_connector_basic(is_async: bool):

1288

1289

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

1289 -def test_external_prefix_cache_metrics(is_async: bool):

1291 +@pytest.mark.parametrize("local_cache_hits", [False, True])

1292 +def test_external_prefix_cache_metrics(is_async: bool, local_cache_hits: bool):

1293 """

1294 Verify connector prefix cache metrics are updated

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

1296 """

1297

1298 + BLOCK_SIZE = 16

1299 + if local_cache_hits:

1300 + NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE * 2 # 32 tokens

1301 + NUM_LOCAL_HITS = NUM_MATCHED_NEW_TOKENS * 2 # 64 tokens

1302 + NUM_REQUESTS = 1

1303 + NUM_TOKENS = NUM_LOCAL_HITS * 2 # 128 tokens

1304 + else:

1305 + NUM_MATCHED_NEW_TOKENS = 4

1306 + NUM_LOCAL_HITS = 0

1307 + NUM_REQUESTS = 2

1308 + NUM_TOKENS = 8 # 8 tokens

1309 +

1310 # Setup Scheduler.

1296 - NUM_MATCHED_NEW_TOKENS = 4

1311 scheduler = create_scheduler(

1298 - enable_prefix_caching=False,

1312 + enable_prefix_caching=local_cache_hits,

1313 use_kv_connector=mock_kv(

1314 matched_tokens=NUM_MATCHED_NEW_TOKENS, is_async=is_async

1315 ),

1316 + block_size=BLOCK_SIZE,

1317 )

1318

1304 - # --- Prepare simple requests ---

1305 - NUM_REQUESTS = 2

1306 - NUM_TOKENS = 8

1319 + if local_cache_hits:

1320 + # First, establish local cache by running a request to completion

1321 + requests = create_requests(

1322 + num_requests=1,

1323 + num_tokens=NUM_LOCAL_HITS,

1324 + max_tokens=2,

1325 + block_size=BLOCK_SIZE,

1326 + )

1327 + req_ids = []

1328 + req_to_index = {}

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

1330 + scheduler.add_request(request)

1331 + req_ids.append(request.request_id)

1332 + req_to_index[request.request_id] = i

1333 +

1334 + if is_async:

1335 + _step_until_kv_transfer_finished(scheduler, req_ids)

1336 +

1337 + # Run first request to completion to establish local cache

1338 + output = scheduler.schedule()

1339 + MODEL_RUNNER_OUTPUT = ModelRunnerOutput(

1340 + req_ids=req_ids,

1341 + req_id_to_index=req_to_index,

1342 + sampled_token_ids=[[1000]] * len(req_ids),

1343 + logprobs=None,

1344 + prompt_logprobs_dict={},

1345 + pooler_output=[],

1346 + )

1347 + _step_until_done(scheduler, output, MODEL_RUNNER_OUTPUT)

1348 + _ = scheduler.schedule()

1349 +

1350 + # --- Prepare test requests ---

1351 MAX_TOKENS = 2

1352 requests = create_requests(

1353 num_requests=NUM_REQUESTS,

1354 num_tokens=NUM_TOKENS,

1355 max_tokens=MAX_TOKENS,

1356 + block_size=BLOCK_SIZE,

1357 )

1358 req_ids = []

1359 req_to_index = {}

@@ -1317,8 +1362,9 @@ def test_external_prefix_cache_metrics(is_async: bool):

1362 req_ids.append(request.request_id)

1363 req_to_index[request.request_id] = i

1364

1365 + initial_ecos = None

1366 if is_async:

1321 - _step_until_kv_transfer_finished(scheduler, req_ids)

1367 + initial_ecos = _step_until_kv_transfer_finished(scheduler, req_ids)

1368

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

1370 output = scheduler.schedule()

@@ -1338,10 +1384,23 @@ def test_external_prefix_cache_metrics(is_async: bool):

1384 assert ecos is not None and len(ecos) > 0

1385 assert ecos[0].scheduler_stats is not None

1386

1341 - external_stats = ecos[0].scheduler_stats.connector_prefix_cache_stats

1387 + if local_cache_hits:

1388 + # For async, local cache stats come from the first step

1389 + if initial_ecos:

1390 + local_stats = initial_ecos[0].scheduler_stats.prefix_cache_stats

1391 + else:

1392 + local_stats = ecos[0].scheduler_stats.prefix_cache_stats

1393 + assert local_stats is not None

1394 + assert local_stats.queries == NUM_TOKENS * NUM_REQUESTS

1395 + assert local_stats.hits == NUM_LOCAL_HITS * NUM_REQUESTS

1396 +

1397 + if initial_ecos:

1398 + external_stats = initial_ecos[0].scheduler_stats.connector_prefix_cache_stats

1399 + else:

1400 + external_stats = ecos[0].scheduler_stats.connector_prefix_cache_stats

1401 assert external_stats is not None

1402

1344 - assert external_stats.queries == NUM_TOKENS * NUM_REQUESTS

1403 + assert external_stats.queries == (NUM_TOKENS - NUM_LOCAL_HITS) * NUM_REQUESTS

1404 assert external_stats.hits == NUM_MATCHED_NEW_TOKENS * NUM_REQUESTS

1405 assert external_stats.requests == NUM_REQUESTS

1406 assert external_stats.preempted_requests == 0

tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py

+27 -1

@@ -281,6 +281,17 @@ def test_sync_fail_invalid_blocks_evicted(fail_scheduler: Scheduler):

281 f"(hash should be None), but hash is still {block.block_hash}"

282 )

283

284 + # Verify connector prefix cache stats:

285 + # - queries = num_prompt_tokens (total tokens not in local cache)

286 + # - hits = num_external_computed_tokens (tokens loaded externally)

287 + assert engine_outputs.scheduler_stats is not None

288 + stats = engine_outputs.scheduler_stats

289 + assert stats.connector_prefix_cache_stats is not None

290 + conn_stats = stats.connector_prefix_cache_stats

291 + assert conn_stats.requests == 1

292 + assert conn_stats.queries == num_prompt_tokens

293 + assert conn_stats.hits == num_external_computed_tokens

294 +

295

296 def test_async_recompute_blocks_not_cached_when_invalid(

297 recompute_scheduler: Scheduler,

@@ -364,7 +375,9 @@ def test_async_recompute_blocks_not_cached_when_invalid(

375 with patch.object(

376 recompute_scheduler.kv_cache_manager, "evict_blocks", evict_blocks_spy

377 ):

367 - recompute_scheduler.update_from_output(scheduler_output, model_runner_output)

378 + outputs = recompute_scheduler.update_from_output(

379 + scheduler_output, model_runner_output

380 + )

381

382 # verify evict_blocks was NOT called (async blocks excluded from eviction)

383 assert len(evict_blocks_calls) == 0, (

@@ -386,6 +399,19 @@ def test_async_recompute_blocks_not_cached_when_invalid(

399 f"Block {invalid_block_id} hash should be None but is {block.block_hash}"

400 )

401

402 + # Verify connector prefix cache stats:

403 + # - queries = num_prompt_tokens (total tokens not in local cache)

404 + # - hits = num_external_computed_tokens (tokens loaded externally)

405 + assert len(outputs) == 1

406 + engine_outputs = next(iter(outputs.values()))

407 + assert engine_outputs.scheduler_stats is not None

408 + stats = engine_outputs.scheduler_stats

409 + assert stats.connector_prefix_cache_stats is not None

410 + conn_stats = stats.connector_prefix_cache_stats

411 + assert conn_stats.requests == 1

412 + assert conn_stats.queries == num_prompt_tokens

413 + assert conn_stats.hits == num_external_computed_tokens

414 +

415 # now simulate async transfer completing

416 model_runner_output_2 = create_model_runner_output(

417 reqs=[],

vllm/v1/core/sched/scheduler.py

+19 -20

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

586

587 num_external_computed_tokens = 0

588 load_kv_async = False

589 + connector_prefix_cache_queries, connector_prefix_cache_hits = 0, 0

590

591 # Get already-cached tokens.

592 if request.num_computed_tokens == 0:

@@ -613,6 +614,11 @@ class Scheduler(SchedulerInterface):

614 request.num_external_computed_tokens = ext_tokens

615 num_external_computed_tokens = ext_tokens

616

617 + connector_prefix_cache_queries = (

618 + request.num_tokens - num_new_local_computed_tokens

619 + )

620 + connector_prefix_cache_hits = num_external_computed_tokens

621 +

622 # Total computed tokens (local + external).

623 num_computed_tokens = (

624 num_new_local_computed_tokens + num_external_computed_tokens

@@ -728,6 +734,15 @@ class Scheduler(SchedulerInterface):

734 self.kv_cache_manager.get_blocks(request_id),

735 num_external_computed_tokens,

736 )

737 + if (

738 + self.connector_prefix_cache_stats is not None

739 + and connector_prefix_cache_queries != 0

740 + ):

741 + self.connector_prefix_cache_stats.record(

742 + num_tokens=connector_prefix_cache_queries,

743 + num_hits=connector_prefix_cache_hits,

744 + preempted=request.num_preemptions > 0,

745 + )

746

747 # Request was already popped from self.waiting

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

@@ -739,8 +754,6 @@ class Scheduler(SchedulerInterface):

754 request.status = RequestStatus.WAITING_FOR_REMOTE_KVS

755 continue

756

742 - self._update_connector_prefix_cache_stats(request)

743 -

757 self.running.append(request)

758 if self.log_stats:

759 request.record_event(

@@ -1805,7 +1818,10 @@ class Scheduler(SchedulerInterface):

1818 return None

1819 prefix_cache_stats = self.kv_cache_manager.make_prefix_cache_stats()

1820 assert prefix_cache_stats is not None

1808 - connector_prefix_cache_stats = self._make_connector_prefix_cache_stats()

1821 + connector_prefix_cache_stats: PrefixCacheStats | None = None

1822 + if self.connector_prefix_cache_stats is not None:

1823 + connector_prefix_cache_stats = self.connector_prefix_cache_stats

1824 + self.connector_prefix_cache_stats = PrefixCacheStats()

1825 eviction_events = (

1826 self.kv_metrics_collector.drain_events()

1827 if self.kv_metrics_collector is not None

@@ -1866,23 +1882,6 @@ class Scheduler(SchedulerInterface):

1882 # KV Connector Related Methods

1883 ########################################################################

1884

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

1870 - if self.connector_prefix_cache_stats is None:

1871 - return

1872 -

1873 - self.connector_prefix_cache_stats.record(

1874 - num_tokens=request.num_tokens,

1875 - num_hits=request.num_external_computed_tokens,

1876 - preempted=request.num_preemptions > 0,

1877 - )

1878 -

1879 - def _make_connector_prefix_cache_stats(self) -> PrefixCacheStats | None:

1880 - if self.connector_prefix_cache_stats is None:

1881 - return None

1882 - stats = self.connector_prefix_cache_stats

1883 - self.connector_prefix_cache_stats = PrefixCacheStats()

1884 - return stats

1885 -

1885 def get_kv_connector(self) -> KVConnectorBase_V1 | None:

1886 return self.connector

1887