From 398a596ed249c14b76806900894304a62d653603 Mon Sep 17 00:00:00 2001 From: weiguihua2 Date: Tue, 9 Dec 2025 01:33:48 +0800 Subject: [PATCH] [MP executor] fix get device count for multi node of mp executor feature (#30042) Signed-off-by: weiguihua2 --- vllm/distributed/device_communicators/shm_broadcast.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index 052df19e34d72..114516ff07a1f 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -27,6 +27,7 @@ from zmq import ( # type: ignore import vllm.envs as envs from vllm.distributed.utils import StatelessProcessGroup, sched_yield from vllm.logger import init_logger +from vllm.platforms import current_platform from vllm.utils.network_utils import ( get_ip, get_open_port, @@ -632,7 +633,7 @@ class MessageQueue: The MessageQueue instance for the calling process, and a list of handles (only non-empty for the reader process). """ - local_size = torch.cuda.device_count() + local_size = current_platform.device_count() rank = dist.get_rank() same_node = rank // local_size == reader_rank // local_size buffer_io = MessageQueue(