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 4, 22:30 UTC
[CPU][BugFix] Fix inter-node pipeline parallel (#40150)
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com>
Fadi Arafeh committed
Apr 20, 2026 at 10:21 UTC
2aab9acf48b5f10a4ee0b29f152fc497a46c956d
2 files changed
+17
-2
vllm/distributed/device_communicators/cpu_communicator.py
+13
@@ -45,6 +45,9 @@ class CpuCommunicator(DeviceCommunicatorBase):
45
unique_name,
46
)
47
48
+ # send/recv tensor_dict is only supported through the SHM communicator backend
49
+ self.supports_tensor_dict = isinstance(self.dist_module, _CPUSHMDistributed)
50
+
51
if self.use_all2all:
52
if self.all2all_backend != "naive": # type: ignore[has-type]
53
logger.warning(
@@ -143,12 +146,22 @@ class CpuCommunicator(DeviceCommunicatorBase):
146
tensor_dict: dict[str, torch.Tensor | Any],
147
dst: int,
148
) -> None:
149
+ if not self.supports_tensor_dict:
150
+ raise NotImplementedError(
151
+ "CpuCommunicator does not support tensor dict fastpath with "
152
+ "torch.distributed backend."
153
+ )
154
return self.dist_module.send_tensor_dict(tensor_dict, dst)
155
156
def recv_tensor_dict(
157
self,
158
src: int,
159
) -> dict[str, torch.Tensor | Any]:
160
+ if not self.supports_tensor_dict:
161
+ raise NotImplementedError(
162
+ "CpuCommunicator does not support tensor dict fastpath with "
163
+ "torch.distributed backend."
164
+ )
165
return self.dist_module.recv_tensor_dict(src)
166
167
def dispatch_router_logits(
vllm/distributed/parallel_state.py
+4
-2
@@ -394,8 +394,10 @@ class GroupCoordinator:
394
current_platform.is_tpu() or current_platform.use_custom_op_collectives()
395
)
396
397
- self.use_cpu_custom_send_recv = current_platform.is_cpu() and hasattr(
398
- torch.ops._C, "init_shm_manager"
397
+ self.use_cpu_custom_send_recv = (
398
+ current_platform.is_cpu()
399
+ and self.device_communicator
400
+ and getattr(self.device_communicator, "supports_tensor_dict", False)
401
)
402
403
def create_mq_broadcaster(