[Bugfix][KV Offload] Ignore pending chunks in invalid sliding windows (#56486)
Signed-off-by: Liqian Lin <jacklin78911@gmail.com>
Liqian Lin committed
Sep 15, 2026 at 10:48 UTC
e45bb43984b60fbbddd6d2ba4acd6e4a5525f49c
2 files changed
+86
-4
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+77
-1
@@ -74,7 +74,7 @@ from vllm.v1.kv_offload.base import (
74
make_offload_key,
75
)
76
from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
77
-from vllm.v1.outputs import KVConnectorOutput
77
+from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
78
from vllm.v1.request import RequestStatus
79
80
@@ -1708,6 +1708,15 @@ class TestMaximalPrefixLookup:
1708
1709
1710
class TestSlidingWindowLookup:
1711
+ def test_pending_chunk_behind_gap_does_not_defer(self):
1712
+ """A pending suffix too short for the window cannot improve the hit."""
1713
+ sched = _make_scheduler_with_lookup(
1714
+ {1: LookupResult.HIT, 2: LookupResult.HIT, 4: LookupResult.HIT_PENDING}
1715
+ )
1716
+ assert (
1717
+ sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX) == 2
1718
+ )
1719
+
1720
def test_all_hit_exact_window(self):
1721
sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT})
1722
assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) == 2
@@ -1826,6 +1835,73 @@ class TestSlidingWindowLookup:
1835
# ---------------------------------------------------------------------------
1836
1837
1838
+@pytest.mark.parametrize("async_scheduling", [False, True])
1839
+def test_swa_load_does_not_wait_for_unusable_pending_suffix(
1840
+ request_runner, async_scheduling
1841
+):
1842
+ """A CPU store beyond a missing SWA chunk must not block a ready prefix."""
1843
+ groups = [
1844
+ KVCacheGroupSpec(
1845
+ ["full"],
1846
+ FullAttentionSpec(
1847
+ block_size=4, num_kv_heads=1, head_size=1, dtype=torch.float32
1848
+ ),
1849
+ ),
1850
+ KVCacheGroupSpec(
1851
+ ["swa"],
1852
+ SlidingWindowSpec(
1853
+ block_size=4,
1854
+ num_kv_heads=1,
1855
+ head_size=1,
1856
+ dtype=torch.float32,
1857
+ sliding_window=8,
1858
+ ),
1859
+ ),
1860
+ ]
1861
+ runner = request_runner(
1862
+ block_size=4,
1863
+ num_gpu_blocks=32,
1864
+ async_scheduling=async_scheduling,
1865
+ kv_cache_groups=groups,
1866
+ )
1867
+ sched = runner.connector_scheduler
1868
+ manager = CPUOffloadingManager(num_chunks=16)
1869
+ sched.manager = manager
1870
+ runner.new_request(token_ids=list(range(17)))
1871
+ state = sched._req_status["0"]
1872
+ state.update_offload_keys()
1873
+ full, swa = (g.offload_keys for g in state.group_states)
1874
+ ready = manager.prepare_store(full[:4] + swa[:2], state.req_context)
1875
+ assert ready is not None
1876
+ manager.complete_store(ready.keys_to_store, state.req_context)
1877
+ pending = manager.prepare_store(swa[3:4], state.req_context)
1878
+ assert pending is not None
1879
+
1880
+ # The last SWA chunk is still being written; the preceding one is missing.
1881
+ output = runner.scheduler.schedule()
1882
+ metadata = output.kv_connector_metadata
1883
+ assert isinstance(metadata, OffloadingConnectorMetadata)
1884
+ [(job_id, load)] = metadata.load_jobs.items()
1885
+ assert isinstance(load.dst_spec, GPULoadStoreSpec)
1886
+ assert load.dst_spec.group_sizes == [2, 2]
1887
+ assert manager.lookup(swa[3], state.req_context) is LookupResult.HIT_PENDING
1888
+
1889
+ runner.scheduler.update_from_output(
1890
+ output,
1891
+ ModelRunnerOutput.with_kv_conn_output_only(
1892
+ KVConnectorOutput(
1893
+ finished_recving={"0"},
1894
+ kv_connector_worker_meta=OffloadingWorkerMetadata(
1895
+ completed_jobs={job_id: 1}
1896
+ ),
1897
+ )
1898
+ ),
1899
+ )
1900
+ resumed = runner.scheduler.schedule()
1901
+ assert resumed.num_scheduled_tokens == {"0": 9}
1902
+ assert manager.lookup(swa[3], state.req_context) is LookupResult.HIT_PENDING
1903
+
1904
+
1905
@pytest.mark.parametrize("alignment_chunk_count", [4, 8, 64])
1906
@pytest.mark.parametrize("sliding_window_chunks", [1, 2, 3])
1907
@pytest.mark.parametrize("is_eagle_group", [False, True])
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+9
-3
@@ -685,6 +685,7 @@ class OffloadingConnectorScheduler:
685
The first run may need a larger window for a partial rightmost chunk.
686
Returns 0 on miss, None if the backend deferred a lookup."""
687
defer_lookup = False
688
+ pending_in_window = False
689
consecutive_hits = 0
690
required_window = initial_window_size or sliding_window_size
691
for idx in range(len(keys) - 1, -1, -1):
@@ -695,7 +696,7 @@ class OffloadingConnectorScheduler:
696
# Block is in cache, just not readable yet — counts
697
# as hit for the consecutive streak. Don't break:
698
# keep scanning to let manager kick off async lookups.
698
- defer_lookup = True
699
+ pending_in_window = True
700
consecutive_hits += 1
701
case LookupResult.RETRY:
702
# Block location uncertain — does not count as hit.
@@ -703,13 +704,18 @@ class OffloadingConnectorScheduler:
704
# async lookups.
705
defer_lookup = True
706
consecutive_hits = 0
707
+ pending_in_window = False
708
required_window = sliding_window_size
709
case LookupResult.MISS:
710
consecutive_hits = 0
711
+ # This gap rules out the incomplete window to its right.
712
+ pending_in_window = False
713
required_window = sliding_window_size
714
if consecutive_hits == required_window:
711
- return idx + required_window if not defer_lookup else None
712
- return consecutive_hits if not defer_lookup else None
715
+ return (
716
+ None if defer_lookup or pending_in_window else idx + required_window
717
+ )
718
+ return None if defer_lookup or pending_in_window else consecutive_hits
719
720
def _touch(self, req_status: RequestOffloadState):
721
for group_config, group_state in zip(