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 7, 00:15 UTC
[bugfix] fix nvlink for nixl/ucx (#36475)
Signed-off-by: youkaichao <youkaichao@gmail.com>
youkaichao committed
Mar 10, 2026 at 07:49 UTC
f85b4eda3a22fedd885ef31650c825d56867587e
1 file changed
+13
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+13
@@ -1141,6 +1141,19 @@ class NixlConnectorWorker:
1141
expected_engine_id: str,
1142
) -> dict[int, str]:
1143
"""Do a NIXL handshake with a remote instance."""
1144
+
1145
+ # the first time we connect to a remote agent.
1146
+ # be careful, the handshake happens in a background thread.
1147
+ # it does not have an active cuda context until any cuda runtime
1148
+ # call is made. when UCX fails to find a valid cuda context, it will
1149
+ # disable any cuda ipc communication, essentially disabling any NVLink
1150
+ # communication.
1151
+ # when we are using device buffers, we need to set the device
1152
+ # explicitly to make sure the handshake background thread has a valid
1153
+ # cuda context.
1154
+ if not self.use_host_buffer:
1155
+ current_platform.set_device(self.device_id)
1156
+
1157
# When target instance TP > local TP, we need to perform multiple
1158
# handshakes. Do it in a single background job for simplicity.
1159
# Regardless, only handshake with the remote TP rank(s) that current