[BugFix][Mooncake] Use global data_parallel_index for the DP engine index (#48061)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Yifan Qiao committed
Aug 4, 2026 at 17:10 UTC
ffee32460d7acdcfe774a468bd7e5cdaac40b562
3 files changed
+3
-16
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+1
-1
@@ -209,6 +209,7 @@ def _make_vllm_config(
209
parallel_config=SimpleNamespace(
210
pipeline_parallel_size=1,
211
rank=rank,
212
+ data_parallel_index=0,
213
decode_context_parallel_size=decode_context_parallel_size,
214
prefill_context_parallel_size=1,
215
),
@@ -270,7 +271,6 @@ def _patch_worker_runtime(
271
# DCP groups are contiguous splits of the TP group (see
272
# parallel_state.py), so dcp_rank == tp_rank % dcp_size.
273
dcp_group = SimpleNamespace(world_size=dcp_size, rank_in_group=tp_rank % dcp_size)
273
- monkeypatch.setattr(worker, "get_mooncake_dp_engine_index", lambda _: 0)
274
monkeypatch.setattr(worker, "get_tensor_model_parallel_rank", lambda: tp_rank)
275
monkeypatch.setattr(worker, "get_tensor_model_parallel_world_size", lambda: tp_size)
276
monkeypatch.setattr(worker, "get_pcp_group", lambda: single_rank_group)
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py
-10
@@ -8,7 +8,6 @@ import uvicorn
8
from fastapi import FastAPI, HTTPException
9
from pydantic import BaseModel
10
11
-from vllm.config import ParallelConfig
11
from vllm.distributed.kv_transfer.kv_connector.utils import EngineId
12
from vllm.logger import init_logger
13
@@ -17,15 +16,6 @@ WorkerAddr = str
16
logger = init_logger(__name__)
17
18
20
-def get_mooncake_dp_engine_index(parallel_config: ParallelConfig) -> int:
21
- """Return the per-engine DP index used for Mooncake side channels."""
22
- if parallel_config.local_engines_only:
23
- assert parallel_config.data_parallel_rank_local is not None
24
- return parallel_config.data_parallel_rank_local
25
-
26
- return parallel_config.data_parallel_index
27
-
28
-
19
class RegisterWorkerPayload(BaseModel):
20
engine_id: EngineId
21
dp_rank: int
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+2
-5
@@ -37,9 +37,6 @@ from vllm.distributed import (
37
)
38
from vllm.distributed.kv_events import BlockStored
39
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import rdma_utils
40
-from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import (
41
- get_mooncake_dp_engine_index,
42
-)
40
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( # noqa: E501
41
ExternalCachedBlockPool,
42
MooncakeStoreCoordinator,
@@ -1156,7 +1153,7 @@ class MooncakeStoreWorker:
1153
model_config = vllm_config.model_config
1154
parallel_config = vllm_config.parallel_config
1155
1159
- self.dp_rank = get_mooncake_dp_engine_index(parallel_config)
1156
+ self.dp_rank = parallel_config.data_parallel_index
1157
self.tp_rank = get_tensor_model_parallel_rank()
1158
self.tp_size = get_tensor_model_parallel_world_size()
1159
self.pp_size = parallel_config.pipeline_parallel_size
@@ -1973,7 +1970,7 @@ class LookupKeyClient:
1970
def get_zmq_rpc_path_lookup(vllm_config: VllmConfig) -> str:
1971
"""Construct IPC path for ZMQ lookup socket."""
1972
assert vllm_config.kv_transfer_config is not None
1976
- dp_rank = get_mooncake_dp_engine_index(vllm_config.parallel_config)
1973
+ dp_rank = vllm_config.parallel_config.data_parallel_index
1974
base_url = envs.VLLM_RPC_BASE_PATH
1975
rpc_port = 0
1976
hostname = socket.gethostname()