Add TPU to device config

2026-03-28 18:29:08 +08:00 · 2024-04-01 08:23:44 +00:00 · 2024-04-01 08:23:44 +00:00 · 38e3d33a62
commit 38e3d33a62
parent 02e614d922
1 changed files with 3 additions and 0 deletions
--- a/vllm/config.py
+++ b/vllm/config.py
@ -596,6 +596,9 @@ class DeviceConfig:
        # Some device types require processing inputs on CPU
        if self.device_type in ["neuron"]:
            self.device = torch.device("cpu")
+        elif self.device_type in ["tpu"]:
+            # Will be set by `xm.xla_device()`
+            self.device = None
        else:
            # Set device with device type
            self.device = torch.device(self.device_type)