@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / da63274d9f

Fork 0

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 14, 05:30 UTC

[Bugfix][NIXL] Fix Async Scheduler timeout issue (#25808)

Signed-off-by: NickLucche <nlucches@redhat.com>

Nicolò Lucchesi committed Sep 27, 2025 at 21:17 UTC da63274d9f3d06ba5815b5c8786a7194923a0234

1 file changed +28 -3

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

+28 -3

@@ -105,6 +105,7 @@ class NixlConnectorMetadata(KVConnectorMetadata):

105 self.reqs_to_recv: dict[ReqId, ReqMeta] = {}

106 self.reqs_to_save: dict[ReqId, ReqMeta] = {}

107 self.reqs_to_send: dict[ReqId, float] = {}

108 + self.reqs_in_batch: set[ReqId] = set()

109

110 def add_new_req(

111 self,

@@ -278,6 +279,7 @@ class NixlConnectorScheduler:

279 self._reqs_need_save: dict[ReqId, tuple[Request, list[int]]] = {}

280 # Reqs to send and their expiration time

281 self._reqs_need_send: dict[ReqId, float] = {}

282 + self._reqs_in_batch: set[ReqId] = set()

283

284 def get_num_new_matched_tokens(

285 self, request: "Request",

@@ -324,6 +326,9 @@ class NixlConnectorScheduler:

326

327 if not params:

328 return

329 +

330 + if params.get("do_remote_decode"):

331 + self._reqs_in_batch.add(request.request_id)

332 if self.use_host_buffer and params.get("do_remote_decode"):

333 # NOTE: when accelerator is not directly supported by Nixl,

334 # prefilled blocks need to be saved to host memory before transfer.

@@ -373,6 +378,8 @@ class NixlConnectorScheduler:

378 request_id=req_id,

379 local_block_ids=block_ids,

380 kv_transfer_params=req.kv_transfer_params,

381 + load_remote_cache=True,

382 + save_to_host=False,

383 )

384

385 for req_id, (req, block_ids) in self._reqs_need_save.items():

@@ -386,10 +393,12 @@ class NixlConnectorScheduler:

393 )

394

395 meta.reqs_to_send = self._reqs_need_send

396 + meta.reqs_in_batch = self._reqs_in_batch

397

398 # Clear the list once workers start the transfers

399 self._reqs_need_recv.clear()

400 self._reqs_need_save.clear()

401 + self._reqs_in_batch = set()

402 self._reqs_need_send = {}

403

404 return meta

@@ -546,6 +555,8 @@ class NixlConnectorWorker:

555 self._recving_transfers = defaultdict[ReqId, list[Transfer]](list)

556 # Track the expiration time of requests that are waiting to be sent.

557 self._reqs_to_send: dict[ReqId, float] = {}

558 + # Set of requests that have been part of a batch, regardless of status.

559 + self._reqs_to_process: set[ReqId] = set()

560

561 # Background thread for handling new handshake requests.

562 self._nixl_handshake_listener_t: Optional[threading.Thread] = None

@@ -1082,6 +1093,7 @@ class NixlConnectorWorker:

1093 "Releasing expired KV blocks for request %s which were "

1094 "retrieved by %d decode worker(s) within %d seconds.", req_id,

1095 count, envs.VLLM_NIXL_ABORT_REQUEST_TIMEOUT)

1096 + self._reqs_to_process.remove(req_id)

1097 del self._reqs_to_send[req_id]

1098 done_sending.add(req_id)

1099

@@ -1097,7 +1109,8 @@ class NixlConnectorWorker:

1109 for notifs in self.nixl_wrapper.get_new_notifs().values():

1110 for notif in notifs:

1111 req_id, tp_ratio = notif.decode("utf-8").rsplit(":", 1)

1100 - if req_id not in self._reqs_to_send:

1112 + if (req_id not in self._reqs_to_send

1113 + and req_id not in self._reqs_to_process):

1114 logger.error(

1115 "Potentially invalid KV blocks for "

1116 "unrecognized request %s were retrieved by "

@@ -1110,7 +1123,8 @@ class NixlConnectorWorker:

1123 tp_ratio):

1124 notified_req_ids.add(req_id)

1125 del self.consumer_notification_counts_by_req[req_id]

1113 - del self._reqs_to_send[req_id]

1126 + self._reqs_to_process.remove(req_id)

1127 + self._reqs_to_send.pop(req_id, None)

1128 return notified_req_ids

1129

1130 def _pop_done_transfers(

@@ -1171,8 +1185,19 @@ class NixlConnectorWorker:

1185 while not self._ready_requests.empty():

1186 self._read_blocks_for_req(*self._ready_requests.get_nowait())

1187

1188 + # Keep around the requests that have been part of a batch. This is

1189 + # needed because async scheduling pushes the misalignment between the

1190 + # moment in which requests expiration is set (P side) and the moment in

1191 + # which blocks are read from D. As P can now more easily lag behind D

1192 + # while processing the next batch, we make sure to only set an

1193 + # expiration for requests that have not been read from D yet.

1194 + for req_id in metadata.reqs_in_batch:

1195 + self._reqs_to_process.add(req_id)

1196 +

1197 # Add to requests that are waiting to be read and track expiration.

1175 - self._reqs_to_send.update(metadata.reqs_to_send)

1198 + for req_id, expiration_time in metadata.reqs_to_send.items():

1199 + if req_id in self._reqs_to_process:

1200 + self._reqs_to_send[req_id] = expiration_time

1201

1202 def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):

1203 logger.debug(