fix(vllm): preserve hybrid block tables for external KV

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

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(

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",

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],)

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"

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