A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix] OffloadingConnector: respect skip_reading_prefix_cache flag (#44592)
Signed-off-by: Hsiao-Yuan Chen <hy.c@Hsiao-YuandeMacBook-Pro.local> Signed-off-by: littlecircle0730 <littlecircle0730@gmail.com> Signed-off-by: littlecircle0730 <43994952+littlecircle0730@users.noreply.github.com> Co-authored-by: Hsiao-Yuan Chen <hy.c@Hsiao-YuandeMacBook-Pro.local> Co-authored-by: Or Ozeri <or@ozery.com>
littlecircle0730 committed
Jun 11, 2026 at 19:20 UTC
4bc83323f2ea8e85c87ae5fb5ff2d792a8f61f9d
3 files changed
+61
-2
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+51
@@ -1381,3 +1381,54 @@ def test_stale_sliding_window_block_after_prepare_store_failure(
1381
expected_stored=(2, 3),
1382
expected_flushed=(2, 3) if not async_scheduling else (),
1383
)
1384
+
1385
+
1386
+@pytest.mark.parametrize("async_scheduling", [True, False])
1387
+def test_skip_reading_prefix_cache(request_runner, async_scheduling: bool):
1388
+ """When skip_reading_prefix_cache=True, the offloading connector must not
1389
+ load any blocks from CPU even if a matching prefix is cached there."""
1390
+ block_size = 4
1391
+ block_size_factor = 3
1392
+ offloaded_block_size = block_size * block_size_factor
1393
+ num_gpu_blocks = 100
1394
+
1395
+ runner = request_runner(
1396
+ block_size=block_size,
1397
+ num_gpu_blocks=num_gpu_blocks,
1398
+ async_scheduling=async_scheduling,
1399
+ block_size_factor=block_size_factor,
1400
+ )
1401
+
1402
+ # Populate the CPU offload cache with one block.
1403
+ runner.new_request(token_ids=[0] * offloaded_block_size)
1404
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
1405
+ generate_store_output(keys)
1406
+ )
1407
+ runner.run(
1408
+ decoded_tokens=[EOS_TOKEN_ID],
1409
+ expected_stored=(0, 1, 2),
1410
+ expected_flushed=(0, 1, 2) if not async_scheduling else (),
1411
+ )
1412
+
1413
+ # Reset GPU prefix cache so the next request cannot hit locally.
1414
+ runner.scheduler.reset_prefix_cache()
1415
+
1416
+ # New request with identical tokens but skip_reading_prefix_cache=True.
1417
+ # The offloading connector must not load anything from CPU, but must
1418
+ # still offload the freshly computed blocks (state management intact).
1419
+ runner.new_request(
1420
+ token_ids=[0] * offloaded_block_size,
1421
+ skip_reading_prefix_cache=True,
1422
+ )
1423
+ runner.manager.prepare_store.side_effect = lambda keys, req_context: (
1424
+ generate_store_output(keys)
1425
+ )
1426
+ runner.run(
1427
+ decoded_tokens=[EOS_TOKEN_ID],
1428
+ expected_loaded=(), # no CPU loads must happen
1429
+ expected_stored=(0, 1, 2), # tokens still offloaded to CPU
1430
+ expected_flushed=(0, 1, 2) if not async_scheduling else (),
1431
+ )
1432
+
1433
+ # The external lookup must have been completely skipped.
1434
+ runner.manager.lookup.assert_not_called()
tests/v1/kv_connector/unit/offloading_connector/utils.py
+5
-1
@@ -324,10 +324,14 @@ class RequestRunner:
324
self,
325
token_ids: list[int],
326
kv_transfer_params: dict | None = None,
327
+ skip_reading_prefix_cache: bool = False,
328
):
329
self.req_id += 1
330
330
- sampling_params = SamplingParams(max_tokens=1000)
331
+ sampling_params = SamplingParams(
332
+ max_tokens=1000,
333
+ skip_reading_prefix_cache=skip_reading_prefix_cache or None,
334
+ )
335
sampling_params.update_from_generation_config({}, EOS_TOKEN_ID)
336
337
req = Request(
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+5
-1
@@ -571,7 +571,11 @@ class OffloadingConnectorScheduler:
571
req_status.update_offload_keys()
572
req_status.num_locally_computed_tokens = num_computed_tokens
573
574
- num_hit_tokens = self._lookup(req_status)
574
+ num_hit_tokens: int | None
575
+ if request.skip_reading_prefix_cache:
576
+ num_hit_tokens = 0
577
+ else:
578
+ num_hit_tokens = self._lookup(req_status)
579
req_status.update_num_hit_blocks(num_computed_tokens + (num_hit_tokens or 0))
580
581
self._touch(req_status)