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, 12:30 UTC
[NIXL][BUG FIX] Fix a bug for PD with host_buffer after merging 29665 (#30420)
Signed-off-by: Chendi Xue <chendi.xue@intel.com> Signed-off-by: Mark McLoughlin <markmc@redhat.com> Co-authored-by: Mark McLoughlin <markmc@redhat.com>
Chendi.Xue committed
Dec 14, 2025 at 09:38 UTC
ae2e503dda693b6b7ab9052ec61e012a3c730f2f
2 files changed
+62
-45
tests/v1/kv_connector/unit/test_nixl_connector.py
+6
-6
@@ -461,7 +461,7 @@ class TestNixlHandshake:
461
metadata = NixlConnectorMetadata()
462
if num_xfers > 0:
463
num_xfers -= 1
464
- metadata.add_new_req(
464
+ metadata.add_new_req_to_recv(
465
request_id=request_id,
466
local_block_ids=[num_xfers + 1, num_xfers + 2, num_xfers + 3],
467
kv_transfer_params={
@@ -532,7 +532,7 @@ class TestNixlHandshake:
532
vllm_config, connector.engine_id
533
)
534
metadata = NixlConnectorMetadata()
535
- metadata.add_new_req(
535
+ metadata.add_new_req_to_recv(
536
request_id="id",
537
local_block_ids=[1, 2, 3],
538
kv_transfer_params={
@@ -588,7 +588,7 @@ class TestNixlHandshake:
588
metadata = NixlConnectorMetadata()
589
total_reqs = 5
590
for i in range(total_reqs):
591
- metadata.add_new_req(
591
+ metadata.add_new_req_to_recv(
592
request_id=f"id_{i}",
593
local_block_ids=[1, 2, 3],
594
kv_transfer_params={
@@ -752,7 +752,7 @@ def test_kv_connector_stats(dist_init):
752
# Create transfer metadata
753
request_id = "test_req_for_stats"
754
metadata = NixlConnectorMetadata()
755
- metadata.add_new_req(
755
+ metadata.add_new_req_to_recv(
756
request_id=request_id,
757
local_block_ids=[1, 2, 3],
758
kv_transfer_params={
@@ -1515,7 +1515,7 @@ def test_handshake_failure_returns_finished(dist_init):
1515
1516
request_id = "test_handshake_fail"
1517
metadata = NixlConnectorMetadata()
1518
- metadata.add_new_req(
1518
+ metadata.add_new_req_to_recv(
1519
request_id=request_id,
1520
local_block_ids=[1, 2, 3],
1521
kv_transfer_params={
@@ -1565,7 +1565,7 @@ def test_transfer_setup_failure_returns_finished(dist_init):
1565
1566
request_id = "test_transfer_fail"
1567
metadata = NixlConnectorMetadata()
1568
- metadata.add_new_req(
1568
+ metadata.add_new_req_to_recv(
1569
request_id=request_id,
1570
local_block_ids=[7, 8, 9],
1571
kv_transfer_params={
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+56
-39
@@ -202,17 +202,22 @@ def compute_nixl_compatibility_hash(
202
return compat_hash
203
204
205
+@dataclass
206
+class RemoteMeta:
207
+ block_ids: list[int]
208
+ host: str
209
+ port: int
210
+ engine_id: str
211
+ request_id: str
212
+
213
+
214
@dataclass
215
class ReqMeta:
216
local_block_ids: list[int]
217
# To be used when logical block size does not match the kernel block size
218
local_physical_block_ids: list[int]
210
- remote_block_ids: list[int]
211
- remote_host: str
212
- remote_port: int
213
- remote_engine_id: str
214
- remote_request_id: str
219
tp_size: int
220
+ remote: RemoteMeta | None = None
221
222
223
class NixlConnectorMetadata(KVConnectorMetadata):
@@ -223,31 +228,43 @@ class NixlConnectorMetadata(KVConnectorMetadata):
228
self.reqs_in_batch: set[ReqId] = set()
229
self.reqs_not_processed: set[ReqId] = set()
230
226
- def add_new_req(
231
+ def _add_new_req(
232
self,
228
- request_id: ReqId,
233
local_block_ids: list[int],
234
kv_transfer_params: dict[str, Any],
231
- load_remote_cache: bool = True,
232
- save_to_host: bool = False,
233
- ):
234
- # save and load are mutually exclusive
235
- assert load_remote_cache ^ save_to_host
236
- _req = ReqMeta(
235
+ ) -> ReqMeta:
236
+ return ReqMeta(
237
local_block_ids=local_block_ids,
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"],
239
# P workers don't need to receive tp_size from proxy here.
240
tp_size=kv_transfer_params.get("tp_size", 1),
241
)
247
- if save_to_host:
248
- self.reqs_to_save[request_id] = _req
249
- if load_remote_cache:
250
- self.reqs_to_recv[request_id] = _req
242
+
243
+ def add_new_req_to_save(
244
+ self,
245
+ request_id: ReqId,
246
+ local_block_ids: list[int],
247
+ kv_transfer_params: dict[str, Any],
248
+ ):
249
+ self.reqs_to_save[request_id] = self._add_new_req(
250
+ local_block_ids, kv_transfer_params
251
+ )
252
+
253
+ def add_new_req_to_recv(
254
+ self,
255
+ request_id: ReqId,
256
+ local_block_ids: list[int],
257
+ kv_transfer_params: dict[str, Any],
258
+ ):
259
+ req = self._add_new_req(local_block_ids, kv_transfer_params)
260
+ req.remote = RemoteMeta(
261
+ block_ids=kv_transfer_params["remote_block_ids"],
262
+ engine_id=kv_transfer_params["remote_engine_id"],
263
+ request_id=kv_transfer_params["remote_request_id"],
264
+ host=kv_transfer_params["remote_host"],
265
+ port=kv_transfer_params["remote_port"],
266
+ )
267
+ self.reqs_to_recv[request_id] = req
268
269
270
class NixlConnector(KVConnectorBase_V1):
@@ -666,22 +683,18 @@ class NixlConnectorScheduler:
683
# Loop through scheduled reqs and convert to ReqMeta.
684
for req_id, (req, block_ids) in self._reqs_need_recv.items():
685
assert req.kv_transfer_params is not None
669
- meta.add_new_req(
686
+ meta.add_new_req_to_recv(
687
request_id=req_id,
688
local_block_ids=block_ids,
689
kv_transfer_params=req.kv_transfer_params,
673
- load_remote_cache=True,
674
- save_to_host=False,
690
)
691
692
for req_id, (req, block_ids) in self._reqs_need_save.items():
693
assert req.kv_transfer_params is not None
679
- meta.add_new_req(
694
+ meta.add_new_req_to_save(
695
request_id=req_id,
696
local_block_ids=block_ids,
697
kv_transfer_params=req.kv_transfer_params,
683
- load_remote_cache=False,
684
- save_to_host=True,
698
)
699
700
meta.reqs_to_send = self._reqs_need_send
@@ -1124,10 +1137,11 @@ class NixlConnectorWorker:
1137
# Do NIXL handshake in background and add to _ready_requests when done.
1138
fut = self._handshake_futures.get(remote_engine_id)
1139
if fut is None:
1140
+ assert meta.remote is not None
1141
fut = self._handshake_initiation_executor.submit(
1142
self._nixl_handshake,
1129
- meta.remote_host,
1130
- meta.remote_port,
1143
+ meta.remote.host,
1144
+ meta.remote.port,
1145
meta.tp_size,
1146
remote_engine_id,
1147
)
@@ -1774,6 +1788,7 @@ class NixlConnectorWorker:
1788
# clean up metadata for completed requests
1789
meta = self._recving_metadata.pop(req_id, None)
1790
assert meta is not None, f"{req_id} not found in recving_metadata list"
1791
+ assert meta.remote is not None
1792
if self.use_host_buffer:
1793
self.sync_recved_kv_to_device(req_id, meta)
1794
if self.enable_permute_local_kv:
@@ -1781,7 +1796,7 @@ class NixlConnectorWorker:
1796
1797
# post processing for heteroblocksize
1798
block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(
1784
- meta.remote_engine_id
1799
+ meta.remote.engine_id
1800
)
1801
if (
1802
not self.use_mla
@@ -1916,17 +1931,18 @@ class NixlConnectorWorker:
1931
meta.local_physical_block_ids = self._logical_to_kernel_block_ids(
1932
meta.local_block_ids
1933
)
1919
- meta.remote_block_ids = self._logical_to_kernel_block_ids(
1920
- meta.remote_block_ids
1934
+ assert meta.remote is not None
1935
+ meta.remote.block_ids = self._logical_to_kernel_block_ids(
1936
+ meta.remote.block_ids
1937
)
1922
- remote_engine_id = meta.remote_engine_id
1938
+ remote_engine_id = meta.remote.engine_id
1939
logger.debug(
1940
"start_load_kv for request %s from remote engine %s. "
1941
"Num local_block_ids: %s. Num remote_block_ids: %s. ",
1942
req_id,
1943
remote_engine_id,
1944
len(meta.local_physical_block_ids),
1929
- len(meta.remote_block_ids),
1945
+ len(meta.remote.block_ids),
1946
)
1947
# always store metadata for failure recovery
1948
self._recving_metadata[req_id] = meta
@@ -1965,17 +1981,18 @@ class NixlConnectorWorker:
1981
self._reqs_to_send[req_id] = expiration_time
1982
1983
def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
1984
+ assert meta.remote is not None
1985
logger.debug(
1986
"Remote agent %s available, calling _read_blocks for req %s",
1970
- meta.remote_engine_id,
1987
+ meta.remote.engine_id,
1988
req_id,
1989
)
1990
self._read_blocks(
1991
request_id=req_id,
1975
- dst_engine_id=meta.remote_engine_id,
1976
- remote_request_id=meta.remote_request_id,
1992
+ dst_engine_id=meta.remote.engine_id,
1993
+ remote_request_id=meta.remote.request_id,
1994
local_block_ids=meta.local_physical_block_ids,
1978
- remote_block_ids=meta.remote_block_ids,
1995
+ remote_block_ids=meta.remote.block_ids,
1996
)
1997
1998
def _read_blocks(