@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d8de4ae322

[Bugfix][KVOffload] P2P tier declares REQUEST_LEVEL on the producer leg (#52912)

Signed-off-by: Liran Schour <lirans@il.ibm.com>

liranschour committed Aug 31, 2026 at 13:07 UTC d8de4ae322b54e518b2de771fe28673c6e9664d9
5 files changed +282 -9
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+235 -4
@@ -1293,12 +1293,21 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool
1293
1294 def _make_scheduler_with_lookup(
1295 lookup_results: dict[int, LookupResult],
1296 + default: LookupResult = LookupResult.MISS,
1297 ) -> OffloadingConnectorScheduler:
1297 - """Create an OffloadingConnectorScheduler with a mocked manager.lookup."""
1298 + """Create an OffloadingConnectorScheduler with a mocked manager.lookup.
1299 +
1300 + Keys are addressed by their integer hash; `default` answers anything else.
1301 + """
1302 manager = MagicMock(spec=OffloadingManager)
1299 - manager.lookup.side_effect = lambda key, req_context: lookup_results.get(
1300 - int(get_offload_block_hash(key).decode()), LookupResult.MISS
1301 - )
1303 +
1304 + def lookup(key, req_context):
1305 + block_hash = get_offload_block_hash(key)
1306 + if not block_hash.isdigit():
1307 + return default
1308 + return lookup_results.get(int(block_hash.decode()), default)
1309 +
1310 + manager.lookup.side_effect = lookup
1311
1312 scheduler = object.__new__(OffloadingConnectorScheduler)
1313 scheduler.manager = manager
@@ -1322,6 +1331,32 @@ def _maximal_lookup(sched, keys, start_chunk_idx: int = 0):
1331 )
1332
1333
1334 +# Lookups issued, and end index returned, by a window-1 scan over 3 keys all
1335 +# resolving to the same result. A result that keeps the streak alive ends the
1336 +# scan at the first key; one that resets it makes the scan walk every key,
1337 +# which is what widens the demanded set. A new member has to be added here.
1338 +_SCAN_BEHAVIOR = {
1339 + LookupResult.HIT: (1, 3),
1340 + LookupResult.HIT_PENDING: (1, None),
1341 + LookupResult.RETRY: (3, None),
1342 + LookupResult.MISS: (3, 0),
1343 +}
1344 +
1345 +
1346 +@pytest.mark.parametrize("result", list(LookupResult))
1347 +def test_scan_behavior_declared_for_every_lookup_result(result: LookupResult):
1348 + """Whether a result keeps a sliding-window streak alive decides how wide
1349 + the demanded chunk set gets, so every member needs deliberate behavior."""
1350 + assert result in _SCAN_BEHAVIOR, f"{result} has no declared scan behavior"
1351 + expected_lookups, expected_end = _SCAN_BEHAVIOR[result]
1352 +
1353 + keys = to_keys([1, 2, 3])
1354 + sched = _make_scheduler_with_lookup(dict.fromkeys([1, 2, 3], result))
1355 +
1356 + assert sched._sliding_window_lookup(keys, 1, _EMPTY_REQ_CTX) == expected_end
1357 + assert len(sched.manager.lookup.call_args_list) == expected_lookups
1358 +
1359 +
1360 class TestMaximalPrefixLookup:
1361 def test_all_hit(self):
1362 sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT})
@@ -1556,6 +1591,62 @@ class TestSlidingWindowLookup:
1591 )
1592
1593
1594 +# ---------------------------------------------------------------------------
1595 +# Tests for SWA store pruning vs. load demand
1596 +# ---------------------------------------------------------------------------
1597 +
1598 +
1599 +@pytest.mark.parametrize("alignment_chunk_count", [4, 8, 64])
1600 +@pytest.mark.parametrize("sliding_window_chunks", [1, 2, 3])
1601 +@pytest.mark.parametrize("is_eagle_group", [False, True])
1602 +@pytest.mark.parametrize("tail_chunks", [0, 1, 5])
1603 +def test_sliding_window_demand_is_store_reachable(
1604 + alignment_chunk_count: int,
1605 + sliding_window_chunks: int,
1606 + is_eagle_group: bool,
1607 + tail_chunks: int,
1608 +):
1609 + """Every chunk the load path looks up must be one the store path keeps.
1610 +
1611 + `is_store_reachable_swa_chunk` prunes chunks `_sliding_window_lookup` can
1612 + never query, but the two mirror each other by hand: this pins the relation,
1613 + where `test_is_store_reachable_swa_chunk` pins only the predicate's values.
1614 + A looked-up key is a demanded key, since a secondary-tier hit queues a
1615 + promotion immediately.
1616 + """
1617 + storable_chunk_count = alignment_chunk_count + tail_chunks
1618 + # Both sides widen by one for an unverified EAGLE group.
1619 + required_window = sliding_window_chunks + int(is_eagle_group)
1620 +
1621 + sched = _make_scheduler_with_lookup(
1622 + {i: LookupResult.HIT_PENDING for i in range(storable_chunk_count)}
1623 + )
1624 + sched._sliding_window_lookup(
1625 + to_keys(range(storable_chunk_count)), required_window, _EMPTY_REQ_CTX
1626 + )
1627 +
1628 + demanded = {
1629 + int(get_offload_block_hash(lookup_call.args[0]).decode())
1630 + for lookup_call in sched.manager.lookup.call_args_list
1631 + }
1632 + assert demanded, "the scan must query something"
1633 + unsuppliable = sorted(
1634 + chunk_idx
1635 + for chunk_idx in demanded
1636 + if not is_store_reachable_swa_chunk(
1637 + chunk_idx,
1638 + storable_chunk_count,
1639 + alignment_chunk_count,
1640 + sliding_window_chunks,
1641 + is_eagle_group,
1642 + )
1643 + )
1644 + assert not unsuppliable, (
1645 + f"chunks {unsuppliable} are demanded by the load path but pruned by the "
1646 + f"store path, so a warm producer can never supply them"
1647 + )
1648 +
1649 +
1650 @pytest.mark.parametrize("async_scheduling", [True, False])
1651 def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling: bool):
1652 """With REQUEST_LEVEL policy, all blocks are stored — including prefix hits."""
@@ -2308,6 +2399,146 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):
2399 )
2400
2401
2402 +def _full_and_swa_groups(
2403 + full_attn_block_size: int, swa_block_size: int, sliding_window: int
2404 +) -> list[KVCacheGroupSpec]:
2405 + """Group 0 full attention (MLA-like), group 1 sliding window."""
2406 + return [
2407 + KVCacheGroupSpec(
2408 + ["layer0"],
2409 + FullAttentionSpec(
2410 + block_size=full_attn_block_size,
2411 + num_kv_heads=1,
2412 + head_size=1,
2413 + dtype=torch.float32,
2414 + ),
2415 + ),
2416 + KVCacheGroupSpec(
2417 + ["layer1"],
2418 + SlidingWindowSpec(
2419 + block_size=swa_block_size,
2420 + num_kv_heads=1,
2421 + head_size=1,
2422 + dtype=torch.float32,
2423 + sliding_window=sliding_window,
2424 + ),
2425 + ),
2426 + ]
2427 +
2428 +
2429 +def _demanded_keys(runner, offload_keys_per_group: list[list], num_tokens: int) -> set:
2430 + """Keys a cold consumer's lookup scan would demand for `num_tokens`.
2431 +
2432 + Drives the real scan functions with every key resolving as a promoting
2433 + secondary tier resolves it.
2434 + """
2435 + scan = _make_scheduler_with_lookup({}, default=LookupResult.HIT_PENDING)
2436 + for group_config, offload_keys in zip(
2437 + runner.connector_scheduler.config.kv_group_configs, offload_keys_per_group
2438 + ):
2439 + num_chunks = num_tokens // group_config.tokens_per_chunk
2440 + keys = offload_keys[:num_chunks]
2441 + window = group_config.sliding_window_size_in_chunks
2442 + if window is None:
2443 + _maximal_lookup(scan, keys)
2444 + else:
2445 + scan._sliding_window_lookup(keys, window, _EMPTY_REQ_CTX)
2446 + return {lookup_call.args[0] for lookup_call in scan.manager.lookup.call_args_list}
2447 +
2448 +
2449 +@pytest.mark.parametrize("warmth", ["cold", "gpu_warm", "primary_warm"])
2450 +@pytest.mark.parametrize("with_swa_group", [False, True])
2451 +def test_request_level_supply_covers_consumer_demand(
2452 + request_runner, warmth: str, with_swa_group: bool
2453 +):
2454 + """A REQUEST_LEVEL producer must offer every key a consumer will demand.
2455 +
2456 + Same relation as `test_sliding_window_demand_is_store_reachable`, but at
2457 + the level where the key set is really built: `_build_store_jobs` prunes
2458 + chunks with a null GPU block and unreachable SWA chunks, while the demand
2459 + comes from an independent scan.
2460 + """
2461 + full_attn_block_size = 16
2462 + swa_block_size = 4
2463 + num_tokens = 32
2464 +
2465 + kv_cache_groups = (
2466 + _full_and_swa_groups(full_attn_block_size, swa_block_size, sliding_window=8)
2467 + if with_swa_group
2468 + else None
2469 + )
2470 + runner = request_runner(
2471 + block_size=swa_block_size if with_swa_group else full_attn_block_size,
2472 + num_gpu_blocks=200,
2473 + async_scheduling=False,
2474 + kv_cache_groups=kv_cache_groups,
2475 + )
2476 +
2477 + def store_everything_offered():
2478 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (
2479 + generate_store_output(keys)
2480 + )
2481 +
2482 + # `_run` rather than `run`, which asserts on an expected GPU-block set this
2483 + # test deliberately does not hard-code.
2484 + if warmth != "cold":
2485 + runner.new_request(token_ids=[0] * num_tokens)
2486 + store_everything_offered()
2487 + runner._run(decoded_tokens=[0], complete_transfers=True)
2488 + runner._run(decoded_tokens=[EOS_TOKEN_ID], complete_transfers=True)
2489 + if warmth == "primary_warm":
2490 + # Warm only in the primary tier, so the producer must load first.
2491 + runner.scheduler.reset_prefix_cache()
2492 +
2493 + # The producer leg: same prompt, REQUEST_LEVEL, so prefix hits stay in the
2494 + # store path instead of being skipped.
2495 + runner.manager.reset_mock()
2496 + store_everything_offered()
2497 + runner.manager.on_new_request.return_value = RequestOffloadingContext(
2498 + policy=OffloadPolicy.REQUEST_LEVEL
2499 + )
2500 + if warmth == "primary_warm":
2501 + # HIT is final; HIT_PENDING would defer the request on every step.
2502 + runner.manager.lookup.return_value = LookupResult.HIT
2503 + runner.new_request(token_ids=[0] * num_tokens)
2504 +
2505 + offload_keys_per_group: list[list] = []
2506 +
2507 + def capture_offload_keys() -> None:
2508 + for req_status in runner.connector_scheduler._req_status.values():
2509 + offload_keys_per_group.clear()
2510 + offload_keys_per_group.extend(
2511 + list(group_state.offload_keys)
2512 + for group_state in req_status.group_states
2513 + )
2514 +
2515 + runner._run(
2516 + decoded_tokens=[0],
2517 + complete_transfers=True,
2518 + post_step_fn=capture_offload_keys,
2519 + )
2520 + runner._run(
2521 + decoded_tokens=[EOS_TOKEN_ID],
2522 + complete_transfers=True,
2523 + post_step_fn=capture_offload_keys,
2524 + )
2525 +
2526 + assert offload_keys_per_group, "no request state was captured"
2527 + supplied = {
2528 + key
2529 + for store_call in runner.manager.prepare_store.call_args_list
2530 + for key in store_call.args[0]
2531 + }
2532 + demanded = _demanded_keys(runner, offload_keys_per_group, num_tokens)
2533 +
2534 + assert demanded, "the scan must demand something"
2535 + missing = demanded - supplied
2536 + assert not missing, (
2537 + f"{len(missing)} of {len(demanded)} demanded keys are never offered to "
2538 + f"prepare_store, so a peer fetching them would wait out the load timeout"
2539 + )
2540 +
2541 +
2542 @pytest.mark.parametrize("async_scheduling", [True, False])
2543 def test_stale_sliding_window_block_after_prepare_store_failure(
2544 request_runner, async_scheduling: bool
tests/v1/kv_connector/unit/offloading_connector/utils.py
+1 -1
@@ -65,7 +65,7 @@ def to_key(int_hash: int) -> OffloadKey:
65 return make_offload_key(str(int_hash).encode(), 0)
66
67
68 -def to_keys(int_hashes: list[int]) -> list[OffloadKey]:
68 +def to_keys(int_hashes: Iterable[int]) -> list[OffloadKey]:
69 return [to_key(i) for i in int_hashes]
70
71
tests/v1/kv_offload/tiering/p2p/test_manager.py
+28 -1
@@ -17,7 +17,12 @@ import pytest
17
18 from vllm.utils.hashing import sha256
19 from vllm.v1.core.kv_cache_utils import DEFAULT_NONE_HASH_SEED, init_none_hash
20 -from vllm.v1.kv_offload.base import LookupResult, ReqContext, ScheduleEndContext
20 +from vllm.v1.kv_offload.base import (
21 + LookupResult,
22 + OffloadPolicy,
23 + ReqContext,
24 + ScheduleEndContext,
25 +)
26 from vllm.v1.kv_offload.tiering.base import JobResult, TransferJob
27 from vllm.v1.kv_offload.tiering.p2p import manager as manager_module
28 from vllm.v1.kv_offload.tiering.p2p.manager import (
@@ -240,6 +245,28 @@ class TestLookup:
245 assert mgr.lookup(b"key", ctx) is LookupResult.MISS
246
247
248 +# ---------------------------------------------------------------------------
249 +# Tests for on_new_request offload policy
250 +# ---------------------------------------------------------------------------
251 +
252 +
253 +@pytest.mark.parametrize(
254 + "kv_params,expected",
255 + [
256 + (_remote_decoder_kv_params(), OffloadPolicy.REQUEST_LEVEL),
257 + (None, OffloadPolicy.BLOCK_LEVEL),
258 + (_remote_prefiller_kv_params(), OffloadPolicy.BLOCK_LEVEL),
259 + ({"remote_decoder": {}}, OffloadPolicy.BLOCK_LEVEL),
260 + ],
261 + ids=["producer", "plain", "consumer", "producer_no_id"],
262 +)
263 +def test_on_new_request_policy(monkeypatch, kv_params, expected):
264 + """Only a producer leg carrying a kv_request_id widens to REQUEST_LEVEL."""
265 + mgr = _make_manager()
266 + monkeypatch.setattr(mgr, "_get_or_create_session", lambda peer_id: None)
267 + assert mgr.on_new_request(_req_context(kv_params=kv_params)).policy is expected
268 +
269 +
270 # ---------------------------------------------------------------------------
271 # Tests for serve_external_requests
272 # ---------------------------------------------------------------------------
tests/v1/kv_offload/tiering/test_tiering_offloading.py
+10 -3
@@ -1063,10 +1063,17 @@ class TestTieringOffloadingManager:
1063 self.manager.on_request_finished(ctx)
1064 assert ctx.req_id not in self.manager._req_state
1065
1066 + @pytest.mark.parametrize("new_ids", [(), (3, 4)], ids=["fully_warm", "mixed"])
1067 def test_prepare_store_cascades_existing_blocks_to_request_level_tiers(
1067 - self, manager_setup
1068 + self, manager_setup, new_ids
1069 ):
1069 - """prepare_store cascades hit blocks to request-level tiers only."""
1070 + """prepare_store cascades hit blocks to request-level tiers only.
1071 +
1072 + The fully-warm case is the producer shape from issue #52808: nothing is
1073 + new, so primary returns an empty keys_to_store and no GPU->primary
1074 + transfer runs. A request-level tier must still be handed the blocks,
1075 + otherwise a peer waiting on them has nothing to fetch.
1076 + """
1077 # Store some blocks to primary first
1078 existing_blocks = to_keys(range(3))
1079 self._start_request()
@@ -1093,7 +1100,7 @@ class TestTieringOffloadingManager:
1100 )
1101
1102 # Call prepare_store with existing + new blocks
1096 - new_blocks = to_keys(range(3, 5))
1103 + new_blocks = to_keys(new_ids)
1104 all_blocks = existing_blocks + new_blocks
1105 result = self.manager.prepare_store(all_blocks, ctx)
1106 assert result is not None
vllm/v1/kv_offload/tiering/p2p/manager.py
+8
@@ -22,6 +22,7 @@ from vllm.v1.core.kv_cache_utils import get_none_hash_seed
22 from vllm.v1.kv_offload.base import (
23 LookupResult,
24 OffloadKey,
25 + OffloadPolicy,
26 ReqContext,
27 RequestOffloadingContext,
28 ScheduleEndContext,
@@ -361,11 +362,18 @@ class P2PSecondaryTierManager(SecondaryTierManager):
362 prefiller side, sessions are created when the consumer's inbound
363 connection arrives in _accept_new_peers — submit_store no longer
364 pre-creates anything.
365 +
366 + Producer-leg requests (``remote_decoder`` present) ask for
367 + REQUEST_LEVEL: the peer needs every block of the request, not just
368 + the ones this request computed.
369 """
370 _annotate_req_context(req_context)
371 source = req_context.get_state(P2PSourceInfo)
372 if source is not None:
373 self._get_or_create_session(source.peer_id)
374 + dest = req_context.get_state(P2PDestInfo)
375 + if dest is not None and dest.kv_request_id:
376 + return RequestOffloadingContext(policy=OffloadPolicy.REQUEST_LEVEL)
377 return RequestOffloadingContext()
378
379 @override