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 11, 21:15 UTC
[MP executor] fix get device count for multi node of mp executor feature (#30042)
Signed-off-by: weiguihua2 <weiguihua2@huawei.com>
weiguihua2 committed
Dec 9, 2025 at 01:33 UTC
398a596ed249c14b76806900894304a62d653603
1 file changed
+2
-1
vllm/distributed/device_communicators/shm_broadcast.py
+2
-1
@@ -27,6 +27,7 @@ from zmq import ( # type: ignore
27
import vllm.envs as envs
28
from vllm.distributed.utils import StatelessProcessGroup, sched_yield
29
from vllm.logger import init_logger
30
+from vllm.platforms import current_platform
31
from vllm.utils.network_utils import (
32
get_ip,
33
get_open_port,
@@ -632,7 +633,7 @@ class MessageQueue:
633
The MessageQueue instance for the calling process,
634
and a list of handles (only non-empty for the reader process).
635
"""
635
- local_size = torch.cuda.device_count()
636
+ local_size = current_platform.device_count()
637
rank = dist.get_rank()
638
same_node = rank // local_size == reader_rank // local_size
639
buffer_io = MessageQueue(