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, 04:45 UTC
[PD] Skip `tp_size` exchange with rank0 (#19413)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Jun 26, 2025 at 05:04 UTC
2582683566ed676a811f4311f1048f0b323676b8
2 files changed
+72
-66
tests/v1/kv_connector/unit/test_nixl_connector.py
+23
-6
@@ -7,6 +7,8 @@ from collections import defaultdict
7
from typing import Optional
8
from unittest.mock import patch
9
10
+import pytest
11
+
12
from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
13
KVConnectorRole, NixlAgentMetadata, NixlConnector, NixlConnectorMetadata,
14
NixlConnectorWorker)
@@ -161,7 +163,8 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
163
super().__init__(*args, **kwargs)
164
self._hand_shake_latency = hand_shake_latency
165
164
- def _nixl_handshake(self, host: str, port: int) -> dict[int, str]:
166
+ def _nixl_handshake(self, host: str, port: int,
167
+ remote_tp_size: int) -> dict[int, str]:
168
# Mimic slow _nixl_handshake, as well as bypass zmq communication.
169
time.sleep(self._hand_shake_latency)
170
# These should've been done in register_kv_caches(), called by
@@ -177,10 +180,10 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
180
agent_metadata=FakeNixlWrapper.AGENT_METADATA,
181
kv_caches_base_addr=[0],
182
num_blocks=1,
180
- tp_size=1,
183
block_len=self.block_len,
184
attn_backend_name=self.backend_name,
183
- ))
185
+ ),
186
+ remote_tp_size=remote_tp_size)
187
return {0: remote_agent_name}
188
189
@@ -233,6 +236,8 @@ class TestNixlHandshake:
236
"localhost",
237
"remote_port":
238
1234,
239
+ "remote_tp_size":
240
+ 1,
241
})
242
connector.bind_connector_metadata(metadata)
243
@@ -259,13 +264,23 @@ class TestNixlHandshake:
264
@patch(
265
"vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
266
FakeNixlWrapper)
267
+ @pytest.mark.parametrize("decode_tp_size, prefill_tp_size", [
268
+ (1, 1),
269
+ (2, 1),
270
+ (4, 2),
271
+ (4, 4),
272
+ ])
273
def test_async_load_kv(
263
- self,
264
- # dist_init is a fixture that initializes the distributed environment.
265
- dist_init):
274
+ self,
275
+ # Fixture that initializes the distributed environment.
276
+ dist_init,
277
+ # Simulate consumer-producer TP sizes.
278
+ decode_tp_size,
279
+ prefill_tp_size):
280
"""Test that NixlConnector's start_load_kv should be non-blocking."""
281
282
vllm_config = create_vllm_config()
283
+ vllm_config.parallel_config.tensor_parallel_size = decode_tp_size
284
285
# Test worker role in decode server.
286
connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
@@ -280,6 +295,7 @@ class TestNixlHandshake:
295
FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
296
"remote_host": "localhost",
297
"remote_port": 1234,
298
+ "remote_tp_size": prefill_tp_size,
299
})
300
connector.bind_connector_metadata(metadata)
301
@@ -329,6 +345,7 @@ class TestNixlHandshake:
345
FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
346
"remote_host": "localhost",
347
"remote_port": 1234,
348
+ "remote_tp_size": 1,
349
})
350
connector.bind_connector_metadata(metadata)
351
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+49
-60
@@ -62,7 +62,6 @@ class NixlAgentMetadata(
62
agent_metadata: bytes
63
kv_caches_base_addr: list[int]
64
num_blocks: int
65
- tp_size: int
65
block_len: int
66
attn_backend_name: str
67
@@ -73,7 +72,8 @@ class ReqMeta:
72
remote_block_ids: list[int]
73
remote_host: str
74
remote_port: int
76
- remote_engine_id: EngineId
75
+ remote_engine_id: str
76
+ tp_size: int
77
78
79
class NixlConnectorMetadata(KVConnectorMetadata):
@@ -93,6 +93,8 @@ class NixlConnectorMetadata(KVConnectorMetadata):
93
remote_engine_id=kv_transfer_params["remote_engine_id"],
94
remote_host=kv_transfer_params["remote_host"],
95
remote_port=kv_transfer_params["remote_port"],
96
+ # P workers don't need to receive tp_size from proxy here.
97
+ tp_size=kv_transfer_params.get("tp_size", 1),
98
)
99
100
@@ -330,7 +332,7 @@ class NixlConnectorScheduler:
332
remote_engine_id=self.engine_id,
333
remote_host=self.side_channel_host,
334
remote_port=self.side_channel_port,
333
- )
335
+ tp_size=self.vllm_config.parallel_config.tensor_parallel_size)
336
337
338
class NixlConnectorWorker:
@@ -473,7 +475,8 @@ class NixlConnectorWorker:
475
"Connection listener got unexpected message %s", msg)
476
sock.send_multipart((identity, b"", encoded_data))
477
476
- def _nixl_handshake(self, host: str, port: int) -> dict[int, str]:
478
+ def _nixl_handshake(self, host: str, port: int,
479
+ remote_tp_size: int) -> dict[int, str]:
480
"""Do a NIXL handshake with a remote instance."""
481
482
start_time = time.perf_counter()
@@ -482,7 +485,7 @@ class NixlConnectorWorker:
485
# a hack to keep us moving. We will switch when moving to etcd
486
# or where we have a single ZMQ socket in the scheduler.
487
485
- def handshake(path: str, rank: int) -> tuple[NixlAgentMetadata, str]:
488
+ def handshake(path: str, rank: int) -> str:
489
# Send query for the request.
490
with zmq_ctx(zmq.REQ, path) as sock:
491
sock.send(GET_META_MSG)
@@ -492,33 +495,25 @@ class NixlConnectorWorker:
495
got_metadata_time = time.perf_counter()
496
497
# Register Remote agent.
495
- remote_agent_name = self.add_remote_agent(metadata, rank)
498
+ remote_agent_name = self.add_remote_agent(
499
+ metadata, rank, remote_tp_size)
500
setup_agent_time = time.perf_counter()
501
502
logger.debug("NIXL handshake: get metadata took: %s",
503
got_metadata_time - start_time)
504
logger.debug("NIXL handshake: add agent took: %s",
505
setup_agent_time - got_metadata_time)
502
- return metadata, remote_agent_name
506
+ return remote_agent_name
507
504
- # Handshake with remote agent-rank0 first to get the tp_size of remote
505
- path = make_zmq_path("tcp", host, port)
506
- logger.debug("Querying master rank metadata on path: %s", path)
507
- rank_to_agent_name: dict[int, str] = {}
508
- metadata, rank_to_agent_name[0] = handshake(path, 0)
509
-
510
- # Handshake only with the other TP remote the current local rank will
508
+ # Handshake only with the remote TP rank that current local rank will
509
# pull from. With homogeneous TP it happens to be the same rank_i.
512
- tp_ratio = self._tp_size[self.engine_id] // metadata.tp_size
510
+ tp_ratio = self._tp_size[self.engine_id] // remote_tp_size
511
p_remote_rank = self.tp_rank // tp_ratio
514
- if p_remote_rank > 0:
515
- path = make_zmq_path("tcp", host, port + p_remote_rank)
516
- logger.debug("Querying metadata on path: %s at remote rank %s",
517
- path, p_remote_rank)
518
- _, rank_to_agent_name[p_remote_rank] = handshake(
519
- path, p_remote_rank)
520
-
521
- return rank_to_agent_name
512
+ path = make_zmq_path("tcp", host, port + p_remote_rank)
513
+ logger.debug("Querying metadata on path: %s at remote rank %s", path,
514
+ p_remote_rank)
515
+ # Remote rank -> agent name.
516
+ return {p_remote_rank: handshake(path, p_remote_rank)}
517
518
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
519
"""Register the KV Cache data in nixl."""
@@ -645,7 +640,6 @@ class NixlConnectorWorker:
640
agent_metadata=self.nixl_wrapper.get_agent_metadata(),
641
kv_caches_base_addr=self.kv_caches_base_addr[self.engine_id],
642
num_blocks=self.num_blocks,
648
- tp_size=self.world_size,
643
block_len=self.block_len,
644
attn_backend_name=self.backend_name)
645
ready_event = threading.Event()
@@ -659,7 +653,8 @@ class NixlConnectorWorker:
653
654
def add_remote_agent(self,
655
nixl_agent_meta: NixlAgentMetadata,
662
- remote_tp_rank: int = 0) -> str:
656
+ remote_tp_rank: int = 0,
657
+ remote_tp_size: int = 1) -> str:
658
"""
659
Add the remote NIXL agent and prepare the descriptors for reading cache
660
blocks from remote.
@@ -704,9 +699,9 @@ class NixlConnectorWorker:
699
return self._remote_agents[engine_id][remote_tp_rank]
700
701
if engine_id in self._tp_size:
707
- assert self._tp_size[engine_id] == nixl_agent_meta.tp_size
702
+ assert self._tp_size[engine_id] == remote_tp_size
703
else:
709
- self._tp_size[engine_id] = nixl_agent_meta.tp_size
704
+ self._tp_size[engine_id] = remote_tp_size
705
# We may eventually enable this after asserting equality in cache
706
# layout and close outputs.
707
assert nixl_agent_meta.attn_backend_name == self.backend_name
@@ -756,33 +751,31 @@ class NixlConnectorWorker:
751
# rank. With heterogeneous TP, prepare the descriptors by splitting the
752
# P KV cache along kv_head dim, of D worker's kv_head size (D>P).
753
# Eg. PTP1 DTP2 => P0 KV:[block0-KV_0 | block0-KV_1..].
759
- p_remote_tp_rank = self.tp_rank // tp_ratio
754
# Only register the remote's descriptors if current rank pulls from it.
761
- if p_remote_tp_rank == remote_tp_rank:
762
- self.kv_caches_base_addr[
763
- engine_id] = nixl_agent_meta.kv_caches_base_addr
764
- rank_offset = self.tp_rank % tp_ratio * self.block_len \
765
- if not (self.use_mla or is_kv_replicated) else 0
766
- # Register all remote blocks, but only the corresponding kv heads.
767
- for base_addr in nixl_agent_meta.kv_caches_base_addr:
768
- for block_id in range(nixl_agent_meta.num_blocks):
769
- block_offset = block_id * nixl_agent_meta.block_len
770
- # For each block, grab the heads chunk belonging to rank_i
771
- # of size remote_nheads // tp_ratio, which correspond to
772
- # self.block_len == remote_block_len//tp_ratio bytes.
773
- addr = base_addr + block_offset + rank_offset
774
- # (addr, len, device id)
775
- blocks_data.append((addr, self.block_len, remote_tp_rank))
776
- logger.debug(
777
- "Created %s blocks for dst engine %s with remote rank %s and "
778
- "local rank %s", len(blocks_data), engine_id, remote_tp_rank,
779
- self.tp_rank)
755
+ self.kv_caches_base_addr[
756
+ engine_id] = nixl_agent_meta.kv_caches_base_addr
757
+ rank_offset = self.tp_rank % tp_ratio * self.block_len \
758
+ if not (self.use_mla or is_kv_replicated) else 0
759
+ # Register all remote blocks, but only the corresponding kv heads.
760
+ for base_addr in nixl_agent_meta.kv_caches_base_addr:
761
+ for block_id in range(nixl_agent_meta.num_blocks):
762
+ block_offset = block_id * nixl_agent_meta.block_len
763
+ # For each block, grab the heads chunk belonging to rank_i
764
+ # of size remote_nheads // tp_ratio, which correspond to
765
+ # self.block_len == remote_block_len//tp_ratio bytes.
766
+ addr = base_addr + block_offset + rank_offset
767
+ # (addr, len, device id)
768
+ blocks_data.append((addr, self.block_len, remote_tp_rank))
769
+ logger.debug(
770
+ "Created %s blocks for dst engine %s with remote rank %s and "
771
+ "local rank %s", len(blocks_data), engine_id, remote_tp_rank,
772
+ self.tp_rank)
773
781
- # Register with NIXL.
782
- descs = self.nixl_wrapper.get_xfer_descs(blocks_data, "VRAM")
783
- self.dst_xfer_side_handles[
784
- engine_id] = self.nixl_wrapper.prep_xfer_dlist(
785
- remote_agent_name, descs)
774
+ # Register with NIXL.
775
+ descs = self.nixl_wrapper.get_xfer_descs(blocks_data, "VRAM")
776
+ self.dst_xfer_side_handles[
777
+ engine_id] = self.nixl_wrapper.prep_xfer_dlist(
778
+ remote_agent_name, descs)
779
780
return remote_agent_name
781
@@ -917,7 +910,7 @@ class NixlConnectorWorker:
910
if fut is None:
911
fut = self._handshake_initiation_executor.submit(
912
self._nixl_handshake, meta.remote_host,
920
- meta.remote_port)
913
+ meta.remote_port, meta.tp_size)
914
self._handshake_futures[remote_engine_id] = fut
915
916
def done_callback(f: Future[dict[int, str]],
@@ -957,13 +950,9 @@ class NixlConnectorWorker:
950
remote_block_ids=meta.remote_block_ids,
951
)
952
960
- def _read_blocks(
961
- self,
962
- local_block_ids: list[int],
963
- remote_block_ids: list[int],
964
- dst_engine_id: str,
965
- request_id: str,
966
- ):
953
+ def _read_blocks(self, local_block_ids: list[int],
954
+ remote_block_ids: list[int], dst_engine_id: str,
955
+ request_id: str):
956
# NOTE(rob): having the staging blocks be on the READER side is
957
# not going to work well (since we will have to call rearrange tensors).
958
# after we detect the txn is complete (which means we cannot make the