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, 03:30 UTC
[Misc] Minor refactor of NIXL background handshake (#20068)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Jul 1, 2025 at 13:40 UTC
650d5dbd04e92f5043a11e4a4d86d4f39ee1b694
1 file changed
+32
-28
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+32
-28
@@ -515,6 +515,33 @@ class NixlConnectorWorker:
515
# Remote rank -> agent name.
516
return {p_remote_rank: handshake(path, p_remote_rank)}
517
518
+ def _background_nixl_handshake(self, req_id: str,
519
+ remote_engine_id: EngineId, meta: ReqMeta):
520
+ # Do NIXL handshake in background and add to _ready_requests when done.
521
+ fut = self._handshake_futures.get(remote_engine_id)
522
+ if fut is None:
523
+ fut = self._handshake_initiation_executor.submit(
524
+ self._nixl_handshake, meta.remote_host, meta.remote_port,
525
+ meta.tp_size)
526
+ self._handshake_futures[remote_engine_id] = fut
527
+
528
+ def done_callback(f: Future[dict[int, str]], eid=remote_engine_id):
529
+ with self._handshake_lock:
530
+ del self._handshake_futures[eid]
531
+ try:
532
+ self._remote_agents[eid] = f.result()
533
+ except Exception:
534
+ logger.exception("Handshake with %s failed", eid)
535
+
536
+ fut.add_done_callback(done_callback)
537
+
538
+ # TODO: handle failure state of future in the
539
+ # callback, we want to fail the request in this case.
540
+ def request_ready(_f: Future[Any], entry=(req_id, meta)):
541
+ self._ready_requests.put(entry)
542
+
543
+ fut.add_done_callback(request_ready)
544
+
545
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
546
"""Register the KV Cache data in nixl."""
547
@@ -902,37 +929,14 @@ class NixlConnectorWorker:
929
remote_engine_id, len(meta.local_block_ids),
930
len(meta.remote_block_ids))
931
if remote_engine_id not in self._remote_agents:
905
- # Being optimistic to assume engine is usually ready, apply
906
- # lock only when the optimistic check fails.
932
+ # Initiate handshake with remote engine to exchange metadata.
933
with self._handshake_lock:
934
if remote_engine_id not in self._remote_agents:
909
- fut = self._handshake_futures.get(remote_engine_id)
910
- if fut is None:
911
- fut = self._handshake_initiation_executor.submit(
912
- self._nixl_handshake, meta.remote_host,
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]],
917
- eid=remote_engine_id):
918
- with self._handshake_lock:
919
- del self._handshake_futures[eid]
920
- try:
921
- self._remote_agents[eid] = f.result()
922
- except Exception:
923
- logger.exception(
924
- "Handshake with %s failed", eid)
925
-
926
- fut.add_done_callback(done_callback)
927
-
928
- # TODO: handle failure state of future in the
929
- # callback, we want to fail the request in this case.
930
- def request_ready(_f: Future[Any],
931
- entry=(req_id, meta)):
932
- self._ready_requests.put(entry)
933
-
934
- fut.add_done_callback(request_ready)
935
+ self._background_nixl_handshake(
936
+ req_id, remote_engine_id, meta)
937
continue
938
+
939
+ # Handshake already completed, start async read xfer.
940
self._read_blocks_for_req(req_id, meta)
941
942
# Start transfers for requests whose handshakes have now finished.