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 15, 11:00 UTC
[Core] Use platform-agnostic device control for DP engine core (#17245)
Signed-off-by: Jade Zheng <zheng.shoujian@outlook.com>
Jade Zheng committed
May 13, 2025 at 03:09 UTC
289199feb6616a27d97d15cde470772e1585bb84
4 files changed
+30
-39
vllm/platforms/cuda.py
+4
-22
@@ -34,24 +34,6 @@ pynvml = import_pynvml()
34
torch.backends.cuda.enable_cudnn_sdp(False)
35
36
37
-def device_id_to_physical_device_id(device_id: int) -> int:
38
- if "CUDA_VISIBLE_DEVICES" in os.environ:
39
- device_ids = os.environ["CUDA_VISIBLE_DEVICES"].split(",")
40
- if device_ids == [""]:
41
- msg = (
42
- "CUDA_VISIBLE_DEVICES is set to empty string, which means"
43
- " GPU support is disabled. If you are using ray, please unset"
44
- " the environment variable `CUDA_VISIBLE_DEVICES` inside the"
45
- " worker/actor. "
46
- "Check https://github.com/vllm-project/vllm/issues/8402 for"
47
- " more information.")
48
- raise RuntimeError(msg)
49
- physical_device_id = device_ids[device_id]
50
- return int(physical_device_id)
51
- else:
52
- return device_id
53
-
54
-
37
def with_nvml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]:
38
39
@wraps(fn)
@@ -338,7 +320,7 @@ class NvmlCudaPlatform(CudaPlatformBase):
320
device_id: int = 0
321
) -> Optional[DeviceCapability]:
322
try:
341
- physical_device_id = device_id_to_physical_device_id(device_id)
323
+ physical_device_id = cls.device_id_to_physical_device_id(device_id)
324
handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
325
major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle)
326
return DeviceCapability(major=major, minor=minor)
@@ -360,20 +342,20 @@ class NvmlCudaPlatform(CudaPlatformBase):
342
@classmethod
343
@with_nvml_context
344
def get_device_name(cls, device_id: int = 0) -> str:
363
- physical_device_id = device_id_to_physical_device_id(device_id)
345
+ physical_device_id = cls.device_id_to_physical_device_id(device_id)
346
return cls._get_physical_device_name(physical_device_id)
347
348
@classmethod
349
@with_nvml_context
350
def get_device_uuid(cls, device_id: int = 0) -> str:
369
- physical_device_id = device_id_to_physical_device_id(device_id)
351
+ physical_device_id = cls.device_id_to_physical_device_id(device_id)
352
handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
353
return pynvml.nvmlDeviceGetUUID(handle)
354
355
@classmethod
356
@with_nvml_context
357
def get_device_total_memory(cls, device_id: int = 0) -> int:
376
- physical_device_id = device_id_to_physical_device_id(device_id)
358
+ physical_device_id = cls.device_id_to_physical_device_id(device_id)
359
handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
360
return int(pynvml.nvmlDeviceGetMemoryInfo(handle).total)
361
vllm/platforms/interface.py
+19
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
import enum
3
+import os
4
import platform
5
import random
6
from platform import uname
@@ -161,6 +162,24 @@ class Platform:
162
def is_sleep_mode_available(self) -> bool:
163
return self._enum == PlatformEnum.CUDA
164
165
+ @classmethod
166
+ def device_id_to_physical_device_id(cls, device_id: int):
167
+ if cls.device_control_env_var in os.environ:
168
+ device_ids = os.environ[cls.device_control_env_var].split(",")
169
+ if device_ids == [""]:
170
+ msg = (f"{cls.device_control_env_var} is set to empty string, "
171
+ "which means current platform support is disabled. If "
172
+ "you are using ray, please unset the environment "
173
+ f"variable `{cls.device_control_env_var}` inside the "
174
+ "worker/actor. Check "
175
+ "https://github.com/vllm-project/vllm/issues/8402 for "
176
+ "more information.")
177
+ raise RuntimeError(msg)
178
+ physical_device_id = device_ids[device_id]
179
+ return int(physical_device_id)
180
+ else:
181
+ return device_id
182
+
183
@classmethod
184
def get_attn_backend_cls(cls, selected_backend: _Backend, head_size: int,
185
dtype: torch.dtype, kv_cache_dtype: Optional[str],
vllm/platforms/rocm.py
+1
-10
@@ -95,15 +95,6 @@ def with_amdsmi_context(fn):
95
return wrapper
96
97
98
-def device_id_to_physical_device_id(device_id: int) -> int:
99
- if "CUDA_VISIBLE_DEVICES" in os.environ:
100
- device_ids = os.environ["CUDA_VISIBLE_DEVICES"].split(",")
101
- physical_device_id = device_ids[device_id]
102
- return int(physical_device_id)
103
- else:
104
- return device_id
105
-
106
-
98
@cache
99
def on_mi250_mi300() -> bool:
100
GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
@@ -238,7 +229,7 @@ class RocmPlatform(Platform):
229
@with_amdsmi_context
230
@lru_cache(maxsize=8)
231
def get_device_name(cls, device_id: int = 0) -> str:
241
- physical_device_id = device_id_to_physical_device_id(device_id)
232
+ physical_device_id = cls.device_id_to_physical_device_id(device_id)
233
handle = amdsmi_get_processor_handles()[physical_device_id]
234
asic_info = amdsmi_get_gpu_asic_info(handle)
235
device_name: str = asic_info["device_id"]
vllm/v1/engine/core.py
+6
-7
@@ -622,13 +622,12 @@ class DPEngineCoreProc(EngineCoreProc):
622
assert 0 <= local_dp_rank <= dp_rank < dp_size
623
624
from vllm.platforms import current_platform
625
- if current_platform.is_cuda_alike():
626
- from vllm.platforms.cuda import device_id_to_physical_device_id
627
- tp_size = vllm_config.parallel_config.tensor_parallel_size
628
- os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(
629
- str(device_id_to_physical_device_id(i))
630
- for i in range(local_dp_rank * tp_size, (local_dp_rank + 1) *
631
- tp_size))
625
+ device_control_env_var = current_platform.device_control_env_var
626
+ tp_size = vllm_config.parallel_config.tensor_parallel_size
627
+ os.environ[device_control_env_var] = ",".join(
628
+ str(current_platform.device_id_to_physical_device_id(i))
629
+ for i in range(local_dp_rank * tp_size, (local_dp_rank + 1) *
630
+ tp_size))
631
632
self.local_dp_rank = local_dp_rank
633
self.dp_group = vllm_config.parallel_config.stateless_init_dp_group()