[Bugfix][KV Connector] Disable Mooncake TP put-striding when DCP > 1 (#45371)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Jingyi Yang <girasoleyang@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Yifan Qiao committed
Jun 18, 2026 at 15:13 UTC
41dcf49ca52ab25178ca8869298275b1787f328a
2 files changed
+87
-8
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+80
-7
@@ -175,14 +175,17 @@ class _FakeModelConfig:
175
176
177
def _make_vllm_config(
178
- *, extra_config: dict[str, object] | None = None
178
+ *,
179
+ extra_config: dict[str, object] | None = None,
180
+ rank: int = 0,
181
+ decode_context_parallel_size: int = 1,
182
) -> SimpleNamespace:
183
return SimpleNamespace(
184
model_config=_FakeModelConfig(),
185
parallel_config=SimpleNamespace(
186
pipeline_parallel_size=1,
184
- rank=0,
185
- decode_context_parallel_size=1,
187
+ rank=rank,
188
+ decode_context_parallel_size=decode_context_parallel_size,
189
prefill_context_parallel_size=1,
190
),
191
kv_transfer_config=_FakeKVTransferConfig(extra_config=extra_config),
@@ -231,13 +234,23 @@ def _install_fake_mooncake(monkeypatch, store_instance: MagicMock):
234
return FakeReplicateConfig
235
236
234
-def _patch_worker_runtime(monkeypatch, *, local_ip: str = "10.0.0.7") -> None:
237
+def _patch_worker_runtime(
238
+ monkeypatch,
239
+ *,
240
+ local_ip: str = "10.0.0.7",
241
+ tp_rank: int = 0,
242
+ tp_size: int = 1,
243
+ dcp_size: int = 1,
244
+) -> None:
245
single_rank_group = SimpleNamespace(world_size=1, rank_in_group=0)
246
+ # DCP groups are contiguous splits of the TP group (see
247
+ # parallel_state.py), so dcp_rank == tp_rank % dcp_size.
248
+ dcp_group = SimpleNamespace(world_size=dcp_size, rank_in_group=tp_rank % dcp_size)
249
monkeypatch.setattr(worker, "get_mooncake_dp_engine_index", lambda _: 0)
237
- monkeypatch.setattr(worker, "get_tensor_model_parallel_rank", lambda: 0)
238
- monkeypatch.setattr(worker, "get_tensor_model_parallel_world_size", lambda: 1)
250
+ monkeypatch.setattr(worker, "get_tensor_model_parallel_rank", lambda: tp_rank)
251
+ monkeypatch.setattr(worker, "get_tensor_model_parallel_world_size", lambda: tp_size)
252
monkeypatch.setattr(worker, "get_pcp_group", lambda: single_rank_group)
240
- monkeypatch.setattr(worker, "get_dcp_group", lambda: single_rank_group)
253
+ monkeypatch.setattr(worker, "get_dcp_group", lambda: dcp_group)
254
monkeypatch.setattr(worker, "get_ip", lambda: local_ip)
255
256
@@ -884,6 +897,66 @@ def test_requester_worker_init_builds_replicate_config_for_preferred_segment(
897
assert w.store_replicate_config.preferred_segment == "10.0.0.7:50053"
898
899
900
+@pytest.mark.parametrize("dcp_size", [1, 4])
901
+def test_worker_put_striding_covers_every_rank_get_namespace(
902
+ tmp_path, monkeypatch, dcp_size
903
+):
904
+ """Every key a rank GETs must have been PUT by some rank.
905
+
906
+ When num_kv_head < tp_size, ranks holding the same KV heads stripe
907
+ their PUTs across one shared key namespace. That dedup is only valid
908
+ when those ranks really share a namespace: with DCP > 1 each rank GETs
909
+ every key from its own ``@dcpN`` namespace, so striding must be
910
+ disabled.
911
+ """
912
+ tp_size = 4
913
+ store = MagicMock()
914
+ store.setup.return_value = 0
915
+ _install_fake_mooncake(monkeypatch, store)
916
+ monkeypatch.setenv(
917
+ "MOONCAKE_CONFIG_PATH",
918
+ _write_mooncake_config(
919
+ tmp_path,
920
+ {
921
+ "metadata_server": "http://metadata/endpoint",
922
+ "protocol": "tcp",
923
+ "device_name": "",
924
+ "master_server_address": "10.0.0.7:50051",
925
+ },
926
+ ),
927
+ )
928
+
929
+ # _FakeModelConfig has num_kv_head=1 < tp_size, which enables striding.
930
+ block_hashes = [f"hash-{i}".encode() for i in range(4)]
931
+ put_keys: set[str] = set()
932
+ get_keys_per_rank: dict[int, set[str]] = {}
933
+ for tp_rank in range(tp_size):
934
+ _patch_worker_runtime(
935
+ monkeypatch, tp_rank=tp_rank, tp_size=tp_size, dcp_size=dcp_size
936
+ )
937
+ w = worker.MooncakeStoreWorker(
938
+ _make_vllm_config(rank=tp_rank, decode_context_parallel_size=dcp_size),
939
+ _make_kv_cache_config(),
940
+ )
941
+ db = w.token_dbs[0]
942
+ token_len = len(block_hashes) * db.block_size
943
+ keys = [
944
+ key.to_string() for _, _, key in db.process_tokens(token_len, block_hashes)
945
+ ]
946
+ assert len(keys) == len(block_hashes)
947
+ # PUT side: mirrors KVCacheStoreSendingThread's striding slice.
948
+ put_keys.update(keys[w.tp_rank % w.put_step :: w.put_step])
949
+ # GET side: KVCacheStoreRecvingThread fetches every key.
950
+ get_keys_per_rank[tp_rank] = set(keys)
951
+
952
+ for tp_rank, rank_keys in get_keys_per_rank.items():
953
+ missing = rank_keys - put_keys
954
+ assert not missing, (
955
+ f"tp_rank={tp_rank} would GET {len(missing)}/{len(rank_keys)} keys "
956
+ f"that no rank PUT (Mooncake OBJECT_NOT_FOUND): {sorted(missing)}"
957
+ )
958
+
959
+
960
# ---------------------------------------------------------------------------
961
# Helpers for register_kv_caches tests
962
# ---------------------------------------------------------------------------
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+7
-1
@@ -972,7 +972,13 @@ class MooncakeStoreWorker:
972
else:
973
self.num_kv_head = model_config.get_total_num_kv_heads()
974
975
- if self.num_kv_head < self.tp_size:
975
+ if self.num_kv_head < self.tp_size and self.dcp_size <= 1:
976
+ # Dedup: TP ranks holding the same KV heads stripe PUTs across
977
+ # one shared key namespace. DCP splits the TP group, so with
978
+ # DCP>1 those ranks have different `@dcpN` namespaces and
979
+ # striping would leave keys unwritten (OBJECT_NOT_FOUND on
980
+ # GET). PCP is outer to TP (pcp_rank is constant within a TP
981
+ # group), so it needs no guard.
982
self.put_step = self.tp_size // self.num_kv_head
983
self.head_or_tp_rank = self.tp_rank // self.put_step
984
else: