[Bugfix][Mooncake] Fix Mooncake lookup prefixes with DCP > 1 (#46855)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Wei Zhao committed
Jun 28, 2026 at 17:36 UTC
95528527eab9077aa4eb1d21ccdc20ef18eb5c95
3 files changed
+86
-10
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+48
@@ -1621,6 +1621,54 @@ def _make_bare_worker(
1621
return worker
1622
1623
1624
+def test_lookup_key_prefixes_cover_dcp_rank_namespaces():
1625
+ worker = _make_bare_worker()
1626
+ worker.tp_size = 4
1627
+ worker.num_kv_head = 1
1628
+ worker.dcp_size = 4
1629
+ worker._init_lookup_key_prefixes()
1630
+
1631
+ assert worker._lookup_expected_per_key == 4
1632
+ assert worker._lookup_key_prefixes[0] == (
1633
+ "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0",
1634
+ "test-model@tp_rank:1@pcp0@dcp1@pp_rank:0@group:0",
1635
+ "test-model@tp_rank:2@pcp0@dcp2@pp_rank:0@group:0",
1636
+ "test-model@tp_rank:3@pcp0@dcp3@pp_rank:0@group:0",
1637
+ )
1638
+
1639
+
1640
+def test_lookup_key_prefixes_cover_pcp_rank_namespaces():
1641
+ worker = _make_bare_worker()
1642
+ worker.tp_size = 4
1643
+ worker.num_kv_head = 1
1644
+ worker.pcp_size = 2
1645
+ worker.dcp_size = 1
1646
+ worker._init_lookup_key_prefixes()
1647
+
1648
+ assert worker._lookup_expected_per_key == 2
1649
+ assert worker._lookup_key_prefixes[0] == (
1650
+ "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0",
1651
+ "test-model@tp_rank:0@pcp1@dcp0@pp_rank:0@group:0",
1652
+ )
1653
+
1654
+
1655
+def test_lookup_requires_all_dcp_rank_namespaces():
1656
+ worker = _make_bare_worker(block_size=16)
1657
+ worker.tp_size = 4
1658
+ worker.num_kv_head = 1
1659
+ worker.dcp_size = 4
1660
+ worker._init_lookup_key_prefixes()
1661
+ worker.store.batch_is_exist.return_value = [1, 1, 0, 1]
1662
+
1663
+ assert worker.lookup(16, [b"a0"]) == 0
1664
+ assert worker.store.batch_is_exist.call_args.args[0] == [
1665
+ "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6130",
1666
+ "test-model@tp_rank:1@pcp0@dcp1@pp_rank:0@group:0@6130",
1667
+ "test-model@tp_rank:2@pcp0@dcp2@pp_rank:0@group:0@6130",
1668
+ "test-model@tp_rank:3@pcp0@dcp3@pp_rank:0@group:0@6130",
1669
+ ]
1670
+
1671
+
1672
def test_lookup_partial_prefix_returns_first_hit_length():
1673
worker = _make_bare_worker()
1674
worker.store.batch_is_exist.return_value = [1, 1, 0]
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
+4
-2
@@ -137,6 +137,8 @@ class PoolKey:
137
key_metadata: KeyMetadata,
138
*,
139
tp_rank: int | None = None,
140
+ pcp_rank: int | None = None,
141
+ dcp_rank: int | None = None,
142
pp_rank: int | None = None,
143
) -> str:
144
"""Return the stable prefix for a Mooncake pool key."""
@@ -145,8 +147,8 @@ class PoolKey:
147
f"{prefix}"
148
f"{key_metadata.model_name}"
149
f"@tp_rank:{key_metadata.tp_rank if tp_rank is None else tp_rank}"
148
- f"@pcp{key_metadata.pcp_rank}"
149
- f"@dcp{key_metadata.dcp_rank}"
150
+ f"@pcp{key_metadata.pcp_rank if pcp_rank is None else pcp_rank}"
151
+ f"@dcp{key_metadata.dcp_rank if dcp_rank is None else dcp_rank}"
152
f"@pp_rank:{key_metadata.pp_rank if pp_rank is None else pp_rank}"
153
f"@group:{key_metadata.group_id}"
154
)
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+34
-8
@@ -1154,17 +1154,43 @@ class MooncakeStoreWorker:
1154
self._init_lookup_key_prefixes()
1155
1156
def _init_lookup_key_prefixes(self) -> None:
1157
- """Precompute per-group key prefixes expanded across TP/PP ranks."""
1158
- tp_count = min(self.tp_size, self.num_kv_head)
1157
+ """Prepare per-group key prefixes across parallel rank namespaces."""
1158
+ # (tp_rank, pcp_rank, dcp_rank, pp_rank) namespaces
1159
+ if self.dcp_size > 1:
1160
+ # DCP reuses the TP workers and splits each TP group into
1161
+ # contiguous DCP groups, so dcp_rank == tp_rank % dcp_size.
1162
+ # Store/load paths do not apply KV-head dedup under DCP
1163
+ rank_namespaces = tuple(
1164
+ (tp_rank, pcp_rank, tp_rank % self.dcp_size, pp_rank)
1165
+ for pcp_rank in range(self.pcp_size)
1166
+ for tp_rank in range(self.tp_size)
1167
+ for pp_rank in range(self.pp_size)
1168
+ )
1169
+ else:
1170
+ # Without DCP, TP ranks that share a KV head write identical KV, so
1171
+ # lookup only needs one TP namespace per unique KV head.
1172
+ tp_count = min(self.tp_size, self.num_kv_head)
1173
+ rank_namespaces = tuple(
1174
+ (tp_rank, pcp_rank, 0, pp_rank)
1175
+ for pcp_rank in range(self.pcp_size)
1176
+ for tp_rank in range(tp_count)
1177
+ for pp_rank in range(self.pp_size)
1178
+ )
1179
+
1180
self._lookup_key_prefixes = tuple(
1181
tuple(
1161
- PoolKey.build_prefix(db.metadata, tp_rank=tp, pp_rank=pp)
1162
- for tp in range(tp_count)
1163
- for pp in range(self.pp_size)
1182
+ PoolKey.build_prefix(
1183
+ db.metadata,
1184
+ tp_rank=tp_rank,
1185
+ pcp_rank=pcp_rank,
1186
+ dcp_rank=dcp_rank,
1187
+ pp_rank=pp_rank,
1188
+ )
1189
+ for tp_rank, pcp_rank, dcp_rank, pp_rank in rank_namespaces
1190
)
1191
for db in self.token_dbs
1192
)
1167
- self._lookup_expected_per_key = tp_count * self.pp_size
1193
+ self._lookup_expected_per_key = len(rank_namespaces)
1194
1195
def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None:
1196
"""Register a cross-layers KV cache tensor.
@@ -1430,12 +1456,12 @@ class MooncakeStoreWorker:
1456
def lookup(self, token_len: int, block_hashes: Sequence[BlockHash]) -> int:
1457
"""Check how many prefix tokens exist in the store.
1458
1433
- Checks across all TP ranks and PP ranks.
1459
+ Checks across all rank-specific key namespaces that may be loaded.
1460
"""
1461
if not block_hashes or token_len <= 0:
1462
return 0
1463
1438
- # Build per-(group, hash) candidate keys expanded across TP/PP.
1464
+ # Build per-(group, hash) candidate keys expanded across rank namespaces.
1465
# candidate_meta stores the (group, hash_bytes) for key slice.
1466
candidate_keys: list[str] = []
1467
candidate_meta: list[tuple[int, bytes]] = []