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, 11:30 UTC
[CPU][Distributed] Fix Enable _CPUSHMDistributed only when TP/PP ranks share the same SHM group name (#34169)
Signed-off-by: Charles Ashby <charlesa.l@hotmail.com>
Charles Ashby committed
Mar 2, 2026 at 04:34 UTC
cbd361fd468c29af00a4443b4f88cc216c6dcfe7
1 file changed
+31
-6
vllm/distributed/device_communicators/cpu_communicator.py
+31
-6
@@ -35,8 +35,15 @@ class CpuCommunicator(DeviceCommunicatorBase):
35
)
36
and hasattr(torch.ops._C, "init_shm_manager")
37
and (unique_name.startswith("tp") or unique_name.startswith("pp"))
38
+ and self._all_group_ranks_share_shm_group_name()
39
):
40
self.dist_module = _CPUSHMDistributed(self)
41
+ elif unique_name.startswith("tp") or unique_name.startswith("pp"):
42
+ logger.info(
43
+ "CPU SHM communicator disabled for group %s: ranks do not share "
44
+ "the same SHM group name, falling back to torch.distributed.",
45
+ unique_name,
46
+ )
47
48
if self.use_all2all:
49
if self.all2all_backend != "naive": # type: ignore[has-type]
@@ -52,6 +59,20 @@ class CpuCommunicator(DeviceCommunicatorBase):
59
self.all2all_manager = NaiveAll2AllManager(self.cpu_group)
60
logger.info("Using naive all2all manager.")
61
62
+ def _all_group_ranks_share_shm_group_name(self) -> bool:
63
+ """
64
+ CPUSHM requires all ranks in this group to agree on one SHM group name.
65
+ This is a lightweight consistency check for VLLM_DIST_IDENT/name inputs.
66
+ """
67
+ local_name = _CPUSHMDistributed.make_group_name(self)
68
+ names: list[str] = [""] * self.world_size
69
+ torch.distributed.all_gather_object(
70
+ names,
71
+ local_name,
72
+ group=self.device_group,
73
+ )
74
+ return len(set(names)) == 1
75
+
76
def all_reduce(self, input_):
77
self.dist_module.all_reduce(input_, group=self.device_group)
78
return input_
@@ -193,17 +214,21 @@ class CpuCommunicator(DeviceCommunicatorBase):
214
215
class _CPUSHMDistributed:
216
def __init__(self, communicator: CpuCommunicator):
196
- instance_identifier = os.environ["VLLM_DIST_IDENT"]
197
- unique_name = communicator.unique_name
198
- instance_identifier = f"{instance_identifier}-{unique_name}"
217
self.communicator = communicator
218
201
- group_ranks = [str(rank) for rank in self.communicator.ranks]
202
- shm_group_identifier = f"[{'-'.join(group_ranks)}]"
203
- self.group_name = f"{instance_identifier}-{shm_group_identifier}-cpushm"
219
+ self.group_name = self.make_group_name(communicator)
220
221
self.handle = self._init_cpu_shm()
222
223
+ @staticmethod
224
+ def make_group_name(communicator: CpuCommunicator) -> str:
225
+ instance_identifier = os.environ["VLLM_DIST_IDENT"]
226
+ unique_name = communicator.unique_name
227
+ instance_identifier = f"{instance_identifier}-{unique_name}"
228
+ group_ranks = [str(rank) for rank in communicator.ranks]
229
+ shm_group_identifier = f"[{'-'.join(group_ranks)}]"
230
+ return f"{instance_identifier}-{shm_group_identifier}-cpushm"
231
+
232
def _init_cpu_shm(self) -> int:
233
thread_num_tensor = torch.tensor(
234
[torch.get_num_threads()],