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, 17:45 UTC
[Bugfix] Add missing extra_tensors arg to DeviceCommunicatorBase.disp… (#31644)
Signed-off-by: kunzh <zhikun.wu@outlook.com>
kzwrime committed
Jan 6, 2026 at 01:26 UTC
21156ff1999f180c95176d4be9b9efb16dcb8473
2 files changed
+9
-5
vllm/distributed/device_communicators/base_device_communicator.py
+1
@@ -285,6 +285,7 @@ class DeviceCommunicatorBase:
285
hidden_states: torch.Tensor,
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]:
290
"""
291
Dispatch the hidden states and router logits to the appropriate device.
vllm/distributed/device_communicators/xpu_communicator.py
+8
-5
@@ -23,11 +23,11 @@ class XpuCommunicator(DeviceCommunicatorBase):
23
):
24
super().__init__(cpu_group, device, device_group, unique_name)
25
if self.use_all2all:
26
- if self.all2all_backend != "naive":
26
+ if self.all2all_backend != "naive": # type: ignore[has-type]
27
logger.warning(
28
"`%s` all2all manager is not supported on XPU. "
29
"Falling back to `naive` all2all manager for XPU.",
30
- self.all2all_backend,
30
+ self.all2all_backend, # type: ignore[has-type]
31
)
32
self.all2all_backend = "naive"
33
if self.all2all_backend == "naive":
@@ -78,12 +78,15 @@ class XpuCommunicator(DeviceCommunicatorBase):
78
hidden_states: torch.Tensor,
79
router_logits: torch.Tensor,
80
is_sequence_parallel: bool = False,
81
+ extra_tensors: list[torch.Tensor] | None = None,
82
) -> tuple[torch.Tensor, torch.Tensor]:
83
assert self.all2all_manager is not None
83
- hidden_states, router_logits = self.all2all_manager.dispatch(
84
- hidden_states, router_logits, is_sequence_parallel
84
+ return self.all2all_manager.dispatch(
85
+ hidden_states,
86
+ router_logits,
87
+ is_sequence_parallel,
88
+ extra_tensors, # type: ignore[call-arg]
89
)
86
- return hidden_states, router_logits
90
91
def combine(
92
self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False