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 10, 04:00 UTC
[Bugfix] Add CpuCommunicator.dispatch and combine to fix DP+MoE inference (#31867)
Signed-off-by: kunzh <zhikun.wu@outlook.com>
kzwrime committed
Jan 15, 2026 at 12:50 UTC
edadca109cf8b41aed3bb9cad11aaca058a6a4d0
2 files changed
+45
-1
vllm/distributed/device_communicators/base_device_communicator.py
+4
-1
@@ -286,7 +286,10 @@ class DeviceCommunicatorBase:
286
router_logits: torch.Tensor,
287
is_sequence_parallel: bool = False,
288
extra_tensors: list[torch.Tensor] | None = None,
289
- ) -> tuple[torch.Tensor, torch.Tensor]:
289
+ ) -> (
290
+ tuple[torch.Tensor, torch.Tensor]
291
+ | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
292
+ ):
293
"""
294
Dispatch the hidden states and router logits to the appropriate device.
295
This is a no-op in the base class.
vllm/distributed/device_communicators/cpu_communicator.py
+41
@@ -8,11 +8,14 @@ import torch
8
from torch.distributed import ProcessGroup
9
10
from vllm.distributed.utils import pickle
11
+from vllm.logger import init_logger
12
from vllm.platforms import current_platform
13
from vllm.platforms.interface import CpuArchEnum
14
15
from .base_device_communicator import DeviceCommunicatorBase
16
17
+logger = init_logger(__name__)
18
+
19
20
class CpuCommunicator(DeviceCommunicatorBase):
21
def __init__(
@@ -32,6 +35,20 @@ class CpuCommunicator(DeviceCommunicatorBase):
35
):
36
self.dist_module = _CPUSHMDistributed(self)
37
38
+ if self.use_all2all:
39
+ if self.all2all_backend != "naive": # type: ignore[has-type]
40
+ logger.warning(
41
+ "`%s` all2all manager is not supported on CPU. "
42
+ "Falling back to `naive` all2all manager for CPU.",
43
+ self.all2all_backend, # type: ignore[has-type]
44
+ )
45
+ self.all2all_backend = "naive"
46
+ if self.all2all_backend == "naive":
47
+ from .all2all import NaiveAll2AllManager
48
+
49
+ self.all2all_manager = NaiveAll2AllManager(self.cpu_group)
50
+ logger.info("Using naive all2all manager.")
51
+
52
def all_reduce(self, input_):
53
self.dist_module.all_reduce(input_, group=self.device_group)
54
return input_
@@ -110,6 +127,30 @@ class CpuCommunicator(DeviceCommunicatorBase):
127
) -> dict[str, torch.Tensor | Any]:
128
return self.dist_module.recv_tensor_dict(src)
129
130
+ def dispatch( # type: ignore[override]
131
+ self,
132
+ hidden_states: torch.Tensor,
133
+ router_logits: torch.Tensor,
134
+ is_sequence_parallel: bool = False,
135
+ extra_tensors: list[torch.Tensor] | None = None,
136
+ ) -> tuple[torch.Tensor, torch.Tensor]:
137
+ assert self.all2all_manager is not None
138
+ return self.all2all_manager.dispatch(
139
+ hidden_states,
140
+ router_logits,
141
+ is_sequence_parallel,
142
+ extra_tensors, # type: ignore[call-arg]
143
+ )
144
+
145
+ def combine(
146
+ self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
147
+ ) -> torch.Tensor:
148
+ assert self.all2all_manager is not None
149
+ hidden_states = self.all2all_manager.combine(
150
+ hidden_states, is_sequence_parallel
151
+ )
152
+ return hidden_states
153
+
154
155
class _CPUSHMDistributed:
156
def __init__(self, communicator: CpuCommunicator):