[Mooncake] Optimize lookup pool key string construction (#46188)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Wei Zhao committed
Jun 23, 2026 at 23:51 UTC
05a0caba916d3944d2c86a76ca5596f176e25190
3 files changed
+57
-31
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+2
@@ -1279,6 +1279,7 @@ def _make_bare_worker(
1279
scheduler_block_size=block_size,
1280
hash_block_size=block_size,
1281
)
1282
+ worker._init_lookup_key_prefixes()
1283
return worker
1284
1285
@@ -1346,6 +1347,7 @@ def test_lookup_checks_all_potential_swa_hit_boundaries():
1347
hash_block_size=8,
1348
retention_interval=0,
1349
)
1350
+ worker._init_lookup_key_prefixes()
1351
# Candidate order: 3 full-attention chunks, then SWA chunks 3, 7, 11.
1352
# Only the first full chunk and the SWA chunk ending at token 32 exist, so
1353
# lookup should recover a 32-token external prefix hit. A sparse
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
+24
-13
@@ -132,21 +132,32 @@ class PoolKey:
132
)
133
)
134
135
- def to_string(self) -> str:
136
- prefix = (
137
- f"{self.key_metadata.cache_prefix}@"
138
- if self.key_metadata.cache_prefix
139
- else ""
140
- )
135
+ @staticmethod
136
+ def build_prefix(
137
+ key_metadata: KeyMetadata,
138
+ *,
139
+ tp_rank: int | None = None,
140
+ pp_rank: int | None = None,
141
+ ) -> str:
142
+ """Return the stable prefix for a Mooncake pool key."""
143
+ prefix = f"{key_metadata.cache_prefix}@" if key_metadata.cache_prefix else ""
144
return (
145
f"{prefix}"
143
- f"{self.key_metadata.model_name}"
144
- f"@tp_rank:{self.key_metadata.tp_rank}"
145
- f"@pcp{self.key_metadata.pcp_rank}"
146
- f"@dcp{self.key_metadata.dcp_rank}"
147
- f"@pp_rank:{self.key_metadata.pp_rank}"
148
- f"@group:{self.key_metadata.group_id}"
149
- f"@{self.chunk_hash}"
146
+ f"{key_metadata.model_name}"
147
+ 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"@pp_rank:{key_metadata.pp_rank if pp_rank is None else pp_rank}"
151
+ f"@group:{key_metadata.group_id}"
152
+ )
153
+
154
+ @staticmethod
155
+ def build_key_string(key_prefix: str, chunk_hash: str) -> str:
156
+ return f"{key_prefix}@{chunk_hash}"
157
+
158
+ def to_string(self) -> str:
159
+ return self.build_key_string(
160
+ self.build_prefix(self.key_metadata), self.chunk_hash
161
)
162
163
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+31
-18
@@ -1119,6 +1119,20 @@ class MooncakeStoreWorker:
1119
)
1120
for g_idx, g in enumerate(self._kv_cache_groups)
1121
]
1122
+ self._init_lookup_key_prefixes()
1123
+
1124
+ def _init_lookup_key_prefixes(self) -> None:
1125
+ """Precompute per-group key prefixes expanded across TP/PP ranks."""
1126
+ tp_count = min(self.tp_size, self.num_kv_head)
1127
+ self._lookup_key_prefixes = tuple(
1128
+ tuple(
1129
+ PoolKey.build_prefix(db.metadata, tp_rank=tp, pp_rank=pp)
1130
+ for tp in range(tp_count)
1131
+ for pp in range(self.pp_size)
1132
+ )
1133
+ for db in self.token_dbs
1134
+ )
1135
+ self._lookup_expected_per_key = tp_count * self.pp_size
1136
1137
def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None:
1138
"""Register a cross-layers KV cache tensor.
@@ -1381,22 +1395,17 @@ class MooncakeStoreWorker:
1395
return 0
1396
1397
# Build per-(group, hash) candidate keys expanded across TP/PP.
1384
- # candidate_meta[i] is the (group_id, hash_bytes) for candidate_keys[i].
1398
+ # candidate_meta stores the (group, hash_bytes) for key slice.
1399
candidate_keys: list[str] = []
1400
candidate_meta: list[tuple[int, bytes]] = []
1401
lookup_masks = self.coord.lookup_mask(token_len)
1388
- tp_count = min(self.tp_size, self.num_kv_head)
1402
for g_idx, db in enumerate(self.token_dbs):
1403
spec_block_size = db.block_size
1404
lookup_mask = lookup_masks[g_idx]
1405
+ key_prefixes = self._lookup_key_prefixes[g_idx]
1406
group_hashes = self.coord.block_hashes_for_spec(
1407
block_hashes, self._kv_cache_groups[g_idx].kv_cache_spec
1408
)
1395
- metadata_templates = [
1396
- dataclasses.replace(db.metadata, tp_rank=tp, pp_rank=pp)
1397
- for tp in range(tp_count)
1398
- for pp in range(self.pp_size)
1399
- ]
1409
for chunk_id, h in enumerate(group_hashes):
1410
start_idx = chunk_id * spec_block_size
1411
if start_idx >= token_len:
@@ -1405,11 +1414,12 @@ class MooncakeStoreWorker:
1414
chunk_id >= len(lookup_mask) or not lookup_mask[chunk_id]
1415
):
1416
continue
1408
- h_hex = h.hex()
1409
- h_bytes = bytes(h)
1410
- for md in metadata_templates:
1411
- candidate_keys.append(PoolKey(md, h_hex).to_string())
1412
- candidate_meta.append((g_idx, h_bytes))
1417
+ hash_hex = h.hex()
1418
+ for key_prefix in key_prefixes:
1419
+ candidate_keys.append(
1420
+ PoolKey.build_key_string(key_prefix, hash_hex)
1421
+ )
1422
+ candidate_meta.append((g_idx, bytes(h)))
1423
1424
if not candidate_keys:
1425
return 0
@@ -1434,12 +1444,15 @@ class MooncakeStoreWorker:
1444
return 0
1445
1446
# A (group, hash) is "present" only when every TP*PP rank has it.
1437
- expected_per_key = max(1, tp_count * self.pp_size)
1438
- present_count: dict[tuple[int, bytes], int] = {}
1439
- for gh, exists in zip(candidate_meta, res, strict=True):
1440
- if exists == 1:
1441
- present_count[gh] = present_count.get(gh, 0) + 1
1442
- exists_set = {gh for gh, c in present_count.items() if c >= expected_per_key}
1447
+ ranks_per_candidate = self._lookup_expected_per_key
1448
+ exists_set = {
1449
+ (g_idx, hash_bytes)
1450
+ for i, (g_idx, hash_bytes) in enumerate(candidate_meta)
1451
+ if all(
1452
+ res[i * ranks_per_candidate + j] == 1
1453
+ for j in range(ranks_per_candidate)
1454
+ )
1455
+ }
1456
1457
_masks, hit_length = self.coord.find_longest_cache_hit(
1458
block_hashes, token_len, ExternalCachedBlockPool(exists_set)