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
[Bugfix][Nixl] Fix DP Metadata Handshake (#19008)
Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>
Robert Shaw committed
Jun 1, 2025 at 23:30 UTC
b9f61e13875e1682d3982829006bec26981fde4d
1 file changed
+36
-32
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+36
-32
@@ -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, get_world_group)
22
+ get_tp_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
@@ -172,6 +172,11 @@ class NixlConnectorScheduler:
172
self.vllm_config = vllm_config
173
self.block_size = vllm_config.cache_config.block_size
174
self.engine_id = engine_id
175
+ self.side_channel_host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST
176
+ self.side_channel_port = (
177
+ envs.VLLM_NIXL_SIDE_CHANNEL_PORT +
178
+ vllm_config.parallel_config.data_parallel_rank_local *
179
+ vllm_config.parallel_config.tensor_parallel_size)
180
logger.info("Initializing NIXL Scheduler %s", engine_id)
181
182
# Requests that need to start recv.
@@ -310,8 +315,8 @@ class NixlConnectorScheduler:
315
do_remote_decode=False,
316
remote_block_ids=computed_block_ids,
317
remote_engine_id=self.engine_id,
313
- remote_host=envs.VLLM_NIXL_SIDE_CHANNEL_HOST,
314
- remote_port=envs.VLLM_NIXL_SIDE_CHANNEL_PORT,
318
+ remote_host=self.side_channel_host,
319
+ remote_port=self.side_channel_port,
320
)
321
322
@@ -330,11 +335,19 @@ class NixlConnectorWorker:
335
# Map of engine_id -> agent_name.
336
self._remote_agents: dict[str, str] = {}
337
338
+ # NIXL handshake port.
339
+ # NOTE(rob): Within a DP group, each DP rank gets its own
340
+ # base port (which is sent in the KVTransferParams).
341
+ # Each TP rank listens/queries on the base_port + tp_rank.
342
+ self.side_channel_port = (
343
+ envs.VLLM_NIXL_SIDE_CHANNEL_PORT +
344
+ vllm_config.parallel_config.data_parallel_rank_local *
345
+ vllm_config.parallel_config.tensor_parallel_size)
346
+
347
# Metadata.
348
self.engine_id = engine_id
335
- self.rank = get_tensor_model_parallel_rank()
349
+ self.tp_rank = get_tensor_model_parallel_rank()
350
self.world_size = get_tensor_model_parallel_world_size()
337
- self.world_rank = get_world_group().rank_in_group
351
self.tp_group = get_tp_group()
352
353
# KV Caches and nixl tracking data.
@@ -383,16 +396,11 @@ class NixlConnectorWorker:
396
397
@staticmethod
398
def _nixl_handshake_listener(metadata: NixlAgentMetadata,
386
- ready_event: threading.Event,
387
- world_rank: int):
399
+ ready_event: threading.Event, base_port: int,
400
+ tp_rank: int):
401
"""Background thread for getting new NIXL handshakes."""
402
# NOTE(rob): this is a simple implementation. We will move
390
- # to a better approach like an ETCD server in the future.
391
-
392
- # NOTE(rob): to support heterogeneous TP, we will have to
393
- # move this into the scheduler rather than worker, since
394
- # each rank needs the metadata of all other ranks (whereas
395
- # in this setup, each rank only gets one other rank's meta.
403
+ # to a better approach via HTTP endpoint soon.
404
405
encoder = msgspec.msgpack.Encoder()
406
encoded_data = encoder.encode(metadata)
@@ -402,11 +410,7 @@ class NixlConnectorWorker:
410
411
# Listen for new requests for metadata.
412
host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST
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.
408
- port = envs.VLLM_NIXL_SIDE_CHANNEL_PORT + world_rank
409
- path = make_zmq_path("tcp", host, port)
413
+ path = make_zmq_path("tcp", host, base_port + tp_rank)
414
logger.debug("Starting listening on path: %s", path)
415
with zmq_ctx(zmq.ROUTER, path) as sock:
416
ready_event.set()
@@ -421,10 +425,10 @@ class NixlConnectorWorker:
425
"""Do a NIXL handshake with a remote instance."""
426
427
start_time = time.perf_counter()
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.
427
- path = make_zmq_path("tcp", host, port + self.world_rank)
428
+ # NOTE(rob): we need each tp_rank to have a unique port.
429
+ # This is a hack to keep us moving. We will switch when
430
+ # we switch to HTTP-based NIXL metadata exchange.
431
+ path = make_zmq_path("tcp", host, port + self.tp_rank)
432
logger.debug("Querying metadata on path: %s", path)
433
with zmq_ctx(zmq.REQ, path) as sock:
434
# Send query for the request.
@@ -532,7 +536,7 @@ class NixlConnectorWorker:
536
ready_event = threading.Event()
537
self._nixl_handshake_listener_t = threading.Thread(
538
target=self._nixl_handshake_listener,
535
- args=(metadata, ready_event, self.world_rank),
539
+ args=(metadata, ready_event, self.side_channel_port, self.tp_rank),
540
daemon=True,
541
name="nixl_handshake_listener")
542
self._nixl_handshake_listener_t.start()
@@ -556,9 +560,9 @@ class NixlConnectorWorker:
560
block_offset = block_id * self.block_len
561
# (addr, len, device id)
562
blocks_data.append(
559
- (base_addr + block_offset, self.block_len, self.rank))
560
- logger.debug("Created %s blocks for src engine %s and rank %s",
561
- len(blocks_data), self.engine_id, self.rank)
563
+ (base_addr + block_offset, self.block_len, self.tp_rank))
564
+ logger.debug("Created %s blocks for src engine %s and tp_rank %s",
565
+ len(blocks_data), self.engine_id, self.tp_rank)
566
567
# Register with NIXL.
568
descs = self.nixl_wrapper.get_xfer_descs(blocks_data, "VRAM")
@@ -573,9 +577,9 @@ class NixlConnectorWorker:
577
block_offset = block_id * self.block_len
578
# (addr, len, device id)
579
blocks_data.append(
576
- (base_addr + block_offset, self.block_len, self.rank))
577
- logger.debug("Created %s blocks for dst engine %s and rank %s",
578
- len(blocks_data), engine_id, self.rank)
580
+ (base_addr + block_offset, self.block_len, self.tp_rank))
581
+ logger.debug("Created %s blocks for dst engine %s and tp_rank %s",
582
+ len(blocks_data), engine_id, self.tp_rank)
583
584
# Register with NIXL.
585
descs = self.nixl_wrapper.get_xfer_descs(blocks_data, "VRAM")
@@ -600,14 +604,14 @@ class NixlConnectorWorker:
604
if len(done_sending) > 0 or len(done_recving) > 0:
605
logger.debug(
606
"Rank %s, get_finished: %s requests done sending "
603
- "and %s requests done recving", self.rank, len(done_sending),
604
- len(done_recving))
607
+ "and %s requests done recving", self.tp_rank,
608
+ len(done_sending), len(done_recving))
609
610
if self.world_size == 1:
611
return done_sending, done_recving
612
613
# Rank 0: get finished from all other ranks.
610
- if self.rank == 0:
614
+ if self.tp_rank == 0:
615
for req_id in done_sending:
616
self._done_sending_count[req_id] += 1
617
for req_id in done_recving: