[Bugfix][Mooncake] Fix Mooncake lookup prefixes with DCP > 1 (#46855)

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>

tests/v1/kv_connector/unit/test_mooncake_store_worker.py

@@ -1621,6 +1621,54 @@ def _make_bare_worker(
     return worker
 
 
+def test_lookup_key_prefixes_cover_dcp_rank_namespaces():
+    worker = _make_bare_worker()
+    worker.tp_size = 4
+    worker.num_kv_head = 1
+    worker.dcp_size = 4
+    worker._init_lookup_key_prefixes()
+
+    assert worker._lookup_expected_per_key == 4
+    assert worker._lookup_key_prefixes[0] == (
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0",
+        "test-model@tp_rank:1@pcp0@dcp1@pp_rank:0@group:0",
+        "test-model@tp_rank:2@pcp0@dcp2@pp_rank:0@group:0",
+        "test-model@tp_rank:3@pcp0@dcp3@pp_rank:0@group:0",
+    )
+
+
+def test_lookup_key_prefixes_cover_pcp_rank_namespaces():
+    worker = _make_bare_worker()
+    worker.tp_size = 4
+    worker.num_kv_head = 1
+    worker.pcp_size = 2
+    worker.dcp_size = 1
+    worker._init_lookup_key_prefixes()
+
+    assert worker._lookup_expected_per_key == 2
+    assert worker._lookup_key_prefixes[0] == (
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0",
+        "test-model@tp_rank:0@pcp1@dcp0@pp_rank:0@group:0",
+    )
+
+
+def test_lookup_requires_all_dcp_rank_namespaces():
+    worker = _make_bare_worker(block_size=16)
+    worker.tp_size = 4
+    worker.num_kv_head = 1
+    worker.dcp_size = 4
+    worker._init_lookup_key_prefixes()
+    worker.store.batch_is_exist.return_value = [1, 1, 0, 1]
+
+    assert worker.lookup(16, [b"a0"]) == 0
+    assert worker.store.batch_is_exist.call_args.args[0] == [
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6130",
+        "test-model@tp_rank:1@pcp0@dcp1@pp_rank:0@group:0@6130",
+        "test-model@tp_rank:2@pcp0@dcp2@pp_rank:0@group:0@6130",
+        "test-model@tp_rank:3@pcp0@dcp3@pp_rank:0@group:0@6130",
+    ]
+
+
 def test_lookup_partial_prefix_returns_first_hit_length():
     worker = _make_bare_worker()
     worker.store.batch_is_exist.return_value = [1, 1, 0]

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py

@@ -137,6 +137,8 @@ class PoolKey:
         key_metadata: KeyMetadata,
         *,
         tp_rank: int | None = None,
+        pcp_rank: int | None = None,
+        dcp_rank: int | None = None,
         pp_rank: int | None = None,
     ) -> str:
         """Return the stable prefix for a Mooncake pool key."""
@@ -145,8 +147,8 @@ class PoolKey:
             f"{prefix}"
             f"{key_metadata.model_name}"
             f"@tp_rank:{key_metadata.tp_rank if tp_rank is None else tp_rank}"
-            f"@pcp{key_metadata.pcp_rank}"
-            f"@dcp{key_metadata.dcp_rank}"
+            f"@pcp{key_metadata.pcp_rank if pcp_rank is None else pcp_rank}"
+            f"@dcp{key_metadata.dcp_rank if dcp_rank is None else dcp_rank}"
             f"@pp_rank:{key_metadata.pp_rank if pp_rank is None else pp_rank}"
             f"@group:{key_metadata.group_id}"
         )

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py

@@ -1154,17 +1154,43 @@ class MooncakeStoreWorker:
         self._init_lookup_key_prefixes()
 
     def _init_lookup_key_prefixes(self) -> None:
-        """Precompute per-group key prefixes expanded across TP/PP ranks."""
-        tp_count = min(self.tp_size, self.num_kv_head)
+        """Prepare per-group key prefixes across parallel rank namespaces."""
+        # (tp_rank, pcp_rank, dcp_rank, pp_rank) namespaces
+        if self.dcp_size > 1:
+            # DCP reuses the TP workers and splits each TP group into
+            # contiguous DCP groups, so dcp_rank == tp_rank % dcp_size.
+            # Store/load paths do not apply KV-head dedup under DCP
+            rank_namespaces = tuple(
+                (tp_rank, pcp_rank, tp_rank % self.dcp_size, pp_rank)
+                for pcp_rank in range(self.pcp_size)
+                for tp_rank in range(self.tp_size)
+                for pp_rank in range(self.pp_size)
+            )
+        else:
+            # Without DCP, TP ranks that share a KV head write identical KV, so
+            # lookup only needs one TP namespace per unique KV head.
+            tp_count = min(self.tp_size, self.num_kv_head)
+            rank_namespaces = tuple(
+                (tp_rank, pcp_rank, 0, pp_rank)
+                for pcp_rank in range(self.pcp_size)
+                for tp_rank in range(tp_count)
+                for pp_rank in range(self.pp_size)
+            )
+
         self._lookup_key_prefixes = tuple(
             tuple(
-                PoolKey.build_prefix(db.metadata, tp_rank=tp, pp_rank=pp)
-                for tp in range(tp_count)
-                for pp in range(self.pp_size)
+                PoolKey.build_prefix(
+                    db.metadata,
+                    tp_rank=tp_rank,
+                    pcp_rank=pcp_rank,
+                    dcp_rank=dcp_rank,
+                    pp_rank=pp_rank,
+                )
+                for tp_rank, pcp_rank, dcp_rank, pp_rank in rank_namespaces
             )
             for db in self.token_dbs
         )
-        self._lookup_expected_per_key = tp_count * self.pp_size
+        self._lookup_expected_per_key = len(rank_namespaces)
 
     def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None:
         """Register a cross-layers KV cache tensor.
@@ -1430,12 +1456,12 @@ class MooncakeStoreWorker:
     def lookup(self, token_len: int, block_hashes: Sequence[BlockHash]) -> int:
         """Check how many prefix tokens exist in the store.
 
-        Checks across all TP ranks and PP ranks.
+        Checks across all rank-specific key namespaces that may be loaded.
         """
         if not block_hashes or token_len <= 0:
             return 0
 
-        # Build per-(group, hash) candidate keys expanded across TP/PP.
+        # Build per-(group, hash) candidate keys expanded across rank namespaces.
         # candidate_meta stores the (group, hash_bytes) for key slice.
         candidate_keys: list[str] = []
         candidate_meta: list[tuple[int, bytes]] = []