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, 02:15 UTC
[misc]refactor `Platform.set_device` method (#20262)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Kunshang Ji committed
Jul 9, 2025 at 09:39 UTC
0b407479ef0b744d5ab90b6ad60e6f2a53e54d80
9 files changed
+43
-4
vllm/platforms/cpu.py
+7
@@ -75,6 +75,13 @@ class CpuPlatform(Platform):
75
def get_device_total_memory(cls, device_id: int = 0) -> int:
76
return psutil.virtual_memory().total
77
78
+ @classmethod
79
+ def set_device(cls, device: torch.device) -> None:
80
+ """
81
+ Set the device for the current platform.
82
+ """
83
+ torch.cpu.set_device(device)
84
+
85
@classmethod
86
def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
87
return False
vllm/platforms/cuda.py
+1
-1
@@ -77,7 +77,7 @@ class CudaPlatformBase(Platform):
77
"""
78
Set the device for the current platform.
79
"""
80
- super().set_device(device)
80
+ torch.cuda.set_device(device)
81
# With this trick we can force the device to be set eagerly
82
# see https://github.com/pytorch/pytorch/issues/155668
83
# for why and when it is needed
vllm/platforms/hpu.py
+7
@@ -45,6 +45,13 @@ class HpuPlatform(Platform):
45
def inference_mode(cls):
46
return torch.no_grad()
47
48
+ @classmethod
49
+ def set_device(cls, device: torch.device) -> None:
50
+ """
51
+ Set the device for the current platform.
52
+ """
53
+ torch.hpu.set_device(device)
54
+
55
@classmethod
56
def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
57
vllm/platforms/interface.py
+1
-1
@@ -305,7 +305,7 @@ class Platform:
305
"""
306
Set the device for the current platform.
307
"""
308
- torch.cuda.set_device(device)
308
+ raise NotImplementedError
309
310
@classmethod
311
def pre_register_and_update(cls,
vllm/platforms/rocm.py
+11
@@ -241,6 +241,17 @@ class RocmPlatform(Platform):
241
logger.info("Using ROCmFlashAttention backend.")
242
return "vllm.attention.backends.rocm_flash_attn.ROCmFlashAttentionBackend" # noqa: E501
243
244
+ @classmethod
245
+ def set_device(cls, device: torch.device) -> None:
246
+ """
247
+ Set the device for the current platform.
248
+ """
249
+ torch.cuda.set_device(device)
250
+ # With this trick we can force the device to be set eagerly
251
+ # see https://github.com/pytorch/pytorch/issues/155668
252
+ # for why and when it is needed
253
+ _ = torch.zeros(1, device=device)
254
+
255
@classmethod
256
@lru_cache(maxsize=8)
257
def get_device_capability(cls,
vllm/platforms/tpu.py
+7
@@ -55,6 +55,13 @@ class TpuPlatform(Platform):
55
logger.info("Using Pallas V1 backend.")
56
return "vllm.v1.attention.backends.pallas.PallasAttentionBackend"
57
58
+ @classmethod
59
+ def set_device(cls, device: torch.device) -> None:
60
+ """
61
+ Set the device for the current platform.
62
+ """
63
+ torch.tpu.set_device(device)
64
+
65
@classmethod
66
def get_device_name(cls, device_id: int = 0) -> str:
67
chip_type, _ = device.get_local_chips()
vllm/platforms/xpu.py
+7
@@ -45,6 +45,13 @@ class XPUPlatform(Platform):
45
logger.info("Using Flash Attention backend on V1 engine.")
46
return "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend"
47
48
+ @classmethod
49
+ def set_device(cls, device: torch.device) -> None:
50
+ """
51
+ Set the device for the current platform.
52
+ """
53
+ torch.xpu.set_device(device)
54
+
55
@classmethod
56
def get_device_capability(
57
cls,
vllm/v1/worker/gpu_worker.py
+1
-1
@@ -130,7 +130,7 @@ class Worker(WorkerBase):
130
# This env var set by Ray causes exceptions with graph building.
131
os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
132
self.device = torch.device(f"cuda:{self.local_rank}")
133
- torch.cuda.set_device(self.device)
133
+ current_platform.set_device(self.device)
134
135
_check_if_gpu_supports_dtype(self.model_config.dtype)
136
gc.collect()
vllm/v1/worker/xpu_worker.py
+1
-1
@@ -132,7 +132,7 @@ class XPUWorker(Worker):
132
if self.device_config.device.type == "xpu" and current_platform.is_xpu(
133
):
134
self.device = torch.device(f"xpu:{self.local_rank}")
135
- torch.xpu.set_device(self.device)
135
+ current_platform.set_device(self.device)
136
torch.xpu.empty_cache()
137
self.init_gpu_memory = torch.xpu.get_device_properties(
138
self.local_rank).total_memory