A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 07:15 UTC
[P/D] NixlConnector DP fixes (#18903)
Signed-off-by: Will Eaton <weaton@redhat.com>
Will Eaton committed
May 29, 2025 at 14:08 UTC
64eaf5fe0589a9e9082ce1de6a2e9d22ccb84b6a
3 files changed
+18
-6
vllm/distributed/kv_transfer/kv_connector/factory.py
+2
-1
@@ -70,7 +70,8 @@ class KVConnectorFactory:
70
connector_module = importlib.import_module(connector_module_path)
71
connector_cls = getattr(connector_module, connector_name)
72
assert issubclass(connector_cls, KVConnectorBase_V1)
73
- logger.info("Creating v1 connector with name: %s", connector_name)
73
+ logger.info("Creating v1 connector with name: %s and engine_id: %s",
74
+ connector_name, kv_transfer_config.engine_id)
75
# NOTE(Kuntai): v1 connector is explicitly separated into two roles.
76
# Scheduler connector:
77
# - Co-locate with scheduler process
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+7
-5
@@ -19,7 +19,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
19
KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole)
20
from vllm.distributed.parallel_state import (
21
get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size,
22
- get_tp_group)
22
+ get_tp_group, get_world_group)
23
from vllm.logger import init_logger
24
from vllm.utils import make_zmq_path, make_zmq_socket, round_down
25
from vllm.v1.core.sched.output import SchedulerOutput
@@ -334,6 +334,7 @@ class NixlConnectorWorker:
334
self.engine_id = engine_id
335
self.rank = get_tensor_model_parallel_rank()
336
self.world_size = get_tensor_model_parallel_world_size()
337
+ self.world_rank = get_world_group().rank_in_group
338
self.tp_group = get_tp_group()
339
340
# KV Caches and nixl tracking data.
@@ -382,7 +383,8 @@ class NixlConnectorWorker:
383
384
@staticmethod
385
def _nixl_handshake_listener(metadata: NixlAgentMetadata,
385
- ready_event: threading.Event, rank: int):
386
+ ready_event: threading.Event,
387
+ world_rank: int):
388
"""Background thread for getting new NIXL handshakes."""
389
# NOTE(rob): this is a simple implementation. We will move
390
# to a better approach like an ETCD server in the future.
@@ -403,7 +405,7 @@ class NixlConnectorWorker:
405
# NOTE(rob): we need each rank to have a unique port. This
406
# hack to keeps us moving. We will switch when moving to etcd
407
# or where we have a single ZMQ socket in the scheduler.
406
- port = envs.VLLM_NIXL_SIDE_CHANNEL_PORT + rank
408
+ port = envs.VLLM_NIXL_SIDE_CHANNEL_PORT + world_rank
409
path = make_zmq_path("tcp", host, port)
410
logger.debug("Starting listening on path: %s", path)
411
with zmq_ctx(zmq.ROUTER, path) as sock:
@@ -422,7 +424,7 @@ class NixlConnectorWorker:
424
# NOTE(rob): we need each rank to have a unique port. This is
425
# a hack to keep us moving. We will switch when moving to etcd
426
# or where we have a single ZMQ socket in the scheduler.
425
- path = make_zmq_path("tcp", host, port + self.rank)
427
+ path = make_zmq_path("tcp", host, port + self.world_rank)
428
logger.debug("Querying metadata on path: %s", path)
429
with zmq_ctx(zmq.REQ, path) as sock:
430
# Send query for the request.
@@ -529,7 +531,7 @@ class NixlConnectorWorker:
531
ready_event = threading.Event()
532
self._nixl_handshake_listener_t = threading.Thread(
533
target=self._nixl_handshake_listener,
532
- args=(metadata, ready_event, self.rank),
534
+ args=(metadata, ready_event, self.world_rank),
535
daemon=True,
536
name="nixl_handshake_listener")
537
self._nixl_handshake_listener_t.start()
vllm/v1/engine/core.py
+9
@@ -707,6 +707,15 @@ class DPEngineCoreProc(EngineCoreProc):
707
assert dp_size > 1
708
assert 0 <= local_dp_rank <= dp_rank < dp_size
709
710
+ if vllm_config.kv_transfer_config is not None:
711
+ # modify the engine_id and append the local_dp_rank to it to ensure
712
+ # that the kv_transfer_config is unique for each DP rank.
713
+ vllm_config.kv_transfer_config.engine_id = (
714
+ f"{vllm_config.kv_transfer_config.engine_id}_dp{local_dp_rank}"
715
+ )
716
+ logger.debug("Setting kv_transfer_config.engine_id to %s",
717
+ vllm_config.kv_transfer_config.engine_id)
718
+
719
from vllm.platforms import current_platform
720
device_control_env_var = current_platform.device_control_env_var
721
world_size = vllm_config.parallel_config.world_size