From 8dfaffb700a40fb6fa900dfb955a1c79a4f2b210 Mon Sep 17 00:00:00 2001 From: MOVZX Date: Mon, 19 Oct 2026 00:33:19 +0700 Subject: [PATCH] [Bugfix] Restore xpu_communicator for CUDA graph capture parallel_state.graph_capture() imports xpu_communicator on CUDA too. The SM86 cut removed the file and capture crashed at startup. Co-authored-by: Hinata Signed-off-by: MOVZX --- .../device_communicators/xpu_communicator.py | 296 ++++++++++++++++++ .../ec_connector/mooncake_ec_connector.py | 130 -------- 2 files changed, 296 insertions(+), 130 deletions(-) create mode 100644 vllm/distributed/device_communicators/xpu_communicator.py delete mode 100644 vllm/distributed/ec_transfer/ec_connector/mooncake_ec_connector.py diff --git a/vllm/distributed/device_communicators/xpu_communicator.py b/vllm/distributed/device_communicators/xpu_communicator.py new file mode 100644 index 0000000000..7a5913754f --- /dev/null +++ b/vllm/distributed/device_communicators/xpu_communicator.py @@ -0,0 +1,296 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup + +import vllm.envs as envs +from vllm.logger import init_logger + +from .base_device_communicator import DeviceCommunicatorBase + +logger = init_logger(__name__) + + +class XpuCommunicator(DeviceCommunicatorBase): + def __init__( + self, + cpu_group: ProcessGroup, + device: torch.device | None = None, + device_group: ProcessGroup | None = None, + unique_name: str = "", + use_all2all: bool = False, + ): + super().__init__( + cpu_group, device, device_group, unique_name, use_all2all=use_all2all + ) + self.ca_comm: None = None + if self.use_all2all: + if self.all2all_backend in ("naive", "allgather_reducescatter"): + from .all2all import AgRsAll2AllManager + + self.all2all_manager = AgRsAll2AllManager(self.cpu_group) + logger.info("Using AgRs manager on XPU device.") + + else: # type: ignore[has-type] + logger.warning( + "`%s` all2all manager is not supported on XPU. " + "Falling back to AgRs manager for XPU, " + "which is the Default backend", + self.all2all_backend, # type: ignore[has-type] + ) + from .all2all import AgRsAll2AllManager + + self.all2all_manager = AgRsAll2AllManager(self.cpu_group) + logger.info("Using AgRs manager on XPU device.") + + def _fixed_rank_sum(self, input_: torch.Tensor) -> torch.Tensor: + flat_input = input_.reshape(-1) + gathered = torch.empty( + (self.world_size, flat_input.numel()), + dtype=input_.dtype, + device=input_.device, + ) + dist.all_gather_single(gathered.view(-1), flat_input, group=self.device_group) + output = gathered[0].clone() + for rank in range(1, self.world_size): + output.add_(gathered[rank]) + return output.view(input_.shape) + + def all_reduce(self, input_: torch.Tensor) -> torch.Tensor: + if envs.VLLM_BATCH_INVARIANT: + # XCCL changes its reduction order with message size. Gather first + # and accumulate in rank order, independent of the token batch. + return self._fixed_rank_sum(input_) + + output = input_.clone() + dist.all_reduce(output, group=self.device_group) + return output + + def reduce_scatter(self, input_: torch.Tensor, dim: int = -1): + world_size = self.world_size + + if dim < 0: + # Convert negative dim to positive. + dim += input_.dim() + + # Note: This will produce an incorrect answer if we don't make + # the input_tensor contiguous. Possible bug in reduce_scatter_tensor? + input_tensor = input_.movedim(0, dim).contiguous() + + assert input_tensor.shape[0] % world_size == 0 + chunk_size = input_tensor.shape[0] // world_size + output_shape = (chunk_size,) + input_tensor.shape[1:] + + output = torch.empty( + output_shape, dtype=input_tensor.dtype, device=input_tensor.device + ) + + if envs.VLLM_BATCH_INVARIANT: + reduced = self._fixed_rank_sum(input_tensor) + start = self.rank_in_group * chunk_size + output.copy_(reduced.narrow(0, start, chunk_size)) + else: + dist.reduce_scatter_tensor(output, input_tensor, group=self.device_group) + + # Reshape before returning + return output.movedim(0, dim).contiguous() + + def reduce_scatterv( + self, input_: torch.Tensor, dim: int = -1, sizes: list[int] | None = None + ): + world_size = self.world_size + + if dim < 0: + # Convert negative dim to positive. + dim += input_.dim() + + # Note: This will produce an incorrect answer if we don't make + # the input_tensor contiguous. Possible bug in reduce_scatter_tensor? + input_tensor = input_.movedim(0, dim).contiguous() + + if sizes is not None: + assert len(sizes) == world_size + assert input_tensor.shape[0] == sum(sizes) + chunk_size = sizes[self.rank_in_group] + else: + assert input_tensor.shape[0] % world_size == 0 + chunk_size = input_tensor.shape[0] // world_size + output_shape = (chunk_size,) + input_tensor.shape[1:] + + output = torch.empty( + output_shape, dtype=input_tensor.dtype, device=input_tensor.device + ) + if envs.VLLM_BATCH_INVARIANT: + reduced = self._fixed_rank_sum(input_tensor) + start = ( + self.rank_in_group * chunk_size + if sizes is None + else sum(sizes[: self.rank_in_group]) + ) + output.copy_(reduced.narrow(0, start, chunk_size)) + elif sizes is None or sizes.count(sizes[0]) == len(sizes): + dist.reduce_scatter_tensor(output, input_tensor, group=self.device_group) + elif torch.xpu.device_count() < self.world_size: + # uneven reduce_scatter is not supported when the + # process cannot see every device in the group ( + # e.g. external-LBDP). + max_size = max(sizes) + padded = input_tensor.new_zeros( + (world_size * max_size,) + input_tensor.shape[1:] + ) + for rank, split in enumerate(input_tensor.split(sizes, dim=0)): + padded[rank * max_size : rank * max_size + sizes[rank]] = split + padded_output = input_tensor.new_empty((max_size,) + input_tensor.shape[1:]) + dist.reduce_scatter_tensor(padded_output, padded, group=self.device_group) + output = padded_output[:chunk_size] + else: + input_splits = list(input_tensor.split(sizes, dim=0)) + dist.reduce_scatter(output, input_splits, group=self.device_group) + # Reshape before returning + return output.movedim(0, dim).contiguous() + + def all_gatherv( + self, + input_: torch.Tensor | list[torch.Tensor], + dim: int = 0, + sizes: list[int] | None = None, + ): + if dim != 0: + raise NotImplementedError("only dim 0 all-gatherv is supported") + world_size = self.world_size + + # 'sizes' is not needed if all inputs in the same group have the same + # shape + if sizes is not None and all(s == sizes[0] for s in sizes): + sizes = None + + def _all_gather_single(input_: torch.Tensor, sizes: list[int] | None = None): + input_size = input_.size() + if sizes is not None: + assert len(sizes) == world_size + assert input_.shape[dim] == sizes[self.rank_in_group], ( + f"{input_.shape[dim]} != {sizes[self.rank_in_group]}" + ) + output_size = (sum(sizes),) + input_size[1:] + else: + output_size = (input_size[0] * world_size,) + input_size[1:] + # Allocate output tensor. + output_tensor = torch.empty( + output_size, dtype=input_.dtype, device=input_.device + ) + + if sizes is not None: + all_gather_list = [] + for size in sizes: + all_gather_list.append( + torch.empty( + (size,) + input_.shape[1:], + dtype=input_.dtype, + device=input_.device, + ) + ) + dist.all_gather(all_gather_list, input_, group=self.device_group) + output_tensor = torch.cat(all_gather_list, dim=0) + else: + dist.all_gather([output_tensor], input_, group=self.device_group) + return output_tensor + + if isinstance(input_, torch.Tensor): + return _all_gather_single(input_, sizes) + + output_list = [] + for inp in input_: + output_list.append(_all_gather_single(inp, sizes=sizes)) + return output_list + + def gather( + self, input_: torch.Tensor, dst: int = 0, dim: int = -1 + ) -> torch.Tensor | None: + assert -input_.dim() <= dim < input_.dim(), ( + f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + ) + if dim < 0: + # Convert negative dim to positive. + dim += input_.dim() + # For xpu path, gather doesn't work properly together with ray + # cluster so we use all_gather instead for now. + input_size = input_.size() + # Allocate output tensor. + output_tensor = torch.empty( + (self.world_size,) + input_size, dtype=input_.dtype, device=input_.device + ) + # All-gather. + dist.all_gather_single(output_tensor, input_, group=self.device_group) + if self.rank_in_group == dst: + # Reshape + output_tensor = output_tensor.movedim(0, dim) + output_tensor = output_tensor.reshape( + input_size[:dim] + + (self.world_size * input_size[dim],) + + input_size[dim + 1 :] + ) + else: + output_tensor = None + return output_tensor + + def broadcast(self, input_: torch.Tensor, src: int = 0) -> None: + dist.broadcast(input_, src=src, group=self.device_group) + + def dispatch_router_logits( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + is_sequence_parallel: bool = False, + extra_tensors: list[torch.Tensor] | None = None, + ) -> ( + tuple[torch.Tensor, torch.Tensor] + | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]] + ): + """Dispatch the hidden states and router logits to the appropriate device. + This is a no-op in the base class. + """ + assert self.all2all_manager is not None + return self.all2all_manager.dispatch_router_logits( + hidden_states, + router_logits, + is_sequence_parallel, + extra_tensors, + ) + + def dispatch( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + is_sequence_parallel: bool = False, + extra_tensors: list[torch.Tensor] | None = None, + ) -> ( + tuple[torch.Tensor, torch.Tensor, torch.Tensor] + | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]] + ): + """Dispatch the hidden states and topk weights/ids to the appropriate device. + This is a no-op in the base class. + """ + assert self.all2all_manager is not None + return self.all2all_manager.dispatch( + hidden_states, + topk_weights, + topk_ids, + is_sequence_parallel, + extra_tensors=extra_tensors, + ) + + def combine( + self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False + ) -> torch.Tensor: + """Combine the hidden states and router logits from the appropriate device. + This is a no-op in the base class. + """ + assert self.all2all_manager is not None + return self.all2all_manager.combine( + hidden_states, + is_sequence_parallel, + ) diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_ec_connector.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_ec_connector.py deleted file mode 100644 index 2949742220..0000000000 --- a/vllm/distributed/ec_transfer/ec_connector/mooncake_ec_connector.py +++ /dev/null @@ -1,130 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Encoder-cache connector backed by Mooncake TransferEngine.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Any - -from vllm.distributed.ec_transfer.ec_connector.base import ( - ECConnectorBase, - ECConnectorMetadata, - ECConnectorRole, - ECConnectorWorkerMetadata, -) -from vllm.distributed.ec_transfer.ec_connector.mooncake.metadata import ( - ECMooncakeConnectorMetadata, -) -from vllm.distributed.ec_transfer.ec_connector.mooncake.scheduler import ( - ECMooncakeScheduler, -) -from vllm.distributed.ec_transfer.ec_connector.mooncake.worker import ECMooncakeWorker - -if TYPE_CHECKING: - import torch - - from vllm.config import VllmConfig - from vllm.v1.core.sched.output import SchedulerOutput - from vllm.v1.outputs import ECConnectorOutput - from vllm.v1.request import Request - - -class ECMooncakeConnector(ECConnectorBase): - """Preserve the public API while delegating to one process-role component.""" - - def __init__(self, vllm_config: VllmConfig, role: ECConnectorRole): - super().__init__(vllm_config=vllm_config, role=role) - - self._scheduler: ECMooncakeScheduler | None = None - self._worker: ECMooncakeWorker | None = None - self._closed = False - - if role == ECConnectorRole.SCHEDULER: - self._scheduler = ECMooncakeScheduler(vllm_config) - elif role == ECConnectorRole.WORKER: - self._worker = ECMooncakeWorker(vllm_config) - else: - raise ValueError(f"Unknown EC connector role: {role}") - - def start_worker_services(self) -> None: - assert self._worker is not None - self._worker.start_services() - - def start_save_caches(self, **kwargs: Any) -> None: - assert self._worker is not None - metadata = self._get_connector_metadata() - assert isinstance(metadata, ECMooncakeConnectorMetadata) - self._worker.start_save_caches(metadata, **kwargs) - - def start_load_caches( - self, encoder_cache: dict[str, torch.Tensor], **kwargs: Any - ) -> None: - assert self._worker is not None - metadata = self._get_connector_metadata() - assert isinstance(metadata, ECMooncakeConnectorMetadata) - self._worker.start_load_caches(metadata, encoder_cache, **kwargs) - - def save_caches( - self, encoder_cache: dict[str, torch.Tensor], mm_hash: str, **kwargs: Any - ) -> None: - assert self._worker is not None - self._worker.save_caches(encoder_cache, mm_hash, **kwargs) - - def get_finished( - self, finished_req_ids: set[str] - ) -> tuple[set[str] | None, set[str] | None]: - assert self._worker is not None - return self._worker.get_finished(finished_req_ids) - - def build_connector_worker_meta(self) -> ECConnectorWorkerMetadata | None: - assert self._worker is not None - return self._worker.build_connector_worker_meta() - - def take_unavailable_requests(self) -> set[str]: - assert self._scheduler is not None - return self._scheduler.take_unavailable_requests() - - def has_cache_item(self, identifier: str) -> bool: - assert self._scheduler is not None - return self._scheduler.has_cache_item(identifier) - - def ensure_cache_available( - self, request: Request, num_computed_tokens: int - ) -> bool: - assert self._scheduler is not None - return self._scheduler.ensure_cache_available(request, num_computed_tokens) - - def update_state_after_alloc(self, request: Request, index: int) -> None: - assert self._scheduler is not None - self._scheduler.update_state_after_alloc(request, index) - - def update_state_after_free(self, request: Request, index: int) -> None: - assert self._scheduler is not None - self._scheduler.update_state_after_free(request, index) - - def build_connector_meta( - self, scheduler_output: SchedulerOutput - ) -> ECConnectorMetadata: - assert self._scheduler is not None - return self._scheduler.build_connector_meta(scheduler_output) - - def update_connector_output(self, connector_output: ECConnectorOutput) -> None: - assert self._scheduler is not None - self._scheduler.update_connector_output(connector_output) - - def has_pending_push_work(self) -> bool: - assert self._scheduler is not None - return self._scheduler.has_pending_push_work() - - def request_finished(self, request: Request) -> tuple[bool, dict[str, Any] | None]: - assert self._scheduler is not None - return self._scheduler.request_finished(request) - - def shutdown(self) -> None: - if self._closed: - return - self._closed = True - if self._scheduler is not None: - self._scheduler.close() - if self._worker is not None: - self._worker.close()