From 9d827170a3aa586dfb458bf28d18fd279bdbf580 Mon Sep 17 00:00:00 2001
From: Mengqing Cao <cmq0113@163.com>
Date: Thu, 21 Nov 2024 12:44:20 +0800
Subject: [PATCH] [Platforms] Add `device_type` in `Platform` (#10508)

Signed-off-by: MengqingCao <cmq0113@163.com>
---
 vllm/config.py              | 17 ++---------------
 vllm/platforms/cpu.py       |  1 +
 vllm/platforms/cuda.py      |  1 +
 vllm/platforms/hpu.py       |  1 +
 vllm/platforms/interface.py |  1 +
 vllm/platforms/neuron.py    |  1 +
 vllm/platforms/openvino.py  |  1 +
 vllm/platforms/rocm.py      |  1 +
 vllm/platforms/tpu.py       |  1 +
 vllm/platforms/xpu.py       |  1 +
 10 files changed, 11 insertions(+), 15 deletions(-)

diff --git a/vllm/config.py b/vllm/config.py
index 7522486782cc9..0ed92f370cf50 100644
--- a/vllm/config.py
+++ b/vllm/config.py
@@ -1193,21 +1193,8 @@ class DeviceConfig:
     def __init__(self, device: str = "auto") -> None:
         if device == "auto":
             # Automated device type detection
-            if current_platform.is_cuda_alike():
-                self.device_type = "cuda"
-            elif current_platform.is_neuron():
-                self.device_type = "neuron"
-            elif current_platform.is_hpu():
-                self.device_type = "hpu"
-            elif current_platform.is_openvino():
-                self.device_type = "openvino"
-            elif current_platform.is_tpu():
-                self.device_type = "tpu"
-            elif current_platform.is_cpu():
-                self.device_type = "cpu"
-            elif current_platform.is_xpu():
-                self.device_type = "xpu"
-            else:
+            self.device_type = current_platform.device_type
+            if self.device_type is None:
                 raise RuntimeError("Failed to infer device type")
         else:
             # Device type is assigned explicitly
diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py
index 43cbafe709d84..0c4c916406223 100644
--- a/vllm/platforms/cpu.py
+++ b/vllm/platforms/cpu.py
@@ -19,6 +19,7 @@
 
 class CpuPlatform(Platform):
     _enum = PlatformEnum.CPU
+    device_type: str = "cpu"
 
     @classmethod
     def get_device_name(cls, device_id: int = 0) -> str:
diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py
index d2911ef650743..07562a8c3d71e 100644
--- a/vllm/platforms/cuda.py
+++ b/vllm/platforms/cuda.py
@@ -109,6 +109,7 @@ def device_id_to_physical_device_id(device_id: int) -> int:
 
 class CudaPlatform(Platform):
     _enum = PlatformEnum.CUDA
+    device_type: str = "cuda"
 
     @classmethod
     def get_device_capability(cls, device_id: int = 0) -> DeviceCapability:
diff --git a/vllm/platforms/hpu.py b/vllm/platforms/hpu.py
index 1e0888a30ba96..36d944b3f24b8 100644
--- a/vllm/platforms/hpu.py
+++ b/vllm/platforms/hpu.py
@@ -5,6 +5,7 @@
 
 class HpuPlatform(Platform):
     _enum = PlatformEnum.HPU
+    device_type: str = "hpu"
 
     @classmethod
     def get_default_attn_backend(cls, selected_backend: _Backend) -> _Backend:
diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py
index f4849fa2ccfb0..68abec28ad71e 100644
--- a/vllm/platforms/interface.py
+++ b/vllm/platforms/interface.py
@@ -56,6 +56,7 @@ def to_int(self) -> int:
 
 class Platform:
     _enum: PlatformEnum
+    device_type: str
 
     def is_cuda(self) -> bool:
         return self._enum == PlatformEnum.CUDA
diff --git a/vllm/platforms/neuron.py b/vllm/platforms/neuron.py
index 07d8398eda525..57e3c0dfae84c 100644
--- a/vllm/platforms/neuron.py
+++ b/vllm/platforms/neuron.py
@@ -3,6 +3,7 @@
 
 class NeuronPlatform(Platform):
     _enum = PlatformEnum.NEURON
+    device_type: str = "neuron"
 
     @classmethod
     def get_device_name(cls, device_id: int = 0) -> str:
diff --git a/vllm/platforms/openvino.py b/vllm/platforms/openvino.py
index ad69ced5417b3..130b8eec1b386 100644
--- a/vllm/platforms/openvino.py
+++ b/vllm/platforms/openvino.py
@@ -10,6 +10,7 @@
 
 class OpenVinoPlatform(Platform):
     _enum = PlatformEnum.OPENVINO
+    device_type: str = "openvino"
 
     @classmethod
     def get_default_attn_backend(cls, selected_backend: _Backend) -> _Backend:
diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py
index bb3a49c8b73bc..c62241d8bb47b 100644
--- a/vllm/platforms/rocm.py
+++ b/vllm/platforms/rocm.py
@@ -29,6 +29,7 @@
 
 class RocmPlatform(Platform):
     _enum = PlatformEnum.ROCM
+    device_type: str = "cuda"
 
     @classmethod
     def get_default_attn_backend(cls, selected_backend: _Backend) -> _Backend:
diff --git a/vllm/platforms/tpu.py b/vllm/platforms/tpu.py
index 2a7ca9fb8c576..863875ef5c2d6 100644
--- a/vllm/platforms/tpu.py
+++ b/vllm/platforms/tpu.py
@@ -16,6 +16,7 @@
 
 class TpuPlatform(Platform):
     _enum = PlatformEnum.TPU
+    device_type: str = "tpu"
 
     @classmethod
     def get_default_attn_backend(cls, selected_backend: _Backend) -> _Backend:
diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py
index c3c4746d3cc25..536e17a5f93e8 100644
--- a/vllm/platforms/xpu.py
+++ b/vllm/platforms/xpu.py
@@ -16,6 +16,7 @@
 
 class XPUPlatform(Platform):
     _enum = PlatformEnum.XPU
+    device_type: str = "xpu"
 
     @classmethod
     def get_default_attn_backend(cls, selected_backend: _Backend) -> _Backend: