From 5afea9b33d38090c12f0db4fd88331827cb48d03 Mon Sep 17 00:00:00 2001 From: Advay Monga Date: Tue, 20 Oct 2026 19:09:09 -0400 Subject: [PATCH] [Bugfix][KV Offload] Reset the store cursor when a request restarts from token 0 A request resumed from preemption restarts at token 0, but the connector kept next_stored_chunk_idx from before the preemption. If the CPU tier had evicted its chunks in the meantime, chunks below the cursor were never offered again, so every chunk the request stored afterwards was unreachable behind a miss. Co-authored-by: Claude Fable 5.1 Signed-off-by: Advay Monga --- .../offloading_connector/test_scheduler.py | 25 +++++++++++++------ .../kv_connector/v1/offloading/scheduler.py | 2 ++ 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 4b52f45e1c29..e45d841fc660 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -1429,7 +1429,9 @@ def test_offloading_connector(request_runner, async_scheduling: bool): @pytest.mark.parametrize("async_scheduling", [True, False]) -def test_request_preemption(request_runner, async_scheduling: bool): +@pytest.mark.parametrize("cpu_hit", [True, False]) +def test_request_preemption(request_runner, async_scheduling: bool, cpu_hit: bool): + """Resume reloads the chunks, or re-stores them if the CPU tier lost them.""" block_size = 4 blocks_per_chunk = 3 tokens_per_chunk = block_size * blocks_per_chunk @@ -1480,16 +1482,23 @@ def test_request_preemption(request_runner, async_scheduling: bool): free_block_queue.num_free_blocks = num_free_blocks_empty runner.scheduler.reset_prefix_cache() - # request should now return from preemption - # re-load [0, ..., 8] from the CPU and store [9, 10, 11] - runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 3 runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) - runner.run( - decoded_tokens=[0] * block_size, - expected_loaded=(0, 1, 2, 3, 4, 5, 6, 7, 8), - ) + if cpu_hit: + # request should now return from preemption + # re-load [0, ..., 8] from the CPU and store [9, 10, 11] + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 3 + runner.run( + decoded_tokens=[0] * block_size, + expected_loaded=(0, 1, 2, 3, 4, 5, 6, 7, 8), + ) + else: + # lookup misses: recompute and re-store [0, ..., 8], then store [9, 10, 11] + runner.run( + decoded_tokens=[0] * block_size, + expected_stored=(0, 1, 2, 3, 4, 5, 6, 7, 8), + ) runner.run( decoded_tokens=[EOS_TOKEN_ID], diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 284538f8631e..a95a61f9e41e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -1144,6 +1144,8 @@ def get_num_new_matched_tokens( req_status = self._req_status[request.request_id] for group_state in req_status.group_states: group_state.block_ids.clear() + # Restarting at token 0: offer every chunk again. + group_state.next_stored_chunk_idx = 0 if req_status.transfer_jobs: logger.debug(