[Bugfix][KV Offload] Ignore pending chunks in invalid sliding windows (#56486)

Signed-off-by: Liqian Lin <jacklin78911@gmail.com>

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])

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(