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 11, 23:45 UTC
[NIXL] Add remote_request_id to kv_transfer_params (#29665)
Signed-off-by: Mark McLoughlin <markmc@redhat.com>
Mark McLoughlin committed
Dec 5, 2025 at 17:43 UTC
dff0a2b39475096f5456721bfc8df3c7fea3cc57
3 files changed
+21
-3
tests/v1/kv_connector/unit/test_nixl_connector.py
+6
@@ -470,6 +470,7 @@ class TestNixlHandshake:
470
num_xfers + 6,
471
],
472
"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
473
+ "remote_request_id": f"prefill-{request_id}",
474
"remote_host": "localhost",
475
"remote_port": 1234,
476
"remote_tp_size": 1,
@@ -536,6 +537,7 @@ class TestNixlHandshake:
537
kv_transfer_params={
538
"remote_block_ids": [4, 5, 6],
539
"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
540
+ "remote_request_id": "prefill-id",
541
"remote_host": "localhost",
542
"remote_port": 1234,
543
"remote_tp_size": prefill_tp_size,
@@ -591,6 +593,7 @@ class TestNixlHandshake:
593
kv_transfer_params={
594
"remote_block_ids": [4, 5, 6],
595
"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
596
+ "remote_request_id": f"prefill-id-{i}",
597
"remote_host": "localhost",
598
"remote_port": 1234,
599
"remote_tp_size": 1,
@@ -754,6 +757,7 @@ def test_kv_connector_stats(dist_init):
757
kv_transfer_params={
758
"remote_block_ids": [4, 5, 6],
759
"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
760
+ "remote_request_id": f"prefill-{request_id}",
761
"remote_host": "localhost",
762
"remote_port": 1234,
763
"remote_tp_size": 1,
@@ -1470,6 +1474,7 @@ def test_handshake_failure_returns_finished(dist_init):
1474
kv_transfer_params={
1475
"remote_block_ids": [4, 5, 6],
1476
"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1477
+ "remote_request_id": f"prefill-{request_id}",
1478
"remote_host": "localhost",
1479
"remote_port": 1234,
1480
"remote_tp_size": 1,
@@ -1519,6 +1524,7 @@ def test_transfer_setup_failure_returns_finished(dist_init):
1524
kv_transfer_params={
1525
"remote_block_ids": [10, 11, 12],
1526
"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1527
+ "remote_request_id": f"prefill-{request_id}",
1528
"remote_host": "localhost",
1529
"remote_port": 1234,
1530
"remote_tp_size": 1,
tests/v1/kv_connector/unit/utils.py
+1
@@ -194,6 +194,7 @@ def create_request(
194
do_remote_prefill=True,
195
do_remote_decode=False,
196
remote_engine_id="my-engine-id",
197
+ remote_request_id=f"prefill-{request_id}",
198
remote_block_ids=list(range(num_remote_blocks)),
199
remote_host="my-host",
200
remote_port=1234,
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+14
-3
@@ -71,8 +71,9 @@ ReqId = str
71
#
72
# Version History:
73
# 1: Initial version with compatibility checking
74
+# 2: Add remote_request_id to kv_transfer_params
75
#
75
-NIXL_CONNECTOR_VERSION: int = 1
76
+NIXL_CONNECTOR_VERSION: int = 2
77
78
GET_META_MSG = b"get_meta_msg"
79
@@ -210,6 +211,7 @@ class ReqMeta:
211
remote_host: str
212
remote_port: int
213
remote_engine_id: str
214
+ remote_request_id: str
215
tp_size: int
216
217
@@ -236,6 +238,7 @@ class NixlConnectorMetadata(KVConnectorMetadata):
238
local_physical_block_ids=local_block_ids,
239
remote_block_ids=kv_transfer_params["remote_block_ids"],
240
remote_engine_id=kv_transfer_params["remote_engine_id"],
241
+ remote_request_id=kv_transfer_params["remote_request_id"],
242
remote_host=kv_transfer_params["remote_host"],
243
remote_port=kv_transfer_params["remote_port"],
244
# P workers don't need to receive tp_size from proxy here.
@@ -622,7 +625,12 @@ class NixlConnectorScheduler:
625
if params.get("remote_block_ids"):
626
if all(
627
p in params
625
- for p in ("remote_engine_id", "remote_host", "remote_port")
628
+ for p in (
629
+ "remote_engine_id",
630
+ "remote_request_id",
631
+ "remote_host",
632
+ "remote_port",
633
+ )
634
):
635
# If remote_blocks and num_external_tokens = 0, we have
636
# a full prefix cache hit on the D worker. We need to call
@@ -751,6 +759,7 @@ class NixlConnectorScheduler:
759
do_remote_decode=False,
760
remote_block_ids=block_ids,
761
remote_engine_id=self.engine_id,
762
+ remote_request_id=request.request_id,
763
remote_host=self.side_channel_host,
764
remote_port=self.side_channel_port,
765
tp_size=self.vllm_config.parallel_config.tensor_parallel_size,
@@ -1964,6 +1973,7 @@ class NixlConnectorWorker:
1973
self._read_blocks(
1974
request_id=req_id,
1975
dst_engine_id=meta.remote_engine_id,
1976
+ remote_request_id=meta.remote_request_id,
1977
local_block_ids=meta.local_physical_block_ids,
1978
remote_block_ids=meta.remote_block_ids,
1979
)
@@ -1974,6 +1984,7 @@ class NixlConnectorWorker:
1984
remote_block_ids: list[int],
1985
dst_engine_id: str,
1986
request_id: str,
1987
+ remote_request_id: str,
1988
):
1989
block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(dst_engine_id)
1990
if block_size_ratio > 1:
@@ -2006,7 +2017,7 @@ class NixlConnectorWorker:
2017
# Number of D TP workers that will read from dst P. Propagate tp_ratio
2018
# on notification so that dst worker can wait before freeing blocks.
2019
tp_ratio = self.kv_topo.tp_ratio_from_engine_id(dst_engine_id)
2009
- notif_id = f"{request_id}:{tp_ratio}".encode()
2020
+ notif_id = f"{remote_request_id}:{tp_ratio}".encode()
2021
2022
# Full prefix cache hit: do not need to read remote blocks,
2023
# just notify P worker that we have the blocks we need.