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 14, 05:30 UTC
[XPU] Fix OOM issue for data parallel with Ray backend (#22500)
Signed-off-by: Fanli Lin <fanli.lin@intel.com> Signed-off-by: Fanli Lin <fanli0116@gmail.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
Fanli Lin committed
Aug 27, 2025 at 19:57 UTC
8f0d7eaea87409a54ccaed76995b59c6b0a3d4cf
2 files changed
+49
-13
vllm/v1/engine/core.py
+18
-9
@@ -39,7 +39,8 @@ from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest,
39
EngineCoreRequestType,
40
ReconfigureDistributedRequest, ReconfigureRankType,
41
UtilityOutput, UtilityResult)
42
-from vllm.v1.engine.utils import EngineHandshakeMetadata, EngineZmqAddresses
42
+from vllm.v1.engine.utils import (EngineHandshakeMetadata, EngineZmqAddresses,
43
+ get_device_indices)
44
from vllm.v1.executor.abstract import Executor
45
from vllm.v1.kv_cache_interface import KVCacheConfig
46
from vllm.v1.metrics.stats import SchedulerStats
@@ -1169,22 +1170,30 @@ class DPEngineCoreActor(DPEngineCoreProc):
1170
# https://github.com/ray-project/ray/pull/40461/files#diff-31e8159767361e4bc259b6d9883d9c0d5e5db780fcea4a52ead4ee3ee4a59a78R1860 # noqa: E501
1171
# and get_accelerator_ids_for_accelerator_resource() in worker.py
1172
# of ray.
1172
- self._set_cuda_visible_devices(vllm_config, local_dp_rank)
1173
+ self._set_visible_devices(vllm_config, local_dp_rank)
1174
1175
super().__init__(vllm_config, local_client, "", executor_class,
1176
log_stats)
1177
1177
- def _set_cuda_visible_devices(self, vllm_config: VllmConfig,
1178
- local_dp_rank: int):
1178
+ def _set_visible_devices(self, vllm_config: VllmConfig,
1179
+ local_dp_rank: int):
1180
from vllm.platforms import current_platform
1180
- device_control_env_var = current_platform.device_control_env_var
1181
+ if current_platform.is_xpu():
1182
+ pass
1183
+ else:
1184
+ device_control_env_var = current_platform.device_control_env_var
1185
+ self._set_cuda_visible_devices(vllm_config, local_dp_rank,
1186
+ device_control_env_var)
1187
+
1188
+ def _set_cuda_visible_devices(self, vllm_config: VllmConfig,
1189
+ local_dp_rank: int,
1190
+ device_control_env_var: str):
1191
world_size = vllm_config.parallel_config.world_size
1192
# Set CUDA_VISIBLE_DEVICES or equivalent.
1193
try:
1184
- os.environ[device_control_env_var] = ",".join(
1185
- str(current_platform.device_id_to_physical_device_id(i))
1186
- for i in range(local_dp_rank *
1187
- world_size, (local_dp_rank + 1) * world_size))
1194
+ value = get_device_indices(device_control_env_var, local_dp_rank,
1195
+ world_size)
1196
+ os.environ[device_control_env_var] = value
1197
except IndexError as e:
1198
raise Exception(
1199
f"Error setting {device_control_env_var}: "
vllm/v1/engine/utils.py
+31
-4
@@ -164,19 +164,33 @@ def set_device_control_env_var(vllm_config: VllmConfig,
164
"""
165
world_size = vllm_config.parallel_config.world_size
166
evar = current_platform.device_control_env_var
167
+
168
+ value = get_device_indices(evar, local_dp_rank, world_size)
169
+ with patch.dict(os.environ, values=((evar, value), )):
170
+ yield
171
+
172
+
173
+def get_device_indices(device_control_env_var: str, local_dp_rank: int,
174
+ world_size: int):
175
+ """
176
+ Returns a comma-separated string of device indices for the specified
177
+ data parallel rank.
178
+
179
+ For example, if world_size=2 and local_dp_rank=1, and there are 4 devices,
180
+ this will select devices 2 and 3 for local_dp_rank=1.
181
+ """
182
try:
183
value = ",".join(
184
str(current_platform.device_id_to_physical_device_id(i))
185
for i in range(local_dp_rank * world_size, (local_dp_rank + 1) *
186
world_size))
187
except IndexError as e:
173
- raise Exception(f"Error setting {evar}: "
188
+ raise Exception(f"Error setting {device_control_env_var}: "
189
f"local range: [{local_dp_rank * world_size}, "
190
f"{(local_dp_rank + 1) * world_size}) "
191
"base value: "
177
- f"\"{os.getenv(evar)}\"") from e
178
- with patch.dict(os.environ, values=((evar, value), )):
179
- yield
192
+ f"\"{os.getenv(device_control_env_var)}\"") from e
193
+ return value
194
195
196
class CoreEngineActorManager:
@@ -254,6 +268,19 @@ class CoreEngineActorManager:
268
dp_vllm_config = copy.deepcopy(vllm_config)
269
dp_vllm_config.parallel_config.placement_group = pg
270
local_client = index < local_engine_count
271
+
272
+ # Ray XPU known issue: dpctl initializes the GPU runtime early, so
273
+ # setting device env vars in Ray actor's initialization method
274
+ # will not affect device selection. See:
275
+ # https://github.com/ray-project/ray/blob/master/python/ray/_private/accelerators/intel_gpu.py#L56 # noqa: E501
276
+ if current_platform.is_xpu():
277
+ device_evar = current_platform.device_control_env_var
278
+ device_indices = get_device_indices(device_evar, local_index,
279
+ world_size)
280
+ actor_env_vars = self.env_vars_dict.copy()
281
+ actor_env_vars[device_evar] = device_indices
282
+ runtime_env = RuntimeEnv(env_vars=actor_env_vars)
283
+
284
actor = ray.remote(DPEngineCoreActor).options(
285
scheduling_strategy=PlacementGroupSchedulingStrategy(
286
placement_group=pg,