From e45bb43984b60fbbddd6d2ba4acd6e4a5525f49c Mon Sep 17 00:00:00 2001 From: Liqian Lin Date: Tue, 15 Sep 2026 10:48:38 +0800 Subject: [PATCH] [Bugfix][KV Offload] Ignore pending chunks in invalid sliding windows (#56486) Signed-off-by: Liqian Lin --- .../offloading_connector/test_scheduler.py | 78 ++++++++++++++++++- .../kv_connector/v1/offloading/scheduler.py | 12 ++- 2 files changed, 86 insertions(+), 4 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 26664e3ac5..e7a83cb7b5 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -74,7 +74,7 @@ from vllm.v1.kv_offload.base import ( make_offload_key, ) from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager -from vllm.v1.outputs import KVConnectorOutput +from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput from vllm.v1.request import RequestStatus @@ -1708,6 +1708,15 @@ class TestMaximalPrefixLookup: class TestSlidingWindowLookup: + def test_pending_chunk_behind_gap_does_not_defer(self): + """A pending suffix too short for the window cannot improve the hit.""" + sched = _make_scheduler_with_lookup( + {1: LookupResult.HIT, 2: LookupResult.HIT, 4: LookupResult.HIT_PENDING} + ) + assert ( + sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX) == 2 + ) + def test_all_hit_exact_window(self): sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT}) assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) == 2 @@ -1826,6 +1835,73 @@ class TestSlidingWindowLookup: # --------------------------------------------------------------------------- +@pytest.mark.parametrize("async_scheduling", [False, True]) +def test_swa_load_does_not_wait_for_unusable_pending_suffix( + request_runner, async_scheduling +): + """A CPU store beyond a missing SWA chunk must not block a ready prefix.""" + groups = [ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=4, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + ), + KVCacheGroupSpec( + ["swa"], + SlidingWindowSpec( + block_size=4, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=8, + ), + ), + ] + runner = request_runner( + block_size=4, + num_gpu_blocks=32, + async_scheduling=async_scheduling, + kv_cache_groups=groups, + ) + sched = runner.connector_scheduler + manager = CPUOffloadingManager(num_chunks=16) + sched.manager = manager + runner.new_request(token_ids=list(range(17))) + state = sched._req_status["0"] + state.update_offload_keys() + full, swa = (g.offload_keys for g in state.group_states) + ready = manager.prepare_store(full[:4] + swa[:2], state.req_context) + assert ready is not None + manager.complete_store(ready.keys_to_store, state.req_context) + pending = manager.prepare_store(swa[3:4], state.req_context) + assert pending is not None + + # The last SWA chunk is still being written; the preceding one is missing. + output = runner.scheduler.schedule() + metadata = output.kv_connector_metadata + assert isinstance(metadata, OffloadingConnectorMetadata) + [(job_id, load)] = metadata.load_jobs.items() + assert isinstance(load.dst_spec, GPULoadStoreSpec) + assert load.dst_spec.group_sizes == [2, 2] + assert manager.lookup(swa[3], state.req_context) is LookupResult.HIT_PENDING + + runner.scheduler.update_from_output( + output, + ModelRunnerOutput.with_kv_conn_output_only( + KVConnectorOutput( + finished_recving={"0"}, + kv_connector_worker_meta=OffloadingWorkerMetadata( + completed_jobs={job_id: 1} + ), + ) + ), + ) + resumed = runner.scheduler.schedule() + assert resumed.num_scheduled_tokens == {"0": 9} + assert manager.lookup(swa[3], state.req_context) is LookupResult.HIT_PENDING + + @pytest.mark.parametrize("alignment_chunk_count", [4, 8, 64]) @pytest.mark.parametrize("sliding_window_chunks", [1, 2, 3]) @pytest.mark.parametrize("is_eagle_group", [False, True]) 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 86e3510939..2fb3c5bb90 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -685,6 +685,7 @@ class OffloadingConnectorScheduler: The first run may need a larger window for a partial rightmost chunk. Returns 0 on miss, None if the backend deferred a lookup.""" defer_lookup = False + pending_in_window = False consecutive_hits = 0 required_window = initial_window_size or sliding_window_size for idx in range(len(keys) - 1, -1, -1): @@ -695,7 +696,7 @@ class OffloadingConnectorScheduler: # Block is in cache, just not readable yet — counts # as hit for the consecutive streak. Don't break: # keep scanning to let manager kick off async lookups. - defer_lookup = True + pending_in_window = True consecutive_hits += 1 case LookupResult.RETRY: # Block location uncertain — does not count as hit. @@ -703,13 +704,18 @@ class OffloadingConnectorScheduler: # async lookups. defer_lookup = True consecutive_hits = 0 + pending_in_window = False required_window = sliding_window_size case LookupResult.MISS: consecutive_hits = 0 + # This gap rules out the incomplete window to its right. + pending_in_window = False required_window = sliding_window_size if consecutive_hits == required_window: - return idx + required_window if not defer_lookup else None - return consecutive_hits if not defer_lookup else None + return ( + None if defer_lookup or pending_in_window else idx + required_window + ) + return None if defer_lookup or pending_in_window else consecutive_hits def _touch(self, req_status: RequestOffloadState): for group_config, group_state in zip(