| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from dataclasses import field |
| 5 | from typing import Any, Literal |
| 6 | |
| 7 | import torch |
| 8 | from pydantic import ConfigDict, SkipValidation |
| 9 | |
| 10 | from vllm.config.utils import config |
| 11 | from vllm.utils.hashing import safe_hash |
| 12 | |
| 13 | Device = Literal["auto", "cuda", "cpu", "tpu", "xpu"] |
| 14 | |
| 15 | |
| 16 | @config(config=ConfigDict(arbitrary_types_allowed=True)) |
| 17 | class DeviceConfig: |
| 18 | """Configuration for the device to use for vLLM execution.""" |
| 19 | |
| 20 | device: SkipValidation[Device | torch.device | None] = "auto" |
| 21 | """Device type for vLLM execution. |
| 22 | This parameter is deprecated and will be |
| 23 | removed in a future release. |
| 24 | It will now be set automatically based |
| 25 | on the current platform.""" |
| 26 | device_type: str = field(init=False) |
| 27 | """Device type from the current platform. This is set in |
| 28 | `__post_init__`.""" |
| 29 | |
| 30 | def compute_hash(self) -> str: |
| 31 | """WARNING: Whenever a new field is added to this config, |
| 32 | ensure that it is included in the factors list if |
| 33 | it affects the computation graph. |
| 34 | |
| 35 | Provide a hash that uniquely identifies all the configs |
| 36 | that affect the structure of the computation |
| 37 | graph from input ids/embeddings to the final hidden states, |
| 38 | excluding anything before input ids/embeddings and after |
| 39 | the final hidden states. |
| 40 | """ |
| 41 | # no factors to consider. |
| 42 | # the device/platform information will be summarized |
| 43 | # by torch/vllm automatically. |
| 44 | factors: list[Any] = [] |
| 45 | hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest() |
| 46 | return hash_str |
| 47 | |
| 48 | def __post_init__(self): |
| 49 | if self.device == "auto": |
| 50 | # Automated device type detection |
| 51 | from vllm.platforms import current_platform |
| 52 | |
| 53 | self.device_type = current_platform.device_type |
| 54 | if not self.device_type: |
| 55 | raise RuntimeError( |
| 56 | "Failed to infer device type, please set " |
| 57 | "the environment variable `VLLM_LOGGING_LEVEL=DEBUG` " |
| 58 | "to turn on verbose logging to help debug the issue." |
| 59 | ) |
| 60 | else: |
| 61 | # Device type is assigned explicitly |
| 62 | if isinstance(self.device, str): |
| 63 | self.device_type = self.device |
| 64 | elif isinstance(self.device, torch.device): |
| 65 | self.device_type = self.device.type |
| 66 | |
| 67 | # Some platforms require processing inputs on CPU. |
| 68 | from vllm.platforms import current_platform |
| 69 | |
| 70 | if ( |
| 71 | current_platform.uses_host_device_handling() |
| 72 | and self.device_type == current_platform.device_type |
| 73 | ): |
| 74 | self.device = None |
| 75 | else: |
| 76 | # Set device with device type |
| 77 | self.device = torch.device(self.device_type) |