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 5, 04:45 UTC
[Bugfix][Mooncake] Fix thread-local CUDA context for NVLink transfers in _send_blocks (#39548)
Signed-off-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: Zhewen Li <zhewenli@inferact.ai>
Zhewen Li committed
Apr 14, 2026 at 14:13 UTC
1696c864b90afb9cec925312d0cf12c396517007
1 file changed
+14
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+14
@@ -41,6 +41,7 @@ from vllm.distributed.parallel_state import (
41
)
42
from vllm.forward_context import ForwardContext
43
from vllm.logger import init_logger
44
+from vllm.platforms import current_platform
45
from vllm.utils.network_utils import get_ip, make_zmq_path, make_zmq_socket
46
from vllm.v1.attention.backend import AttentionMetadata
47
from vllm.v1.attention.backends.utils import get_kv_cache_layout
@@ -645,6 +646,10 @@ class MooncakeConnectorWorker:
646
logger.info("Initializing Mooncake Transfer Engine worker %s", engine_id)
647
648
self.vllm_config = vllm_config
649
+ # Capture device BEFORE TransferEngine init — MNNVL's NVLink allocator
650
+ # may change the current CUDA device during engine.initialize().
651
+ self.device_id = torch.accelerator.current_device_index()
652
+ current_platform.set_device(self.device_id)
653
654
self.engine = TransferEngine()
655
self.hostname = get_ip()
@@ -705,9 +710,12 @@ class MooncakeConnectorWorker:
710
# For kv_both, we will act both prefiller and decoder.
711
if not self.is_kv_consumer:
712
# Background threads for sending kvcaches to D.
713
+ # Each pool thread must be bound to the correct CUDA device
714
+ # because CUDA device selection is thread-local.
715
self._sender_executor = ThreadPoolExecutor(
716
max_workers=self.num_sender_workers,
717
thread_name_prefix="vllm-mooncake-sender",
718
+ initializer=self._bind_sender_thread_device,
719
)
720
logger.debug(
721
"Mooncake Prefiller: use %d workers to send kvcaches",
@@ -1193,6 +1201,12 @@ class MooncakeConnectorWorker:
1201
1202
return src_ptrs, dst_ptrs, lengths, err_reqs, err_msg
1203
1204
+ def _bind_sender_thread_device(self) -> None:
1205
+ """ThreadPoolExecutor initializer — binds each pool thread to the
1206
+ correct CUDA device. CUDA device selection is thread-local, so
1207
+ without this, NVLink transfers fail for TP ranks > 0."""
1208
+ current_platform.set_device(self.device_id)
1209
+
1210
def _send_blocks(
1211
self,
1212
remote_session: str,