From 81dabe2e5d613d107577c9b4238fc8d718b92887 Mon Sep 17 00:00:00 2001 From: ahmed xijiaat <52128022+xijiaat@users.noreply.github.com> Date: Tue, 8 Sep 2026 15:15:48 +0800 Subject: [PATCH] [Bugfix][KV Offload] Validate SWA coverage at unaligned cache-hit boundaries (#55712) Signed-off-by: xijiade.aihemaiti <3146335281@qq.com> --- .../offloading_connector/test_scheduler.py | 131 +++++++++++++++++- .../kv_connector/v1/offloading/scheduler.py | 38 ++++- 2 files changed, 161 insertions(+), 8 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 cd9fb1d4ee..d01d0f258a 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -34,6 +34,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( RequestOffloadState, get_sliding_window_size_in_chunks, ) +from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.kv_cache_utils import ( BlockHash, @@ -45,6 +46,10 @@ from vllm.v1.core.sched.output import ( KVConnectorBlockState, SchedulerOutput, ) +from vllm.v1.core.single_type_kv_cache_manager import ( + FullAttentionManager, + SlidingWindowManager, +) from vllm.v1.kv_cache_interface import ( ChunkedLocalAttentionSpec, FullAttentionSpec, @@ -59,6 +64,7 @@ from vllm.v1.kv_offload.base import ( LookupResult, Medium, OffloadingEvent, + OffloadingKVEventsConfig, OffloadingManager, OffloadPolicy, ReqContext, @@ -67,10 +73,106 @@ from vllm.v1.kv_offload.base import ( get_offload_group_idx, make_offload_key, ) +from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import RequestStatus +@pytest.mark.parametrize("boundary", [3840, 3904]) +@pytest.mark.parametrize("eagle", [False, True]) +@pytest.mark.parametrize("left_state", ["present", "missing", "pending", "loading"]) +def test_swa_offload_window_covers_unaligned_hit(boundary, eagle, left_state): + """A smaller SWA group can move the hit inside another group's CPU chunk.""" + groups = [] + for i, (block_size, window) in enumerate([(256, None), (64, 128), (8, 128)]): + kwargs = dict( + block_size=block_size, num_kv_heads=1, head_size=1, dtype=torch.float32 + ) + kv_spec = ( + FullAttentionSpec(**kwargs) + if window is None + else SlidingWindowSpec(**kwargs, sliding_window=window) + ) + groups.append( + KVCacheGroupSpec([f"layer{i}"], kv_spec, is_eagle_group=eagle and i == 1) + ) + manager = CPUOffloadingManager(num_blocks=100) + spec = SimpleNamespace( + tokens_per_block=(256, 64, 8), + tokens_per_hash=8, + blocks_per_chunk=4, + offload_prompt_only=True, + kv_events_config=OffloadingKVEventsConfig( + enable_kv_cache_events=False, self_describing_kv_events=False + ), + get_manager=lambda: manager, + ) + config = SimpleNamespace( + speculative_config=None, + cache_config=SimpleNamespace( + enable_prefix_caching=True, prefix_cache_retention_interval=None + ), + parallel_config=SimpleNamespace(world_size=1, decode_context_parallel_size=1), + ) + sched = OffloadingConnectorScheduler( + spec, + config, + KVCacheConfig(num_blocks=256, kv_cache_tensors=[], kv_cache_groups=groups), + ) + request = MagicMock() + request.request_id = "unaligned" + request.kv_transfer_params = None + request.num_tokens = request.num_prompt_tokens = 4353 + request.block_hashes = [BlockHash(f"h{i}".encode()) for i in range(544)] + request.skip_reading_prefix_cache = False + sched.on_new_request(request) + state = sched._req_status[request.request_id] + state.update_offload_keys() + fa, swa, short = (g.offload_keys for g in state.group_states) + end = (boundary + 255) // 256 + left_key = swa[end - 2] + ready = fa[:4] + swa[end - 1 : end + int(eagle)] + ready += short[boundary // 32 - 4 : boundary // 32] + if left_state in ("present", "loading"): + ready.append(left_key) + store = manager.prepare_store(ready, state.req_context) + assert store is not None and not store.evicted_keys + manager.complete_store(store.keys_to_store, state.req_context) + if left_state == "pending": + assert manager.prepare_store([left_key], state.req_context) is not None + if left_state == "loading": + sched._chunks_being_loaded.add(left_key) + + external, load_async = sched.get_num_new_matched_tokens(request, 0) + if boundary == 3904 and left_state != "present": + assert external == (0 if left_state == "missing" else None) + assert not load_async + return + assert external == boundary and load_async + pool = BlockPool(256, enable_caching=True, hash_block_size=8) + managers = [] + for i, group in enumerate(groups): + cls = FullAttentionManager if i == 0 else SlidingWindowManager + m = cls( + kv_cache_spec=group.kv_cache_spec, + block_pool=pool, + enable_caching=True, + kv_cache_group_id=i, + scheduler_block_size=256, + ) + m.add_local_computed_blocks(request.request_id, [], 0, external) + managers.append(m) + for m in managers: + m.allocate_external_computed_blocks(request.request_id, 0, external) + blocks = KVCacheBlocks(tuple(m.req_to_blocks[request.request_id] for m in managers)) + sched.update_state_after_alloc(request, blocks, external) + assert len(sched._jobs) == 1 + job = next(iter(sched._jobs.values())) + if boundary == 3904: + assert {swa[14], swa[15]} <= job.keys + manager.complete_load(job.keys, state.req_context) + + def _make_partial_tail_scheduler() -> OffloadingConnectorScheduler: vllm_config = _make_vllm_config(extra_config={"self_describing_kv_events": True}) vllm_config.cache_config.prefix_match_unit = 4 @@ -1418,6 +1520,27 @@ def test_scan_behavior_declared_for_every_lookup_result(result: LookupResult): assert len(sched.manager.lookup.call_args_list) == expected_lookups +@pytest.mark.parametrize( + ("middle", "expected"), + [ + (LookupResult.MISS, 2), + (LookupResult.RETRY, None), + (LookupResult.HIT_PENDING, None), + (LookupResult.HIT, 5), + ], +) +def test_sliding_window_unaligned_initial_run(middle, expected): + sched = _make_scheduler_with_lookup({3: middle}, default=LookupResult.HIT) + # If the larger rightmost window misses, the earlier aligned window only + # needs the usual two chunks. Pending/uncertain keys must still defer. + assert ( + sched._sliding_window_lookup( + to_keys([1, 2, 3, 4, 5]), 2, _EMPTY_REQ_CTX, initial_window_size=3 + ) + == expected + ) + + class TestMaximalPrefixLookup: def test_all_hit(self): sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT}) @@ -2915,12 +3038,12 @@ class TestEagle: captured_keys: list = [] orig_sw_lookup = type(sched)._sliding_window_lookup - def capturing_sw_lookup(self_arg, keys, window, req_context): + def capturing_sw_lookup(self_arg, keys, window, req_context, initial_window): captured_keys.append(list(keys)) - return orig_sw_lookup(self_arg, keys, window, req_context) + return orig_sw_lookup(self_arg, keys, window, req_context, initial_window) - sched._sliding_window_lookup = lambda keys, window, req_ctx: ( - capturing_sw_lookup(sched, keys, window, req_ctx) + sched._sliding_window_lookup = lambda keys, window, req_ctx, initial_window: ( + capturing_sw_lookup(sched, keys, window, req_ctx, initial_window) ) req_status = self._make_req_status( 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 050aebe8f5..972eeff2ee 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -104,6 +104,22 @@ class GroupOffloadConfig(NamedTuple): # be excluded from store and load scheduling. is_eagle_group: bool = False + def load_window_size_in_chunks(self, num_tokens: int) -> int | None: + window = self.sliding_window_size_in_chunks + if isinstance(self.kv_cache_spec, SlidingWindowSpec): + assert window is not None + # Lookup rounds the right edge up, but allocation retains the + # window ending at the actual hit boundary, possibly one chunk left. + right_padding = -num_tokens % self.tokens_per_chunk + window = max( + window, + cdiv( + self.kv_cache_spec.sliding_window - 1 + right_padding, + self.tokens_per_chunk, + ), + ) + return window + def get_sliding_window_size_in_chunks( kv_cache_spec: KVCacheSpec, tokens_per_chunk: int @@ -657,12 +673,15 @@ class OffloadingConnectorScheduler: keys: Sequence[OffloadKey], sliding_window_size: int, req_context: ReqContext, + initial_window_size: int | None = None, ) -> int | None: """Return the end index (in `keys`) of the last run of `sliding_window_size` consecutive hits, scanning from the end. + 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 consecutive_hits = 0 + required_window = initial_window_size or sliding_window_size for idx in range(len(keys) - 1, -1, -1): match self.manager.lookup(keys[idx], req_context): case LookupResult.HIT: @@ -679,10 +698,12 @@ class OffloadingConnectorScheduler: # async lookups. defer_lookup = True consecutive_hits = 0 + required_window = sliding_window_size case LookupResult.MISS: consecutive_hits = 0 - if consecutive_hits == sliding_window_size: - return idx + sliding_window_size if not defer_lookup else None + 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 def _touch(self, req_status: RequestOffloadState): @@ -810,10 +831,19 @@ class OffloadingConnectorScheduler: required_window = sliding_window_size_in_chunks if is_eagle_unverified: required_window += 1 + candidate_end = min( + max_hit_size_tokens, + (num_chunks - int(is_eagle_unverified)) * tokens_per_chunk, + ) + initial_window = group_config.load_window_size_in_chunks( + candidate_end + ) + assert initial_window is not None num_hit_chunks = self._sliding_window_lookup( offload_keys, required_window, req_status.req_context, + initial_window + int(is_eagle_unverified), ) if num_hit_chunks == 0: return 0 @@ -864,8 +894,8 @@ class OffloadingConnectorScheduler: self.config.kv_group_configs, req_status.group_states ): tokens_per_chunk = group_config.tokens_per_chunk - sliding_window_size_in_chunks = ( - group_config.sliding_window_size_in_chunks + sliding_window_size_in_chunks = group_config.load_window_size_in_chunks( + num_computed_tokens + num_hit_tokens ) offload_keys = group_state.offload_keys num_chunks = cdiv(