fix(kv_pool): cover physical DCP shards in lookup

Signed-off-by: zmc1997 <40617288+zmc1997@users.noreply.github.com>

Author: zmc1997

tests/ut/distributed/ascend_store/test_metadata.py

@@ -111,12 +111,11 @@ def test_hash_diff(self):
     def test_to_string(self):
         k = PoolKey(self.meta, "hash1")
         s = k.to_string()
-        self.assertIn("llama", s)
-        self.assertIn("@pcp2", s)
-        self.assertIn("@dcp3", s)
-        self.assertIn("@head_or_tp_rank:1", s)
-        self.assertIn("@pp_rank:0", s)
-        self.assertIn("hash1", s)
+        self.assertEqual(
+            s,
+            "llama@pcp:2@dcp:3@head_or_tp_rank:1@pp_rank:0"
+            "@group:0@cache_role:kv@cache_family:default@hash1",
+        )
 
     def test_pp_ranks_use_distinct_keys(self):
         other_pp_meta = KeyMetadata("llama", 1, 2, 3, 1)
@@ -148,6 +147,7 @@ def test_to_string_contains_layer_id(self):
         meta = KeyMetadata("model", 0, 0, 0, 0)
         k = LayerPoolKey(meta, "h1", 5)
         s = k.to_string()
+        self.assertIn("@pcp:0@dcp:0", s)
         self.assertIn("@layer_id:5", s)
         self.assertIn("model", s)
         self.assertTrue(s.endswith("@h1"))

vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/metadata.py

@@ -115,7 +115,7 @@ def __hash__(self):
     def to_string(self):
         return (
             f"{self.key_metadata.model_name}"
-            f"@pcp{self.key_metadata.pcp_rank}@dcp{self.key_metadata.dcp_rank}"
+            f"@pcp:{self.key_metadata.pcp_rank}@dcp:{self.key_metadata.dcp_rank}"
             f"@head_or_tp_rank:{self.key_metadata.head_or_tp_rank}"
             f"@pp_rank:{self.key_metadata.pp_rank}"
             f"@group:{self.key_metadata.kv_cache_group_id}"
@@ -162,7 +162,7 @@ def __hash__(self):
     def to_string(self):
         return (
             f"{self.key_metadata.model_name}"
-            f"@pcp{self.key_metadata.pcp_rank}@dcp{self.key_metadata.dcp_rank}"
+            f"@pcp:{self.key_metadata.pcp_rank}@dcp:{self.key_metadata.dcp_rank}"
             f"@head_or_tp_rank:{self.key_metadata.head_or_tp_rank}"
             f"@group:{self.key_metadata.kv_cache_group_id}"
             f"@cache_role:{self.key_metadata.cache_role}"
@@ -327,7 +327,7 @@ def _get_key_prefix(
             group_metadata = self.metadata[kv_cache_group_id]
             prefix = (
                 f"{group_metadata.model_name}"
-                f"@pcp{group_metadata.pcp_rank}@dcp{group_metadata.dcp_rank}"
+                f"@pcp:{group_metadata.pcp_rank}@dcp:{group_metadata.dcp_rank}"
                 f"@head_or_tp_rank:{group_metadata.head_or_tp_rank}"
                 f"@pp_rank:{group_metadata.pp_rank}"
                 f"@group:{kv_cache_group_id}"

vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/pool_worker.py

@@ -2211,12 +2211,18 @@ def _replace_key_field(key: str, field: str, value: int) -> str:
         return f"{key[:value_start]}{value}{key[value_end:]}"
 
     def _expand_lookup_keys_by_rank(self, keys: list[str], group_id: int) -> list[str]:
+        # All-rank KV pool lookup currently assumes PCP=1.
         expanded: list[str] = []
+        num_head_or_tp_ranks = self.get_group_tp_size(group_id)
+        # Keep each rank shard's block/layer keys contiguous to match
+        # lookup_scheduler()'s [rank_shard][block] result slicing.
         for pp_rank in range(self.pp_size):
-            for tp_rank in range(self.get_group_tp_size(group_id)):
-                for key in keys:
-                    tp_key = self._replace_key_field(key, "head_or_tp_rank", tp_rank)
-                    expanded.append(self._replace_key_field(tp_key, "pp_rank", pp_rank))
+            for dcp_rank in range(self.dcp_size):
+                for head_or_tp_rank in range(num_head_or_tp_ranks):
+                    for key in keys:
+                        rank_key = self._replace_key_field(key, "dcp", dcp_rank)
+                        rank_key = self._replace_key_field(rank_key, "head_or_tp_rank", head_or_tp_rank)
+                        expanded.append(self._replace_key_field(rank_key, "pp_rank", pp_rank))
         return expanded
 
     def _expand_lookup_key_variants(self, key: str, group_id: int, include_all_ranks: bool) -> list[str]: