From 95528527eab9077aa4eb1d21ccdc20ef18eb5c95 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Sun, 28 Jun 2026 17:36:23 -0400 Subject: [PATCH] [Bugfix][Mooncake] Fix Mooncake lookup prefixes with DCP > 1 (#46855) Signed-off-by: wzhao18 --- .../unit/test_mooncake_store_worker.py | 48 +++++++++++++++++++ .../kv_connector/v1/mooncake/store/data.py | 6 ++- .../kv_connector/v1/mooncake/store/worker.py | 42 ++++++++++++---- 3 files changed, 86 insertions(+), 10 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index d6ce200f4c..852c75ae9c 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/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] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py index 7f8168ab38..ef98ec0d4e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py +++ b/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}" ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index aea2d602e7..e60d2f47a4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/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]] = []