| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import logging |
| 4 | import os |
| 5 | import traceback |
| 6 | from itertools import chain |
| 7 | from typing import TYPE_CHECKING |
| 8 | |
| 9 | from vllm import envs |
| 10 | from vllm.plugins import PLATFORM_PLUGINS_GROUP, load_plugins_by_group |
| 11 | from vllm.utils.import_utils import resolve_obj_by_qualname |
| 12 | |
| 13 | from .interface import CpuArchEnum, Platform, PlatformEnum, in_wsl |
| 14 | |
| 15 | logger = logging.getLogger(__name__) |
| 16 | |
| 17 | |
| 18 | def vllm_version_matches_substr(substr: str) -> bool: |
| 19 | """Check to see if the vLLM version matches a substring.""" |
| 20 | from importlib.metadata import PackageNotFoundError, version |
| 21 | |
| 22 | try: |
| 23 | vllm_version = version("vllm") |
| 24 | except PackageNotFoundError as e: |
| 25 | logger.warning( |
| 26 | "The vLLM package was not found, so its version could not be " |
| 27 | "inspected. This may cause platform detection to fail." |
| 28 | ) |
| 29 | raise e |
| 30 | return substr in vllm_version |
| 31 | |
| 32 | |
| 33 | def tpu_platform_plugin() -> str | None: |
| 34 | logger.debug("Checking if TPU platform is available.") |
| 35 | |
| 36 | # Check for Pathways TPU proxy |
| 37 | if envs.VLLM_TPU_USING_PATHWAYS: |
| 38 | logger.debug("Confirmed TPU platform is available via Pathways proxy.") |
| 39 | return "tpu_inference.platforms.tpu_platform.TpuPlatform" |
| 40 | |
| 41 | # Check for libtpu installation |
| 42 | try: |
| 43 | # While it's technically possible to install libtpu on a |
| 44 | # non-TPU machine, this is a very uncommon scenario. Therefore, |
| 45 | # we assume that libtpu is installed only if the machine |
| 46 | # has TPUs. |
| 47 | |
| 48 | import libtpu # noqa: F401 |
| 49 | |
| 50 | logger.debug("Confirmed TPU platform is available.") |
| 51 | return "vllm.platforms.tpu.TpuPlatform" |
| 52 | except Exception as e: |
| 53 | logger.debug("TPU platform is not available because: %s", str(e)) |
| 54 | return None |
| 55 | |
| 56 | |
| 57 | def cuda_platform_plugin() -> str | None: |
| 58 | is_cuda = False |
| 59 | logger.debug("Checking if CUDA platform is available.") |
| 60 | try: |
| 61 | from vllm.utils.import_utils import import_pynvml |
| 62 | |
| 63 | pynvml = import_pynvml() |
| 64 | pynvml.nvmlInit() |
| 65 | try: |
| 66 | # NOTE: Edge case: vllm cpu build on a GPU machine. |
| 67 | # Third-party pynvml can be imported in cpu build, |
| 68 | # we need to check if vllm is built with cpu too. |
| 69 | # Otherwise, vllm will always activate cuda plugin |
| 70 | # on a GPU machine, even if in a cpu build. |
| 71 | is_cuda = ( |
| 72 | pynvml.nvmlDeviceGetCount() > 0 |
| 73 | and not vllm_version_matches_substr("cpu") |
| 74 | ) |
| 75 | if pynvml.nvmlDeviceGetCount() <= 0: |
| 76 | logger.debug("CUDA platform is not available because no GPU is found.") |
| 77 | if vllm_version_matches_substr("cpu"): |
| 78 | logger.debug( |
| 79 | "CUDA platform is not available because vLLM is built with CPU." |
| 80 | ) |
| 81 | if is_cuda: |
| 82 | logger.debug("Confirmed CUDA platform is available.") |
| 83 | finally: |
| 84 | pynvml.nvmlShutdown() |
| 85 | except Exception as e: |
| 86 | logger.debug("Exception happens when checking CUDA platform: %s", str(e)) |
| 87 | if "nvml" not in e.__class__.__name__.lower(): |
| 88 | # If the error is not related to NVML, re-raise it. |
| 89 | raise e |
| 90 | |
| 91 | # CUDA is supported on Jetson, but NVML may not be. |
| 92 | import os |
| 93 | |
| 94 | def cuda_is_jetson() -> bool: |
| 95 | return os.path.isfile("/etc/nv_tegra_release") or os.path.exists( |
| 96 | "/sys/class/tegra-firmware" |
| 97 | ) |
| 98 | |
| 99 | if cuda_is_jetson(): |
| 100 | logger.debug("Confirmed CUDA platform is available on Jetson.") |
| 101 | is_cuda = True |
| 102 | else: |
| 103 | logger.debug("CUDA platform is not available because: %s", str(e)) |
| 104 | |
| 105 | return "vllm.platforms.cuda.CudaPlatform" if is_cuda else None |
| 106 | |
| 107 | |
| 108 | def rocm_platform_plugin() -> str | None: |
| 109 | is_rocm = False |
| 110 | logger.debug("Checking if ROCm platform is available.") |
| 111 | try: |
| 112 | import amdsmi |
| 113 | |
| 114 | amdsmi.amdsmi_init() |
| 115 | try: |
| 116 | if len(amdsmi.amdsmi_get_processor_handles()) > 0: |
| 117 | is_rocm = True |
| 118 | logger.debug("Confirmed ROCm platform is available.") |
| 119 | else: |
| 120 | logger.debug("ROCm platform is not available because no GPU is found.") |
| 121 | finally: |
| 122 | amdsmi.amdsmi_shut_down() |
| 123 | except Exception as e: |
| 124 | logger.debug("ROCm platform is not available because: %s", str(e)) |
| 125 | |
| 126 | if not is_rocm and in_wsl(): |
| 127 | try: |
| 128 | import torch |
| 129 | |
| 130 | if ( |
| 131 | not vllm_version_matches_substr("cpu") |
| 132 | and getattr(torch.version, "hip", None) |
| 133 | and torch.accelerator.is_available() |
| 134 | ): |
| 135 | is_rocm = True |
| 136 | logger.debug("Confirmed ROCm platform is available in WSL via PyTorch.") |
| 137 | except Exception as e: |
| 138 | logger.debug("WSL ROCm fallback detection failed because: %s", str(e)) |
| 139 | |
| 140 | return "vllm.platforms.rocm.RocmPlatform" if is_rocm else None |
| 141 | |
| 142 | |
| 143 | def xpu_platform_plugin() -> str | None: |
| 144 | is_xpu = False |
| 145 | logger.debug("Checking if XPU platform is available.") |
| 146 | try: |
| 147 | import torch |
| 148 | |
| 149 | if torch.distributed.is_xccl_available(): |
| 150 | dist_backend = "xccl" |
| 151 | from vllm.platforms.xpu import XPUPlatform |
| 152 | |
| 153 | XPUPlatform.dist_backend = dist_backend |
| 154 | logger.debug("Confirmed %s backend is available.", XPUPlatform.dist_backend) |
| 155 | |
| 156 | if hasattr(torch, "xpu") and torch.xpu.is_available(): |
| 157 | is_xpu = True |
| 158 | logger.debug("Confirmed XPU platform is available.") |
| 159 | except Exception as e: |
| 160 | logger.debug("XPU platform is not available because: %s", str(e)) |
| 161 | |
| 162 | return "vllm.platforms.xpu.XPUPlatform" if is_xpu else None |
| 163 | |
| 164 | |
| 165 | def _is_amd_zen_cpu() -> bool: |
| 166 | """Detect AMD CPU with AVX-512 via /proc/cpuinfo.""" |
| 167 | if not os.path.exists("/proc/cpuinfo"): |
| 168 | return False |
| 169 | with open("/proc/cpuinfo") as f: |
| 170 | cpuinfo = f.read() |
| 171 | return "AuthenticAMD" in cpuinfo and "avx512" in cpuinfo |
| 172 | |
| 173 | |
| 174 | def cpu_platform_plugin() -> str | None: |
| 175 | logger.debug("Checking if CPU platform is available.") |
| 176 | is_cpu = envs.VLLM_TARGET_DEVICE == "cpu" |
| 177 | if is_cpu: |
| 178 | logger.debug( |
| 179 | "Confirmed CPU platform is available because " |
| 180 | "VLLM_TARGET_DEVICE is set to CPU." |
| 181 | ) |
| 182 | else: |
| 183 | try: |
| 184 | is_cpu = vllm_version_matches_substr("cpu") |
| 185 | if is_cpu: |
| 186 | logger.debug( |
| 187 | "Confirmed CPU platform is available because vLLM is built " |
| 188 | "with CPU." |
| 189 | ) |
| 190 | if not is_cpu: |
| 191 | import sys |
| 192 | |
| 193 | is_cpu = sys.platform.startswith("darwin") |
| 194 | if is_cpu: |
| 195 | logger.debug( |
| 196 | "Confirmed CPU platform is available because the machine " |
| 197 | "is MacOS." |
| 198 | ) |
| 199 | except Exception as e: |
| 200 | logger.debug("CPU platform is not available because: %s", str(e)) |
| 201 | |
| 202 | if not is_cpu: |
| 203 | return None |
| 204 | |
| 205 | if _is_amd_zen_cpu(): |
| 206 | try: |
| 207 | import zentorch # noqa: F401 |
| 208 | |
| 209 | logger.info( |
| 210 | "AMD Zen CPU detected with zentorch installed, using ZenCpuPlatform." |
| 211 | ) |
| 212 | return "vllm.platforms.zen_cpu.ZenCpuPlatform" |
| 213 | except ImportError: |
| 214 | logger.debug( |
| 215 | "AMD Zen CPU detected but zentorch not installed, " |
| 216 | "falling back to CpuPlatform." |
| 217 | ) |
| 218 | except OSError: |
| 219 | # An ABI-mismatched build fails here with an undefined-symbol |
| 220 | # error; other failures are not known to be safe to recover from. |
| 221 | logger.warning( |
| 222 | "AMD Zen CPU detected but zentorch failed to import, falling " |
| 223 | "back to CpuPlatform. This usually means the zentorch build " |
| 224 | "does not match the installed torch version.", |
| 225 | exc_info=True, |
| 226 | ) |
| 227 | |
| 228 | return "vllm.platforms.cpu.CpuPlatform" |
| 229 | |
| 230 | |
| 231 | builtin_platform_plugins = { |
| 232 | "tpu": tpu_platform_plugin, |
| 233 | "cuda": cuda_platform_plugin, |
| 234 | "rocm": rocm_platform_plugin, |
| 235 | "xpu": xpu_platform_plugin, |
| 236 | "cpu": cpu_platform_plugin, |
| 237 | } |
| 238 | |
| 239 | |
| 240 | def resolve_current_platform_cls_qualname() -> str: |
| 241 | # An explicit CPU target is authoritative. Native CPU-only CI jobs reuse |
| 242 | # an accelerator wheel and can run on accelerator hosts, so probing every |
| 243 | # plugin would otherwise activate both CPU and the host accelerator. |
| 244 | if envs.VLLM_TARGET_DEVICE == "cpu": |
| 245 | cpu_platform_cls_qualname = cpu_platform_plugin() |
| 246 | assert cpu_platform_cls_qualname is not None |
| 247 | logger.debug("Explicitly selected CPU platform.") |
| 248 | return cpu_platform_cls_qualname |
| 249 | |
| 250 | platform_plugins = load_plugins_by_group(PLATFORM_PLUGINS_GROUP) |
| 251 | |
| 252 | activated_plugins = [] |
| 253 | |
| 254 | for name, func in chain(builtin_platform_plugins.items(), platform_plugins.items()): |
| 255 | try: |
| 256 | assert callable(func) |
| 257 | platform_cls_qualname = func() |
| 258 | if platform_cls_qualname is not None: |
| 259 | activated_plugins.append(name) |
| 260 | except Exception: |
| 261 | logger.debug( |
| 262 | "Platform plugin %s failed during detection.", |
| 263 | name, |
| 264 | exc_info=True, |
| 265 | ) |
| 266 | |
| 267 | activated_builtin_plugins = list( |
| 268 | set(activated_plugins) & set(builtin_platform_plugins.keys()) |
| 269 | ) |
| 270 | activated_oot_plugins = list(set(activated_plugins) & set(platform_plugins.keys())) |
| 271 | |
| 272 | if len(activated_oot_plugins) >= 2: |
| 273 | raise RuntimeError( |
| 274 | "Only one platform plugin can be activated, but got: " |
| 275 | f"{activated_oot_plugins}" |
| 276 | ) |
| 277 | elif len(activated_oot_plugins) == 1: |
| 278 | platform_cls_qualname = platform_plugins[activated_oot_plugins[0]]() |
| 279 | logger.info("Platform plugin %s is activated", activated_oot_plugins[0]) |
| 280 | elif len(activated_builtin_plugins) >= 2: |
| 281 | raise RuntimeError( |
| 282 | "Only one platform plugin can be activated, but got: " |
| 283 | f"{activated_builtin_plugins}" |
| 284 | ) |
| 285 | elif len(activated_builtin_plugins) == 1: |
| 286 | platform_cls_qualname = builtin_platform_plugins[activated_builtin_plugins[0]]() |
| 287 | logger.debug( |
| 288 | "Automatically detected platform %s.", activated_builtin_plugins[0] |
| 289 | ) |
| 290 | else: |
| 291 | platform_cls_qualname = "vllm.platforms.interface.UnspecifiedPlatform" |
| 292 | logger.debug("No platform detected, vLLM is running on UnspecifiedPlatform") |
| 293 | return platform_cls_qualname |
| 294 | |
| 295 | |
| 296 | _current_platform = None |
| 297 | _init_trace: str = "" |
| 298 | |
| 299 | if TYPE_CHECKING: |
| 300 | current_platform: Platform |
| 301 | |
| 302 | |
| 303 | def __getattr__(name: str): |
| 304 | if name == "current_platform": |
| 305 | # lazy init current_platform. |
| 306 | # 1. out-of-tree platform plugins need `from vllm.platforms import |
| 307 | # Platform` so that they can inherit `Platform` class. Therefore, |
| 308 | # we cannot resolve `current_platform` during the import of |
| 309 | # `vllm.platforms`. |
| 310 | # 2. when users use out-of-tree platform plugins, they might run |
| 311 | # `import vllm`, some vllm internal code might access |
| 312 | # `current_platform` during the import, and we need to make sure |
| 313 | # `current_platform` is only resolved after the plugins are loaded |
| 314 | # (we have tests for this, if any developer violate this, they will |
| 315 | # see the test failures). |
| 316 | global _current_platform |
| 317 | if _current_platform is None: |
| 318 | platform_cls_qualname = resolve_current_platform_cls_qualname() |
| 319 | _current_platform = resolve_obj_by_qualname(platform_cls_qualname)() |
| 320 | global _init_trace |
| 321 | _init_trace = "".join(traceback.format_stack()) |
| 322 | return _current_platform |
| 323 | elif name in globals(): |
| 324 | return globals()[name] |
| 325 | else: |
| 326 | raise AttributeError(f"No attribute named '{name}' exists in {__name__}.") |
| 327 | |
| 328 | |
| 329 | def __setattr__(name: str, value): |
| 330 | if name == "current_platform": |
| 331 | global _current_platform |
| 332 | _current_platform = value |
| 333 | elif name in globals(): |
| 334 | globals()[name] = value |
| 335 | else: |
| 336 | raise AttributeError(f"No attribute named '{name}' exists in {__name__}.") |
| 337 | |
| 338 | |
| 339 | __all__ = [ |
| 340 | "Platform", |
| 341 | "PlatformEnum", |
| 342 | "current_platform", |
| 343 | "CpuArchEnum", |
| 344 | "_init_trace", |
| 345 | "_is_amd_zen_cpu", |
| 346 | ] |