[Bugfix][KV Offload] Validate SWA coverage at unaligned cache-hit boundaries (#55712)
Signed-off-by: xijiade.aihemaiti <3146335281@qq.com>
ahmed xijiaat committed
Sep 8, 2026 at 15:15 UTC
81dabe2e5d613d107577c9b4238fc8d718b92887
2 files changed
+161
-8
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+127
-4
@@ -34,6 +34,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
34
RequestOffloadState,
35
get_sliding_window_size_in_chunks,
36
)
37
+from vllm.v1.core.block_pool import BlockPool
38
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
39
from vllm.v1.core.kv_cache_utils import (
40
BlockHash,
@@ -45,6 +46,10 @@ from vllm.v1.core.sched.output import (
46
KVConnectorBlockState,
47
SchedulerOutput,
48
)
49
+from vllm.v1.core.single_type_kv_cache_manager import (
50
+ FullAttentionManager,
51
+ SlidingWindowManager,
52
+)
53
from vllm.v1.kv_cache_interface import (
54
ChunkedLocalAttentionSpec,
55
FullAttentionSpec,
@@ -59,6 +64,7 @@ from vllm.v1.kv_offload.base import (
64
LookupResult,
65
Medium,
66
OffloadingEvent,
67
+ OffloadingKVEventsConfig,
68
OffloadingManager,
69
OffloadPolicy,
70
ReqContext,
@@ -67,10 +73,106 @@ from vllm.v1.kv_offload.base import (
73
get_offload_group_idx,
74
make_offload_key,
75
)
76
+from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
77
from vllm.v1.outputs import KVConnectorOutput
78
from vllm.v1.request import RequestStatus
79
80
81
+@pytest.mark.parametrize("boundary", [3840, 3904])
82
+@pytest.mark.parametrize("eagle", [False, True])
83
+@pytest.mark.parametrize("left_state", ["present", "missing", "pending", "loading"])
84
+def test_swa_offload_window_covers_unaligned_hit(boundary, eagle, left_state):
85
+ """A smaller SWA group can move the hit inside another group's CPU chunk."""
86
+ groups = []
87
+ for i, (block_size, window) in enumerate([(256, None), (64, 128), (8, 128)]):
88
+ kwargs = dict(
89
+ block_size=block_size, num_kv_heads=1, head_size=1, dtype=torch.float32
90
+ )
91
+ kv_spec = (
92
+ FullAttentionSpec(**kwargs)
93
+ if window is None
94
+ else SlidingWindowSpec(**kwargs, sliding_window=window)
95
+ )
96
+ groups.append(
97
+ KVCacheGroupSpec([f"layer{i}"], kv_spec, is_eagle_group=eagle and i == 1)
98
+ )
99
+ manager = CPUOffloadingManager(num_blocks=100)
100
+ spec = SimpleNamespace(
101
+ tokens_per_block=(256, 64, 8),
102
+ tokens_per_hash=8,
103
+ blocks_per_chunk=4,
104
+ offload_prompt_only=True,
105
+ kv_events_config=OffloadingKVEventsConfig(
106
+ enable_kv_cache_events=False, self_describing_kv_events=False
107
+ ),
108
+ get_manager=lambda: manager,
109
+ )
110
+ config = SimpleNamespace(
111
+ speculative_config=None,
112
+ cache_config=SimpleNamespace(
113
+ enable_prefix_caching=True, prefix_cache_retention_interval=None
114
+ ),
115
+ parallel_config=SimpleNamespace(world_size=1, decode_context_parallel_size=1),
116
+ )
117
+ sched = OffloadingConnectorScheduler(
118
+ spec,
119
+ config,
120
+ KVCacheConfig(num_blocks=256, kv_cache_tensors=[], kv_cache_groups=groups),
121
+ )
122
+ request = MagicMock()
123
+ request.request_id = "unaligned"
124
+ request.kv_transfer_params = None
125
+ request.num_tokens = request.num_prompt_tokens = 4353
126
+ request.block_hashes = [BlockHash(f"h{i}".encode()) for i in range(544)]
127
+ request.skip_reading_prefix_cache = False
128
+ sched.on_new_request(request)
129
+ state = sched._req_status[request.request_id]
130
+ state.update_offload_keys()
131
+ fa, swa, short = (g.offload_keys for g in state.group_states)
132
+ end = (boundary + 255) // 256
133
+ left_key = swa[end - 2]
134
+ ready = fa[:4] + swa[end - 1 : end + int(eagle)]
135
+ ready += short[boundary // 32 - 4 : boundary // 32]
136
+ if left_state in ("present", "loading"):
137
+ ready.append(left_key)
138
+ store = manager.prepare_store(ready, state.req_context)
139
+ assert store is not None and not store.evicted_keys
140
+ manager.complete_store(store.keys_to_store, state.req_context)
141
+ if left_state == "pending":
142
+ assert manager.prepare_store([left_key], state.req_context) is not None
143
+ if left_state == "loading":
144
+ sched._chunks_being_loaded.add(left_key)
145
+
146
+ external, load_async = sched.get_num_new_matched_tokens(request, 0)
147
+ if boundary == 3904 and left_state != "present":
148
+ assert external == (0 if left_state == "missing" else None)
149
+ assert not load_async
150
+ return
151
+ assert external == boundary and load_async
152
+ pool = BlockPool(256, enable_caching=True, hash_block_size=8)
153
+ managers = []
154
+ for i, group in enumerate(groups):
155
+ cls = FullAttentionManager if i == 0 else SlidingWindowManager
156
+ m = cls(
157
+ kv_cache_spec=group.kv_cache_spec,
158
+ block_pool=pool,
159
+ enable_caching=True,
160
+ kv_cache_group_id=i,
161
+ scheduler_block_size=256,
162
+ )
163
+ m.add_local_computed_blocks(request.request_id, [], 0, external)
164
+ managers.append(m)
165
+ for m in managers:
166
+ m.allocate_external_computed_blocks(request.request_id, 0, external)
167
+ blocks = KVCacheBlocks(tuple(m.req_to_blocks[request.request_id] for m in managers))
168
+ sched.update_state_after_alloc(request, blocks, external)
169
+ assert len(sched._jobs) == 1
170
+ job = next(iter(sched._jobs.values()))
171
+ if boundary == 3904:
172
+ assert {swa[14], swa[15]} <= job.keys
173
+ manager.complete_load(job.keys, state.req_context)
174
+
175
+
176
def _make_partial_tail_scheduler() -> OffloadingConnectorScheduler:
177
vllm_config = _make_vllm_config(extra_config={"self_describing_kv_events": True})
178
vllm_config.cache_config.prefix_match_unit = 4
@@ -1418,6 +1520,27 @@ def test_scan_behavior_declared_for_every_lookup_result(result: LookupResult):
1520
assert len(sched.manager.lookup.call_args_list) == expected_lookups
1521
1522
1523
+@pytest.mark.parametrize(
1524
+ ("middle", "expected"),
1525
+ [
1526
+ (LookupResult.MISS, 2),
1527
+ (LookupResult.RETRY, None),
1528
+ (LookupResult.HIT_PENDING, None),
1529
+ (LookupResult.HIT, 5),
1530
+ ],
1531
+)
1532
+def test_sliding_window_unaligned_initial_run(middle, expected):
1533
+ sched = _make_scheduler_with_lookup({3: middle}, default=LookupResult.HIT)
1534
+ # If the larger rightmost window misses, the earlier aligned window only
1535
+ # needs the usual two chunks. Pending/uncertain keys must still defer.
1536
+ assert (
1537
+ sched._sliding_window_lookup(
1538
+ to_keys([1, 2, 3, 4, 5]), 2, _EMPTY_REQ_CTX, initial_window_size=3
1539
+ )
1540
+ == expected
1541
+ )
1542
+
1543
+
1544
class TestMaximalPrefixLookup:
1545
def test_all_hit(self):
1546
sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT})
@@ -2915,12 +3038,12 @@ class TestEagle:
3038
captured_keys: list = []
3039
orig_sw_lookup = type(sched)._sliding_window_lookup
3040
2918
- def capturing_sw_lookup(self_arg, keys, window, req_context):
3041
+ def capturing_sw_lookup(self_arg, keys, window, req_context, initial_window):
3042
captured_keys.append(list(keys))
2920
- return orig_sw_lookup(self_arg, keys, window, req_context)
3043
+ return orig_sw_lookup(self_arg, keys, window, req_context, initial_window)
3044
2922
- sched._sliding_window_lookup = lambda keys, window, req_ctx: (
2923
- capturing_sw_lookup(sched, keys, window, req_ctx)
3045
+ sched._sliding_window_lookup = lambda keys, window, req_ctx, initial_window: (
3046
+ capturing_sw_lookup(sched, keys, window, req_ctx, initial_window)
3047
)
3048
3049
req_status = self._make_req_status(
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+34
-4
@@ -104,6 +104,22 @@ class GroupOffloadConfig(NamedTuple):
104
# be excluded from store and load scheduling.
105
is_eagle_group: bool = False
106
107
+ def load_window_size_in_chunks(self, num_tokens: int) -> int | None:
108
+ window = self.sliding_window_size_in_chunks
109
+ if isinstance(self.kv_cache_spec, SlidingWindowSpec):
110
+ assert window is not None
111
+ # Lookup rounds the right edge up, but allocation retains the
112
+ # window ending at the actual hit boundary, possibly one chunk left.
113
+ right_padding = -num_tokens % self.tokens_per_chunk
114
+ window = max(
115
+ window,
116
+ cdiv(
117
+ self.kv_cache_spec.sliding_window - 1 + right_padding,
118
+ self.tokens_per_chunk,
119
+ ),
120
+ )
121
+ return window
122
+
123
124
def get_sliding_window_size_in_chunks(
125
kv_cache_spec: KVCacheSpec, tokens_per_chunk: int
@@ -657,12 +673,15 @@ class OffloadingConnectorScheduler:
673
keys: Sequence[OffloadKey],
674
sliding_window_size: int,
675
req_context: ReqContext,
676
+ initial_window_size: int | None = None,
677
) -> int | None:
678
"""Return the end index (in `keys`) of the last run of
679
`sliding_window_size` consecutive hits, scanning from the end.
680
+ The first run may need a larger window for a partial rightmost chunk.
681
Returns 0 on miss, None if the backend deferred a lookup."""
682
defer_lookup = False
683
consecutive_hits = 0
684
+ required_window = initial_window_size or sliding_window_size
685
for idx in range(len(keys) - 1, -1, -1):
686
match self.manager.lookup(keys[idx], req_context):
687
case LookupResult.HIT:
@@ -679,10 +698,12 @@ class OffloadingConnectorScheduler:
698
# async lookups.
699
defer_lookup = True
700
consecutive_hits = 0
701
+ required_window = sliding_window_size
702
case LookupResult.MISS:
703
consecutive_hits = 0
684
- if consecutive_hits == sliding_window_size:
685
- return idx + sliding_window_size if not defer_lookup else None
704
+ required_window = sliding_window_size
705
+ if consecutive_hits == required_window:
706
+ return idx + required_window if not defer_lookup else None
707
return consecutive_hits if not defer_lookup else None
708
709
def _touch(self, req_status: RequestOffloadState):
@@ -810,10 +831,19 @@ class OffloadingConnectorScheduler:
831
required_window = sliding_window_size_in_chunks
832
if is_eagle_unverified:
833
required_window += 1
834
+ candidate_end = min(
835
+ max_hit_size_tokens,
836
+ (num_chunks - int(is_eagle_unverified)) * tokens_per_chunk,
837
+ )
838
+ initial_window = group_config.load_window_size_in_chunks(
839
+ candidate_end
840
+ )
841
+ assert initial_window is not None
842
num_hit_chunks = self._sliding_window_lookup(
843
offload_keys,
844
required_window,
845
req_status.req_context,
846
+ initial_window + int(is_eagle_unverified),
847
)
848
if num_hit_chunks == 0:
849
return 0
@@ -864,8 +894,8 @@ class OffloadingConnectorScheduler:
894
self.config.kv_group_configs, req_status.group_states
895
):
896
tokens_per_chunk = group_config.tokens_per_chunk
867
- sliding_window_size_in_chunks = (
868
- group_config.sliding_window_size_in_chunks
897
+ sliding_window_size_in_chunks = group_config.load_window_size_in_chunks(
898
+ num_computed_tokens + num_hit_tokens
899
)
900
offload_keys = group_state.offload_keys
901
num_chunks = cdiv(