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 3, 21:30 UTC
[XPU] update dp rank w/o env-var isolation (#39856)
Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
liuzhenwei committed
May 12, 2026 at 14:49 UTC
8517cdaf900b13ce81ecfe2de113469ed89e8083
2 files changed
+60
-17
vllm/v1/engine/utils.py
+6
-5
@@ -155,15 +155,16 @@ class CoreEngineProcManager:
155
156
try:
157
for proc, local_dp_rank in zip(self.processes, local_dp_ranks):
158
- # Adjust device control in DP for non-CUDA platforms
159
- # as well as external and ray launchers
160
- # For CUDA platforms, we use torch.accelerator.set_device_index()()
158
+ # Adjust device control in DP for platforms that cannot rely
159
+ # on torch.accelerator.set_device_index(), and for Ray launchers.
160
device_control_context: contextlib.AbstractContextManager[None] = (
161
contextlib.nullcontext()
162
)
163
+ needs_device_env_isolation = not (
164
+ current_platform.is_cuda_alike() or current_platform.is_xpu()
165
+ )
166
if is_dp and (
165
- not current_platform.is_cuda_alike()
166
- or vllm_config.parallel_config.use_ray
167
+ needs_device_env_isolation or vllm_config.parallel_config.use_ray
168
):
169
device_control_context = set_device_control_env_var(
170
vllm_config, local_dp_rank
vllm/v1/worker/xpu_worker.py
+54
-12
@@ -2,7 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import gc
4
import os
5
-from typing import Any
5
6
import torch
7
@@ -40,19 +39,35 @@ class XPUWorker(Worker):
39
assert device_config.device_type == "xpu"
40
assert current_platform.is_xpu()
41
43
- # Torch profiler. Enabled and configured through profiler_config.
44
- self.profiler: Any | None = None
45
- profiler_config = vllm_config.profiler_config
46
- if profiler_config.profiler == "torch":
47
- worker_name = f"{vllm_config.instance_id}-rank-{self.rank}"
48
- self.profiler = TorchProfilerWrapper(
49
- profiler_config,
50
- worker_name=worker_name,
51
- local_rank=self.local_rank,
52
- activities=["CPU", "XPU"],
42
+ def init_device(self):
43
+ # In DP mode, XPU workers see all visible devices.
44
+ # Offset local_rank by the local DP shard.
45
+ parallel_config = self.parallel_config
46
+ if (
47
+ parallel_config.distributed_executor_backend
48
+ not in ("ray", "external_launcher")
49
+ and parallel_config.data_parallel_backend != "ray"
50
+ and parallel_config.nnodes_within_dp == 1
51
+ ):
52
+ dp_local_rank = parallel_config.data_parallel_rank_local
53
+ if dp_local_rank is None:
54
+ dp_local_rank = parallel_config.data_parallel_index
55
+ tp_pp_world_size = (
56
+ parallel_config.pipeline_parallel_size
57
+ * parallel_config.tensor_parallel_size
58
+ )
59
+ self.local_rank += dp_local_rank * tp_pp_world_size
60
+
61
+ visible_device_count = torch.accelerator.device_count()
62
+ assert self.local_rank < visible_device_count, (
63
+ f"DP adjusted local rank {self.local_rank} is out of bounds. "
64
+ )
65
+ assert parallel_config.local_world_size <= visible_device_count, (
66
+ f"local_world_size ({parallel_config.local_world_size}) must "
67
+ f"be less than or equal to the number of visible devices "
68
+ f"({visible_device_count})."
69
)
70
55
- def init_device(self):
71
device = self.device_config.device
72
if (
73
isinstance(device, torch.device)
@@ -117,3 +132,30 @@ class XPUWorker(Worker):
132
if self.rank == 0:
133
# If usage stat is enabled, collect relevant info.
134
report_usage_stats(self.vllm_config)
135
+
136
+ def profile(self, is_start: bool = True, profile_prefix: str | None = None):
137
+ if self.profiler_config is None or self.profiler_config.profiler is None:
138
+ raise RuntimeError(
139
+ "Profiling is not enabled. Please set --profiler-config to enable "
140
+ "profiling. Example: "
141
+ "'--profiler-config.profiler=torch --profiler-config.torch_profiler_dir"
142
+ "=YOUR_DIR_PATH_TO_DUMP_TRACE'"
143
+ )
144
+
145
+ if is_start and self.profiler is None:
146
+ from vllm.distributed.utils import get_worker_rank_suffix
147
+
148
+ rank_suffix = get_worker_rank_suffix(global_rank=self.rank)
149
+ trace_name = (
150
+ f"{profile_prefix}_{rank_suffix}" if profile_prefix else rank_suffix
151
+ )
152
+
153
+ self.profiler = TorchProfilerWrapper(
154
+ self.profiler_config,
155
+ worker_name=trace_name,
156
+ local_rank=self.local_rank,
157
+ activities=["CPU", "XPU"],
158
+ )
159
+ logger.debug("Starting torch profiler with trace name: %s", trace_name)
160
+
161
+ super().profile(is_start=is_start, profile_prefix=profile_prefix)