iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix][KV Offload] Validate SWA coverage at unaligned cache-hit boundaries (#55712)

Browse Source 
Signed-off-by: xijiade.aihemaiti <3146335281@qq.com>
This commit is contained in:
ahmed xijiaat
2026-09-08 10:15:48 +03:00
committed by GitHub
parent 5133e1d285
commit 81dabe2e5d
2 changed files with 161 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+127 -4
View File
@@ -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(
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+34 -4
View File
@@ -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(