From 2705bb2b098b7cf7f03b2b9aec9fd40b244be29d Mon Sep 17 00:00:00 2001 From: Derek Yates Date: Sun, 19 Jul 2026 16:56:42 -0400 Subject: [PATCH] [Bugfix][KV Offloading] Defer stores until block IDs are tracked Bound each group store boundary by storable tokens, available keys, and complete tracked GPU-block chunks. Advance store progress only through the common ready prefix so transient tracking lag is retried. Assisted-by: OpenAI Codex Signed-off-by: Derek Yates --- .../offloading_connector/test_scheduler.py | 34 +++++++++++ .../kv_connector/v1/offloading/scheduler.py | 56 +++++++++++++------ 2 files changed, 73 insertions(+), 17 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 bc5190db00f5..22b7b4787264 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -19,6 +19,7 @@ OffloadingConnectorScheduler, RequestOffloadState, ) +from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, @@ -64,6 +65,39 @@ def test_scheduler_reports_allocation_failure(request_runner): assert reduced[_ConnectorMetricName.ALLOCATION_FAILURE] == 1 +def test_store_defers_chunk_until_block_id_is_tracked(request_runner): + """Store the ready prefix, then retry a chunk once its GPU block ID arrives.""" + runner = request_runner( + block_size=4, + num_gpu_blocks=10, + async_scheduling=True, + ) + runner.new_request(token_ids=[0] * 12) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + + req_status = runner.connector_scheduler._req_status["0"] + req_status.update_offload_keys() + group_state = req_status.group_states[0] + group_state.block_ids.extend([1, 2]) + + scheduler_output = SchedulerOutput.make_empty() + scheduler_output.num_scheduled_tokens = {"0": 12} + first_jobs = runner.connector_scheduler._build_store_jobs(scheduler_output) + + assert len(first_jobs) == 1 + assert next(iter(first_jobs.values())).src_spec.block_ids.tolist() == [1, 2] + assert group_state.next_stored_chunk_idx == 2 + + group_state.block_ids.append(3) + second_jobs = runner.connector_scheduler._build_store_jobs(scheduler_output) + + assert len(second_jobs) == 1 + assert next(iter(second_jobs.values())).src_spec.block_ids.tolist() == [3] + assert group_state.next_stored_chunk_idx == 3 + + @pytest.mark.parametrize("async_scheduling", [True, False]) @pytest.mark.parametrize("prompt_offset", [-1, -2]) def test_last_block_offloaded_at_request_finish( 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 5e98c1266e20..93a8e72a1f7a 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -327,16 +327,15 @@ def storable_chunks( num_chunks = max(0, num_chunks - 1) return num_chunks - def advance_stored_idx(self, num_offloadable_tokens: int) -> None: - # max(): at the prefill->decode transition of a chunk-aligned prompt, - # storable_chunks drops by one (the eagle exclusion kicks in), and the - # index must not move backwards past already-stored chunks. - for group_config, group_state in zip( - self.config.kv_group_configs, self.group_states + def advance_stored_idx(self, num_chunks_by_group: Sequence[int]) -> None: + # Keep the cursor monotonic: the EAGLE exclusion can lower a group's + # storable boundary at the prefill-to-decode transition. + for num_chunks, group_state in zip( + num_chunks_by_group, self.group_states, strict=True ): group_state.next_stored_chunk_idx = max( group_state.next_stored_chunk_idx, - self.storable_chunks(group_config, num_offloadable_tokens), + num_chunks, ) def update_num_hit_chunks(self, num_cached_tokens: int) -> None: @@ -950,14 +949,38 @@ def _build_store_jobs( # Filter out chunks skipped due to sliding window attention / SSM # or unreachable by the load path's alignment constraints. - new_offload_keys: list[OffloadKey] = [] + group_chunk_ends: list[int] = [] for group_config, group_state in zip( self.config.kv_group_configs, req_status.group_states ): - num_chunks = req_status.storable_chunks( + num_storable_chunks = req_status.storable_chunks( group_config, num_offloadable_tokens ) + num_tracked_chunks = len(group_state.block_ids) // blocks_per_chunk + num_chunks = min( + num_storable_chunks, + len(group_state.offload_keys), + num_tracked_chunks, + ) + group_chunk_ends.append(num_chunks) + if num_chunks < num_storable_chunks: + logger.debug( + "Request %s deferring group %d offload chunks: " + "storable=%d keys=%d tracked_blocks=%d", + req_id, + group_config.group_idx, + num_storable_chunks, + len(group_state.offload_keys), + len(group_state.block_ids), + ) + new_offload_keys: list[OffloadKey] = [] + for group_config, group_state, num_chunks in zip( + self.config.kv_group_configs, + req_status.group_states, + group_chunk_ends, + strict=True, + ): start_chunk_idx = group_state.next_stored_chunk_idx if num_chunks <= start_chunk_idx: continue @@ -972,7 +995,6 @@ def _build_store_jobs( + blocks_per_chunk - 1 : num_chunks * blocks_per_chunk : blocks_per_chunk ] - assert len(offload_keys) == len(offload_block_ids) alignment_chunk_count = group_config.alignment_chunk_count tail = group_config.sliding_window_size_in_chunks @@ -996,7 +1018,7 @@ def _build_store_jobs( new_offload_keys.append(offload_key) if not new_offload_keys: - req_status.advance_stored_idx(num_offloadable_tokens) + req_status.advance_stored_idx(group_chunk_ends) self._maybe_cleanup_finished_req(req_id, req_status) continue @@ -1012,7 +1034,7 @@ def _build_store_jobs( continue if not store_output.keys_to_store: - req_status.advance_stored_idx(num_offloadable_tokens) + req_status.advance_stored_idx(group_chunk_ends) self._maybe_cleanup_finished_req(req_id, req_status) continue @@ -1025,15 +1047,15 @@ def _build_store_jobs( src_block_ids: list[int] = [] sliding_window_block_ids: list[int] = [] non_sliding_window_block_ids: list[int] = [] - for group_config, group_state in zip( - self.config.kv_group_configs, req_status.group_states + for group_config, group_state, num_chunks in zip( + self.config.kv_group_configs, + req_status.group_states, + group_chunk_ends, + strict=True, ): is_sliding_window = ( group_config.sliding_window_size_in_chunks is not None ) - num_chunks = req_status.storable_chunks( - group_config, num_offloadable_tokens - ) start_chunk_idx = group_state.next_stored_chunk_idx block_ids = group_state.block_ids num_group_blocks = 0