(feat): optionally disable lookup on PD decode (#50498)
Signed-off-by: Jeff Ma <jeffjma@umich.edu> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
Jeff (Junze) Ma committed
Aug 1, 2026 at 02:41 UTC
3986b967b249468fc1bd052f89a1eeeada6377b4
5 files changed
+67
-3
tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py
+25
@@ -19,6 +19,7 @@ def _make_bare_scheduler(
19
scheduler = object.__new__(MooncakeStoreScheduler)
20
scheduler.kv_role = "kv_both"
21
scheduler.lookup_async = False
22
+ scheduler.enable_lookup = True
23
scheduler._block_size = 16
24
scheduler._hash_block_size = hash_block_size
25
scheduler.enable_partial_hash_hits = enable_partial_hash_hits
@@ -636,6 +637,30 @@ def test_sub_block_prompt_not_looked_up_without_fine_grained():
637
assert "req-0" not in scheduler.load_specs
638
639
640
+def test_disabled_lookup_reports_no_hit_without_querying_client():
641
+ # With enable_lookup=False the connector reports no external hit without
642
+ # consulting the lookup client, so admission is never deferred on a store
643
+ # lookup. Used by instances that only contribute store capacity.
644
+ scheduler = _make_bare_scheduler()
645
+ scheduler.enable_lookup = False
646
+ scheduler.client = _StubLookupClient(hit_tokens=32)
647
+
648
+ request = SimpleNamespace(
649
+ request_id="req-0",
650
+ num_tokens=48,
651
+ block_hashes=[b"h0", b"h1", b"h2"],
652
+ )
653
+
654
+ need_to_allocate, load_async = scheduler.get_num_new_matched_tokens(
655
+ request, num_computed_tokens=0
656
+ )
657
+
658
+ assert need_to_allocate == 0
659
+ assert load_async is False
660
+ assert scheduler.client.num_tokens == []
661
+ assert scheduler.load_specs == {}
662
+
663
+
664
def test_pending_partial_tail_emits_offload_only_reqmeta():
665
# A sub-block prompt never produces a block-aligned save, so the partial-
666
# tail offload arriving this step is emitted as an offload-only ReqMeta
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+1
@@ -1746,6 +1746,7 @@ def _make_bare_worker(
1746
worker.store = MagicMock()
1747
worker.store.register_buffer.return_value = 0
1748
worker.kv_role = kv_role
1749
+ worker._capacity_only = False
1750
worker.block_size = block_size
1751
worker.tp_rank = 0
1752
worker.enable_kv_events = False
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py
+10
-2
@@ -138,9 +138,17 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
138
)
139
assert vllm_config.kv_transfer_config is not None
140
assert kv_cache_config is not None, "kv_cache_config is required"
141
- self._validate_kv_cache_config(vllm_config, kv_cache_config)
142
- self._kv_cache_config = kv_cache_config
141
self.kv_role = vllm_config.kv_transfer_config.kv_role
142
+ # Capacity-only: contributes its segment to the store pool but transfers
143
+ # no KV, so the KV-cache-shape invariants below cannot be reached.
144
+ self._capacity_only = self.kv_role == "kv_consumer" and not (
145
+ vllm_config.kv_transfer_config.kv_connector_extra_config.get(
146
+ "enable_lookup", True
147
+ )
148
+ )
149
+ if not self._capacity_only:
150
+ self._validate_kv_cache_config(vllm_config, kv_cache_config)
151
+ self._kv_cache_config = kv_cache_config
152
self._kv_cache_events: MooncakeStoreKVEvents | None = None
153
154
self.connector_scheduler: MooncakeStoreScheduler | None = None
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py
+5
@@ -60,6 +60,8 @@ class MooncakeStoreScheduler:
60
kvc_extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config
61
self.load_async = kvc_extra_config.get("load_async", True)
62
self.lookup_async = kvc_extra_config.get("lookup_async", False)
63
+ # Skips lookup CPU cost on instances that never load KV from the store.
64
+ self.enable_lookup = kvc_extra_config.get("enable_lookup", True)
65
self.client = LookupKeyClient(vllm_config)
66
67
# Align with the engine's own scheduler_block_size and hash_block_size.
@@ -86,6 +88,9 @@ class MooncakeStoreScheduler:
88
Returns ``(None, False)`` when an async lookup is still in flight,
89
signaling the scheduler to retry this request on a later step.
90
"""
91
+ if not self.enable_lookup:
92
+ return 0, False
93
+
94
# Fine-grained hits may land on a hash boundary inside a block; without
95
# partial hits, prefixes shorter than one physical block are skipped.
96
align = (
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+26
-1
@@ -1172,6 +1172,12 @@ class MooncakeStoreWorker:
1172
self.load_async = vllm_config.kv_transfer_config.kv_connector_extra_config.get(
1173
"load_async", True
1174
)
1175
+ # Mirrors MooncakeStoreConnector._capacity_only.
1176
+ self._capacity_only = self.kv_role == "kv_consumer" and not (
1177
+ vllm_config.kv_transfer_config.kv_connector_extra_config.get(
1178
+ "enable_lookup", True
1179
+ )
1180
+ )
1181
self.cache_config = vllm_config.cache_config
1182
self.block_size, self.hash_block_size = resolve_kv_cache_block_sizes(
1183
kv_cache_config, vllm_config
@@ -1266,6 +1272,17 @@ class MooncakeStoreWorker:
1272
self.kv_connector_stats = MooncakeStoreConnectorStats()
1273
1274
self._kv_cache_config = kv_cache_config
1275
+ self.token_dbs: list[ChunkedTokenDatabase] = []
1276
+
1277
+ # a capacity-only instance does not need below utils
1278
+ if self._capacity_only:
1279
+ logger.info(
1280
+ "Mooncake store in capacity-only mode: segment mounted "
1281
+ "(global_segment_size=%d), KV transfer disabled.",
1282
+ store_config.global_segment_size,
1283
+ )
1284
+ return
1285
+
1286
# Single-group + PCP/DCP > 1: scale the lone group's spec.block_size to
1287
# self.block_size (= scheduler_block_size) so the coordinator's
1288
# ``block_size % hash_block_size == 0`` invariant holds.
@@ -1314,7 +1331,7 @@ class MooncakeStoreWorker:
1331
self._group_tp_replication_factors: tuple[int, ...] = (
1332
self._compute_group_tp_replication_factors()
1333
)
1317
- self.token_dbs: list[ChunkedTokenDatabase] = [
1334
+ self.token_dbs = [
1335
ChunkedTokenDatabase(
1336
dataclasses.replace(
1337
metadata,
@@ -1405,6 +1422,8 @@ class MooncakeStoreWorker:
1422
kv_caches: dict[str, torch.Tensor | list[torch.Tensor]],
1423
) -> None:
1424
"""Register KV cache tensors and start transfer threads."""
1425
+ if self._capacity_only:
1426
+ return
1427
if not kv_caches:
1428
logger.warning("No KV caches to offload.")
1429
return
@@ -1541,6 +1560,9 @@ class MooncakeStoreWorker:
1560
compute is launched on the compute stream) for better
1561
compute-I/O overlap.
1562
"""
1563
+ if self._capacity_only:
1564
+ return set(), set()
1565
+
1566
# Issue async loads
1567
for request in meta.requests:
1568
load_spec = request.load_spec
@@ -1659,6 +1681,9 @@ class MooncakeStoreWorker:
1681
hit covering all ``num_tokens`` is re-derived below the request end so
1682
the last token is recomputed for sampling.
1683
"""
1684
+ if self._capacity_only:
1685
+ return 0
1686
+
1687
token_len = self.coord.align_lookup_length(num_tokens)
1688
if not block_hashes or token_len <= 0:
1689
return 0