From b2662244f399f35f4e90e55f24f06c8d6d741cc9 Mon Sep 17 00:00:00 2001 From: Ketor Date: Thu, 27 Aug 2026 10:04:29 +0800 Subject: [PATCH] fix(vllm): preserve hybrid block tables for external KV --- integration/vllm/src/dfkv_vllm/coordinator.py | 45 ++++++++-------- integration/vllm/src/dfkv_vllm/scheduler.py | 20 +++---- integration/vllm/src/dfkv_vllm/worker.py | 54 ++++++++++++------- .../tests/test_coordinator_stateful_mask.py | 26 +++++++++ .../vllm/tests/test_hybrid_pool_layout.py | 11 ++++ .../tests/test_scheduler_full_block_ids.py | 43 +++++++++++++++ 6 files changed, 145 insertions(+), 54 deletions(-) create mode 100644 integration/vllm/tests/test_coordinator_stateful_mask.py create mode 100644 integration/vllm/tests/test_scheduler_full_block_ids.py diff --git a/integration/vllm/src/dfkv_vllm/coordinator.py b/integration/vllm/src/dfkv_vllm/coordinator.py index 3f89d7bb..529bda45 100644 --- a/integration/vllm/src/dfkv_vllm/coordinator.py +++ b/integration/vllm/src/dfkv_vllm/coordinator.py @@ -202,32 +202,29 @@ def store_mask(self, aligned_token_len: int) -> tuple[list[bool], ...]: if aligned_token_len == 0: return tuple([] for _ in self.kv_cache_groups) - num_chunks_per_group = [ - aligned_token_len // g.kv_cache_spec.block_size - for g in self.kv_cache_groups - ] - - # Fast path: single group or full attn groups or uniform block_sizes - if all( - isinstance(spec, FullAttentionSpec) - or spec.block_size == self.lcm_block_size - for spec, _, _ in self.attention_groups - ): - return tuple([True] * n for n in num_chunks_per_group) - - n_segments = aligned_token_len // self.lcm_block_size dummy_hashes: list[BlockHash] = [_DUMMY_BLOCK_HASH] * ( - self.lcm_block_size // self.hash_block_size - ) - template_masks, _ = self.find_longest_cache_hit( - dummy_hashes, - max_length=self.lcm_block_size, - cached_block_pool=ExternalCachedBlockPool(), - ) - return tuple( - list(template_masks[g]) * n_segments - for g in range(len(self.kv_cache_groups)) + aligned_token_len // self.hash_block_size ) + block_pool = ExternalCachedBlockPool(hash_block_size=self.hash_block_size) + masks: list[list[bool]] = [[] for _ in self.kv_cache_groups] + for idx, (spec, group_ids, manager_cls) in enumerate(self.attention_groups): + hashes = self.block_hashes_for_spec(dummy_hashes, spec) + hit_blocks = _unwrap_hit_blocks( + manager_cls.find_longest_cache_hit( + block_hashes=hashes, + max_length=aligned_token_len, + kv_cache_group_ids=group_ids, + block_pool=cast(BlockPool, block_pool), + kv_cache_spec=spec, + drop_eagle_block=idx in self.eagle_attn_group_indices, + alignment_tokens=self.lcm_block_size, + ) + ) + for group_id, blocks in zip(group_ids, hit_blocks, strict=True): + masks[group_id] = [ + block is not block_pool.null_block for block in blocks + ] + return tuple(masks) def block_hashes_for_spec( self, block_hashes: list[BlockHash], spec: KVCacheSpec diff --git a/integration/vllm/src/dfkv_vllm/scheduler.py b/integration/vllm/src/dfkv_vllm/scheduler.py index 36dd62df..8f9390cf 100644 --- a/integration/vllm/src/dfkv_vllm/scheduler.py +++ b/integration/vllm/src/dfkv_vllm/scheduler.py @@ -148,9 +148,7 @@ def update_state_after_alloc( num_external_tokens: int, ): """Update state after block allocation.""" - local_block_ids: tuple[list[int], ...] = () - if num_external_tokens > 0: - local_block_ids = blocks.get_block_ids() + local_block_ids = blocks.get_block_ids() self._unfinished_requests[request.request_id] = (request, local_block_ids) self._unfinished_request_ids.add(request.request_id) @@ -212,15 +210,13 @@ def build_connector_meta( + scheduler_output.num_scheduled_tokens[request.req_id] ) assert request.req_id in self._unfinished_requests - request_tuple = self._unfinished_requests.get(request.req_id) - request_real = request_tuple[0] # type: ignore[index] - - if isinstance(request.block_ids, tuple): - # Multi-group: preserve per-group structure. - unfolded_block_ids = tuple(b.copy() for b in request.block_ids) - else: - # Single-group legacy: list[int] -> 1-tuple. - unfolded_block_ids = (request.block_ids.copy(),) + request_real, allocated_block_ids = self._unfinished_requests[request.req_id] + # NewRequestData.block_ids contains only blocks allocated in this + # scheduler step. Keep the complete table captured after allocation + # so later chunks can address from sequence offset zero. + unfolded_block_ids = tuple( + group_ids.copy() for group_ids in allocated_block_ids + ) prefill_tokens = _new_req_prefill_tokens(request) request_tracker = RequestTracker( diff --git a/integration/vllm/src/dfkv_vllm/worker.py b/integration/vllm/src/dfkv_vllm/worker.py index b11a9ae2..101a1b5d 100644 --- a/integration/vllm/src/dfkv_vllm/worker.py +++ b/integration/vllm/src/dfkv_vllm/worker.py @@ -59,10 +59,7 @@ from .client_ranks import ENV_NAME as CLIENT_RANKS_ENV from .client_ranks import (ELIDE_ENV, participant, resolve_client_ranks, should_create_client) -from .coordinator import ( - ExternalCachedBlockPool, - DfkvStoreCoordinator, -) +from .coordinator import DfkvStoreCoordinator from .data import ( VLLM_MULTIWR_V2, ChunkedTokenDatabase, @@ -216,6 +213,26 @@ def _key_stripe_identity( def _put_failed_indices(rcs: list[int]) -> list[int]: return [i for i, rc in enumerate(rcs) if rc != 0] +def _logical_block_ids( + mask: Sequence[bool], + block_ids: Sequence[int], +) -> list[int]: + positions = [index for index, selected in enumerate(mask) if selected] + if len(block_ids) < len(positions): + raise ValueError( + f"block table has {len(block_ids)} ids for " + f"{len(positions)} selected state slots" + ) + logical = [-1] * len(mask) + for position, block_id in zip( + positions, + block_ids[-len(positions) :] if positions else (), + strict=True, + ): + logical[position] = int(block_id) + return logical + + # ============================================================ # Transfer Threads @@ -565,8 +582,12 @@ def _handle_request(self, req_meta: ReqMeta): keys: list[bytes] = [] block_hashes: list[BlockHash] = [] group_indices: list[int] = [] + logical_block_ids_per_group: list[list[int]] = [] for g_idx, db in enumerate(self.token_databases): mask = store_masks[g_idx] + logical_block_ids_per_group.append( + _logical_block_ids(mask, block_ids_per_group[g_idx]) + ) for chunk_idx, (start, end, key) in enumerate( db.process_tokens(token_len, req_meta.block_hashes) ): @@ -646,7 +667,7 @@ def _handle_request(self, req_meta: ReqMeta): self.token_databases[g_idx], start, end, - block_ids_per_group[g_idx], + logical_block_ids_per_group[g_idx], ) for start, end, g_idx in zip( starts, ends, group_indices, strict=True @@ -1095,18 +1116,19 @@ def _handle_request(self, req_meta: ReqMeta): block_id_list: list[int] = [] for g_idx, db in enumerate(self.token_databases): mask = load_mask_per_group[g_idx] + logical_block_ids = _logical_block_ids( + mask, req_meta.block_ids[g_idx] + ) for start, end, key in db.process_tokens( token_len, req_meta.block_hashes, mask_num ): chunk_idx = start // db.block_size if chunk_idx >= len(mask) or not mask[chunk_idx]: continue - block_id = int( - req_meta.block_ids[g_idx][start // db.block_size] - ) + block_id = logical_block_ids[chunk_idx] key_list.append(key.to_bytes()) descriptor_chunks.append( - (db, start, end, req_meta.block_ids[g_idx]) + (db, start, end, logical_block_ids) ) block_id_list.append(block_id) @@ -2175,15 +2197,11 @@ def lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int: prefix_chunks += 1 complete_prefix_tokens = prefix_chunks * self.coord.lcm_block_size - if complete_prefix_tokens == 0: - hit_length = 0 - else: - _masks, semantic_hit_length = self.coord.find_longest_cache_hit( - block_hashes, - complete_prefix_tokens, - ExternalCachedBlockPool(exists_set), - ) - hit_length = min(complete_prefix_tokens, semantic_hit_length) + # candidate_keys are derived from the same manager-selected store mask + # used by save and load. Once every required object for the contiguous + # prefix exists, re-running cross-group convergence can only discard + # intentionally sparse stateful groups. + hit_length = complete_prefix_tokens logger.debug( "dfkv lookup: token_len=%d candidates=%d complete_chunks=%d/%d " "-> hit_length=%d", diff --git a/integration/vllm/tests/test_coordinator_stateful_mask.py b/integration/vllm/tests/test_coordinator_stateful_mask.py new file mode 100644 index 00000000..dc3b5efb --- /dev/null +++ b/integration/vllm/tests/test_coordinator_stateful_mask.py @@ -0,0 +1,26 @@ +from types import SimpleNamespace + +from dfkv_vllm.coordinator import DfkvStoreCoordinator + + +def test_equal_block_size_stateful_group_uses_manager_mask(): + class StatefulManager: + @staticmethod + def find_longest_cache_hit(**kwargs): + pool = kwargs["block_pool"] + assert pool.hash_block_size == 16 + return ([pool.null_block, pool._present_block],) + + coordinator = object.__new__(DfkvStoreCoordinator) + coordinator.lcm_block_size = 16 + coordinator.hash_block_size = 16 + coordinator.eagle_attn_group_indices = set() + stateful_spec = SimpleNamespace(block_size=16) + coordinator.kv_cache_groups = [ + SimpleNamespace(kv_cache_spec=stateful_spec), + ] + coordinator.attention_groups = [(stateful_spec, [0], StatefulManager)] + + masks = coordinator.store_mask(32) + + assert masks == ([False, True],) diff --git a/integration/vllm/tests/test_hybrid_pool_layout.py b/integration/vllm/tests/test_hybrid_pool_layout.py index 61b2cc7f..4dfe8aa3 100644 --- a/integration/vllm/tests/test_hybrid_pool_layout.py +++ b/integration/vllm/tests/test_hybrid_pool_layout.py @@ -43,6 +43,17 @@ def _db() -> ChunkedTokenDatabase: return ChunkedTokenDatabase(_METADATA, block_size=16) +def test_logical_block_ids_expand_compact_stateful_table(): + assert worker_module._logical_block_ids([False, False, True], [7]) == [ + -1, + -1, + 7, + ] + assert worker_module._logical_block_ids([True, True], [5, 11]) == [5, 11] + with pytest.raises(ValueError, match="1 ids for 2 selected state slots"): + worker_module._logical_block_ids([True, True], [5]) + + def test_pool_key_uses_cross_runtime_binary_schema(): assert PoolKey(_METADATA, "0123abcd").to_bytes() == ( b"DFKVPOOL\x02" diff --git a/integration/vllm/tests/test_scheduler_full_block_ids.py b/integration/vllm/tests/test_scheduler_full_block_ids.py new file mode 100644 index 00000000..ada27edb --- /dev/null +++ b/integration/vllm/tests/test_scheduler_full_block_ids.py @@ -0,0 +1,43 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + +from dfkv_vllm.scheduler import DfkvStoreScheduler + + +def test_new_request_metadata_uses_complete_allocated_block_table(): + scheduler = object.__new__(DfkvStoreScheduler) + scheduler.kv_role = "kv_both" + scheduler.client = MagicMock() + scheduler.load_specs = {} + scheduler._request_trackers = {} + scheduler._preempted_req_ids = set() + scheduler._unfinished_request_ids = set() + scheduler._block_size = 4 + + request_real = SimpleNamespace(request_id="req-1", block_hashes=[]) + complete_block_ids = ([10, 11], [20, 21]) + scheduler._unfinished_requests = {} + blocks = SimpleNamespace(get_block_ids=lambda: complete_block_ids) + scheduler.update_state_after_alloc(request_real, blocks, num_external_tokens=0) + + scheduled_new_request = SimpleNamespace( + req_id="req-1", + num_computed_tokens=0, + # vLLM's NewRequestData exposes only blocks allocated this step. + block_ids=([11], [21]), + prefill_token_ids=None, + prompt_token_ids=list(range(8)), + ) + scheduler_output = SimpleNamespace( + finished_req_ids=set(), + preempted_req_ids=set(), + scheduled_new_reqs=[scheduled_new_request], + scheduled_cached_reqs=SimpleNamespace(req_ids=[]), + num_scheduled_tokens={"req-1": 8}, + ) + + metadata = scheduler.build_connector_meta(scheduler_output) + + assert len(metadata.requests) == 1 + assert metadata.requests[0].block_ids == complete_block_ids + assert metadata.requests[0].can_save is True