interface.py 20.5 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import contextlib
4
import enum
5
import os
6
import platform
7
import random
8
import sys
9
from datetime import timedelta
10
from typing import TYPE_CHECKING, Any, NamedTuple
11

12
import numpy as np
13
14
import torch

15
16
from vllm.logger import init_logger

17
if TYPE_CHECKING:
18
19
    from torch.distributed import PrefixStore, ProcessGroup

20
    from vllm.attention.backends.registry import AttentionBackendEnum
21
    from vllm.config import VllmConfig
22
    from vllm.config.cache import CacheDType
23
    from vllm.inputs import ProcessorInputs, PromptType
24
25
    from vllm.pooling_params import PoolingParams
    from vllm.sampling_params import SamplingParams
26
    from vllm.utils.argparse_utils import FlexibleArgumentParser
27
else:
28
    FlexibleArgumentParser = object
29

30
31
logger = init_logger(__name__)

32

33
34
def in_wsl() -> bool:
    # Reference: https://github.com/microsoft/WSL/issues/4071
35
    return "microsoft" in " ".join(platform.uname()).lower()
36
37


38
39
40
class PlatformEnum(enum.Enum):
    CUDA = enum.auto()
    ROCM = enum.auto()
41
    TPU = enum.auto()
42
    XPU = enum.auto()
43
    CPU = enum.auto()
44
    OOT = enum.auto()
45
    UNSPECIFIED = enum.auto()
46
47


48
49
50
51
class CpuArchEnum(enum.Enum):
    X86 = enum.auto()
    ARM = enum.auto()
    POWERPC = enum.auto()
52
    S390X = enum.auto()
53
    RISCV = enum.auto()
54
55
56
57
    OTHER = enum.auto()
    UNKNOWN = enum.auto()


58
59
60
61
class DeviceCapability(NamedTuple):
    major: int
    minor: int

62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
    def __lt__(self, other: Any) -> bool:
        if not isinstance(other, DeviceCapability):
            return NotImplemented
        return (self.major, self.minor) < (other.major, other.minor)

    def __le__(self, other: Any) -> bool:
        if not isinstance(other, DeviceCapability):
            return NotImplemented
        return (self.major, self.minor) <= (other.major, other.minor)

    def __eq__(self, other: Any) -> bool:
        if not isinstance(other, DeviceCapability):
            return NotImplemented
        return (self.major, self.minor) == (other.major, other.minor)

    def __ge__(self, other: Any) -> bool:
        if not isinstance(other, DeviceCapability):
            return NotImplemented
        return (self.major, self.minor) >= (other.major, other.minor)

    def __gt__(self, other: Any) -> bool:
        if not isinstance(other, DeviceCapability):
            return NotImplemented
        return (self.major, self.minor) > (other.major, other.minor)

87
88
89
90
91
    def as_version_str(self) -> str:
        return f"{self.major}.{self.minor}"

    def to_int(self) -> int:
        """
92
        Express device capability as an integer `<major><minor>`.
93
94
95
96
97
98
99

        It is assumed that the minor version is always a single digit.
        """
        assert 0 <= self.minor < 10
        return self.major * 10 + self.minor


100
101
class Platform:
    _enum: PlatformEnum
102
    device_name: str
103
    device_type: str
104

105
106
107
108
    # available dispatch keys:
    # check https://github.com/pytorch/pytorch/blob/313dac6c1ca0fa0cde32477509cce32089f8532a/torchgen/model.py#L134 # noqa
    # use "CPU" as a fallback for platforms not registered in PyTorch
    dispatch_key: str = "CPU"
109

110
111
112
113
    # available ray device keys:
    # https://github.com/ray-project/ray/blob/10ba5adadcc49c60af2c358a33bb943fb491a171/python/ray/_private/ray_constants.py#L438 # noqa
    # empty string means the device does not support ray
    ray_device_key: str = ""
114
115
116
117
118
119
120

    # platform-agnostic way to specify the device control environment variable,
    # .e.g. CUDA_VISIBLE_DEVICES for CUDA.
    # hint: search for "get_visible_accelerator_ids_env_var" in
    # https://github.com/ray-project/ray/tree/master/python/ray/_private/accelerators # noqa
    device_control_env_var: str = "VLLM_DEVICE_CONTROL_ENV_VAR_PLACEHOLDER"

121
122
123
124
125
126
    # The torch.compile backend for compiling simple and
    # standalone functions. The default value is "inductor" to keep
    # the same behavior as PyTorch.
    # NOTE: for the forward part of the model, vLLM has another separate
    # compilation strategy.
    simple_compile_backend: str = "inductor"
127

128
129
130
    # The backend used for distributed communication.
    dist_backend: str = ""

131
    supported_quantization: list[str] = []
132

133
134
    additional_env_vars: list[str] = []

135
    _global_graph_pool: Any | None = None
136

137
138
139
140
141
    @property
    def pass_key(self) -> str:
        """Inductor config key for the PassManager custom pass"""
        return "post_grad_custom_post_pass"

142
143
144
145
146
147
148
149
    @property
    def supported_dtypes(self) -> list[torch.dtype]:
        """Returns the supported dtypes for the current platform."""
        # Be careful with the order of the dtypes. The first dtype will
        # be used as the default dtype fallback for the current platform,
        # when encountering unsupported dtypes in "auto" dtype.
        return [torch.bfloat16, torch.float16, torch.float32]

150
151
152
153
154
155
    def is_cuda(self) -> bool:
        return self._enum == PlatformEnum.CUDA

    def is_rocm(self) -> bool:
        return self._enum == PlatformEnum.ROCM

156
157
158
    def is_tpu(self) -> bool:
        return self._enum == PlatformEnum.TPU

159
160
161
    def is_xpu(self) -> bool:
        return self._enum == PlatformEnum.XPU

162
163
164
    def is_cpu(self) -> bool:
        return self._enum == PlatformEnum.CPU

165
166
167
    def is_out_of_tree(self) -> bool:
        return self._enum == PlatformEnum.OOT

168
169
170
    def is_unspecified(self) -> bool:
        return self._enum == PlatformEnum.UNSPECIFIED

171
172
173
    def get_max_output_tokens(self, prompt_len: int) -> int:
        return sys.maxsize

174
    def is_cuda_alike(self) -> bool:
175
        """Stateless version of [torch.cuda.is_available][]."""
176
177
        return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM)

178
    def is_sleep_mode_available(self) -> bool:
179
180
181
182
183
        # TODO: Actually only mi3xx has the sleep mode support now
        # for ROCm, but currently we don't have a way to detect the
        # exact GPU model statelessly here. So we return True for
        # all ROCm platforms for now.
        return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM)
184

185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
    @classmethod
    def get_pass_manager_cls(cls) -> str:
        """
        Get the pass manager class for this platform.
        It will be registered as a custom pass under the current_platform.pass_key.
        """
        return "vllm.compilation.pass_manager.PostGradPassManager"

    @classmethod
    def get_compile_backend(cls) -> str:
        """
        Get the custom compile backend for current platform.
        """
        return cls.simple_compile_backend

200
201
    @classmethod
    def device_id_to_physical_device_id(cls, device_id: int):
202
203
204
        # Treat empty device control env var as unset. This is a valid
        # configuration in Ray setups where the engine is launched in
        # a CPU-only placement group located on a GPU node.
205
206
207
208
        if (
            cls.device_control_env_var in os.environ
            and os.environ[cls.device_control_env_var] != ""
        ):
209
210
211
212
213
214
            device_ids = os.environ[cls.device_control_env_var].split(",")
            physical_device_id = device_ids[device_id]
            return int(physical_device_id)
        else:
            return device_id

215
    @classmethod
216
    def import_kernels(cls) -> None:
217
        """Import any platform-specific C kernels."""
218
219
220
221
222
223
        try:
            import vllm._C  # noqa: F401
        except ImportError as e:
            logger.warning("Failed to import from vllm._C: %r", e)
        with contextlib.suppress(ImportError):
            import vllm._moe_C  # noqa: F401
224

225
    @classmethod
226
227
228
229
230
    def get_vit_attn_backend(
        cls, head_size: int, dtype: torch.dtype
    ) -> "AttentionBackendEnum":
        # Import AttentionBackendEnum here to avoid circular import.
        from vllm.attention.backends.registry import AttentionBackendEnum
231

232
        return AttentionBackendEnum.TORCH_SDPA
233

234
    @classmethod
235
236
    def get_attn_backend_cls(
        cls,
237
        selected_backend: "AttentionBackendEnum",
238
239
        head_size: int,
        dtype: torch.dtype,
240
        kv_cache_dtype: "CacheDType | None",
241
242
243
244
        block_size: int,
        use_mla: bool,
        has_sink: bool,
        use_sparse: bool,
245
        attn_type: str | None = None,
246
    ) -> str:
247
248
        """Get the attention backend class of a device."""
        return ""
249

250
251
252
253
    @classmethod
    def get_device_capability(
        cls,
        device_id: int = 0,
254
    ) -> DeviceCapability | None:
255
        """Stateless version of [torch.cuda.get_device_capability][]."""
256
        return None
257

258
259
260
    @classmethod
    def has_device_capability(
        cls,
261
        capability: tuple[int, int] | int,
262
263
264
265
266
        device_id: int = 0,
    ) -> bool:
        """
        Test whether this platform is compatible with a device capability.

267
        The `capability` argument can either be:
268

269
270
271
        - A tuple `(major, minor)`.
        - An integer `<major><minor>`. (See
        [`DeviceCapability.to_int`][vllm.platforms.interface.DeviceCapability.to_int])
272
273
274
275
276
277
278
279
280
281
        """
        current_capability = cls.get_device_capability(device_id=device_id)
        if current_capability is None:
            return False

        if isinstance(capability, tuple):
            return current_capability >= capability

        return current_capability.to_int() >= capability

282
283
284
    @classmethod
    def is_device_capability(
        cls,
285
        capability: tuple[int, int] | int,
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
        device_id: int = 0,
    ) -> bool:
        """
        Test whether this platform has exactly the specified device capability.

        The `capability` argument can either be:

        - A tuple `(major, minor)`.
        - An integer `<major><minor>`. (See
        [`DeviceCapability.to_int`][vllm.platforms.interface.DeviceCapability.to_int])
        """
        current_capability = cls.get_device_capability(device_id=device_id)
        if current_capability is None:
            return False

        if isinstance(capability, tuple):
            return current_capability == capability

        return current_capability.to_int() == capability

306
307
    @classmethod
    def get_device_name(cls, device_id: int = 0) -> str:
308
309
310
        """Get the name of a device."""
        raise NotImplementedError

311
312
313
314
315
    @classmethod
    def get_device_uuid(cls, device_id: int = 0) -> str:
        """Get the uuid of a device, e.g. the PCI bus ID."""
        raise NotImplementedError

316
317
318
    @classmethod
    def get_device_total_memory(cls, device_id: int = 0) -> int:
        """Get the total memory of a device in bytes."""
319
320
        raise NotImplementedError

321
322
    @classmethod
    def inference_mode(cls):
323
324
325
326
327
328
329
330
        """A device-specific wrapper of `torch.inference_mode`.

        This wrapper is recommended because some hardware backends such as TPU
        do not support `torch.inference_mode`. In such a case, they will fall
        back to `torch.no_grad` by overriding this method.
        """
        return torch.inference_mode(mode=True)

331
    @classmethod
332
    def seed_everything(cls, seed: int | None = None) -> None:
333
334
335
336
337
338
        """
        Set the seed of each random module.
        `torch.manual_seed` will set seed on all devices.

        Loosely based on: https://github.com/Lightning-AI/pytorch-lightning/blob/2.4.0/src/lightning/fabric/utilities/seed.py#L20
        """
339
340
341
342
        if seed is not None:
            random.seed(seed)
            np.random.seed(seed)
            torch.manual_seed(seed)
343

344
345
346
347
348
    @classmethod
    def set_device(cls, device: torch.device) -> None:
        """
        Set the device for the current platform.
        """
349
        raise NotImplementedError
350

351
    @classmethod
352
    def pre_register_and_update(
353
        cls, parser: FlexibleArgumentParser | None = None
354
    ) -> None:
355
        """
356
        Do some pre-registration or update action for the current platform.
357
358
359
360
361
362
363
364
365
366

        This function is called before global VllmConfig is initialized or cli
        arguments are parsed. It's used for out-of-tree platforms to register or
        update the configuration.

        For example, the out-of-tree quantization config can be imported and
        registered here dynamically.
        """
        pass

367
    @classmethod
368
    def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
369
370
371
372
373
374
375
376
377
378
379
        """
        Check and update the configuration for the current platform.

        It can raise an exception if the configuration is not compatible with
        the current platform, or it can update the configuration to make it
        compatible with the current platform.

        The config is passed by reference, so it can be modified in place.
        """
        pass

380
381
382
383
384
385
386
387
388
389
390
391
    @classmethod
    def verify_model_arch(cls, model_arch: str) -> None:
        """
        Verify whether the current platform supports the specified model
        architecture.

        - This will raise an Error or Warning based on the model support on
        the current platform.
        - By default all models are considered supported.
        """
        pass

392
393
394
395
396
    @classmethod
    def verify_quantization(cls, quant: str) -> None:
        """
        Verify whether the quantization is supported by the current platform.
        """
397
        if cls.supported_quantization and quant not in cls.supported_quantization:
398
            raise ValueError(
399
400
                f"{quant} quantization is currently not supported in {cls.device_name}."
            )
401

402
403
404
405
406
407
408
409
410
411
412
413
414
415
    @classmethod
    def get_cpu_architecture(cls) -> CpuArchEnum:
        """
        Determine the CPU architecture of the current system.
        Returns CpuArchEnum indicating the architecture type.
        """
        machine = platform.machine().lower()

        if machine in ("x86_64", "amd64", "i386", "i686"):
            return CpuArchEnum.X86
        elif machine.startswith("arm") or machine.startswith("aarch"):
            return CpuArchEnum.ARM
        elif machine.startswith("ppc"):
            return CpuArchEnum.POWERPC
416
417
        elif machine == "s390x":
            return CpuArchEnum.S390X
418
419
        elif machine.startswith("riscv"):
            return CpuArchEnum.RISCV
420
421
422

        return CpuArchEnum.OTHER if machine else CpuArchEnum.UNKNOWN

423
424
425
426
427
428
    @classmethod
    def is_pin_memory_available(cls) -> bool:
        """Checks whether pin memory is available on the current platform."""
        if in_wsl():
            # Pinning memory in WSL is not supported.
            # https://docs.nvidia.com/cuda/wsl-user-guide/index.html#known-limitations-for-linux-cuda-applications
429
430
431
432
            logger.warning(
                "Using 'pin_memory=False' as WSL is detected. "
                "This may slow down the performance."
            )
433
434
435
            return False
        return True

436
    @classmethod
437
    def get_current_memory_usage(
438
        cls, device: torch.types.Device | None = None
439
    ) -> float:
440
441
442
443
444
        """
        Return the memory usage in bytes.
        """
        raise NotImplementedError

445
446
447
448
449
450
451
    @classmethod
    def get_punica_wrapper(cls) -> str:
        """
        Return the punica wrapper for current platform.
        """
        raise NotImplementedError

452
    @classmethod
453
    def get_infinity_values(cls, dtype: torch.dtype) -> tuple[float, float]:
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
        """
        Return the platform specific values for (-inf, inf)
        """
        return float("-inf"), float("inf")

    @classmethod
    def can_update_inplace(cls) -> bool:
        """
        Checks if the platform allows inplace memory updates
        """
        return True

    @classmethod
    def get_lora_vocab_padding_size(cls) -> int:
        """
        Returns how much padding the LoRA logits need for kernels
        """
        return 256

473
474
475
476
477
    @classmethod
    def get_device_communicator_cls(cls) -> str:
        """
        Get device specific communicator class for distributed communication.
        """
Mengqing Cao's avatar
Mengqing Cao committed
478
        return "vllm.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase"  # noqa
479

480
481
482
483
484
485
486
    @classmethod
    def supports_mx(cls) -> bool:
        """
        Returns whether the current platform supports MX types.
        """
        return False

487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
    @classmethod
    def supports_fp8(cls) -> bool:
        """
        Returns whether the current platform supports FP8 types.
        """
        return False

    @classmethod
    def is_fp8_fnuz(cls) -> bool:
        """
        Returns whether the preferred FP8 type is FNUZ on the current platform.

        There are two representations of FP8, OCP FP8 and FNUZ FP8.
        The OCP specification can be found at https://tinyurl.com/b7jvwpft.
        The FNUZ specification can be found at https://tinyurl.com/5n6hwwu5.

        AMD's MI300 and MI325 have native hardware support for FNUZ. All other
        hardware has converged on the OCP FP8 standard.
        """
        return False

    @classmethod
    def fp8_dtype(cls) -> torch.dtype:
        """
        Returns the preferred FP8 type on the current platform.

        See the documentation for is_fp8_fnuz for details.
        """
        return torch.float8_e4m3fn

517
518
519
520
521
    @classmethod
    def use_all_gather(cls) -> bool:
        """
        Whether to use allgather in LogitsProcessor to gather the logits.
        """
522
        return True
523

524
525
526
527
528
529
530
    @classmethod
    def use_custom_allreduce(cls) -> bool:
        """
        Returns if custom allreduce is supported on the current platform
        """
        return False

531
532
533
534
535
536
537
538
    @classmethod
    def opaque_attention_op(cls) -> bool:
        """
        Returns True if we register attention as one giant opaque custom op
        on the current platform
        """
        return False

539
540
541
    @classmethod
    def validate_request(
        cls,
542
543
544
        prompt: "PromptType",
        params: "SamplingParams | PoolingParams",
        processed_inputs: "ProcessorInputs",
545
546
547
    ) -> None:
        """Raises if this request is unsupported on this platform"""

548
    def __getattr__(self, key: str):
549
        device = getattr(torch, self.device_type, None)
550
551
552
        if device is not None and hasattr(device, key):
            return getattr(device, key)
        else:
553
554
555
556
557
            logger.warning(
                "Current platform %s does not have '%s' attribute.",
                self.device_type,
                key,
            )
558
559
            return None

560
561
    def get_global_graph_pool(self) -> Any:
        """
562
        Return the global graph pool for this platform.
563
564
565
566
567
568
        """
        cls = self.__class__
        if cls._global_graph_pool is None:
            cls._global_graph_pool = self.graph_pool_handle()
        return cls._global_graph_pool

569
    @classmethod
570
    def get_static_graph_wrapper_cls(cls) -> str:
571
        """
572
        Get static graph wrapper class for static graph.
573
        """
574
        return "vllm.compilation.base_static_graph.AbstractStaticGraphWrapper"
575

576
577
578
579
    @classmethod
    def stateless_init_device_torch_dist_pg(
        cls,
        backend: str,
580
        prefix_store: "PrefixStore",
581
582
583
        group_rank: int,
        group_size: int,
        timeout: timedelta,
584
    ) -> "ProcessGroup":
585
586
587
        """
        Init platform-specific torch distributed process group.
        """
588
        raise NotImplementedError
589

590
    @classmethod
591
    def check_if_supports_dtype(cls, dtype: torch.dtype):
592
593
594
595
596
        """
        Check if the dtype is supported by the current platform.
        """
        raise NotImplementedError

597
598
599
600
601
602
603
    @classmethod
    def support_hybrid_kv_cache(cls) -> bool:
        """
        Returns if the hybrid kv cache is supported by the current platform.
        """
        return False

604
605
606
607
608
609
610
    @classmethod
    def support_static_graph_mode(cls) -> bool:
        """
        Returns if the graph mode is supported by the current platform.
        """
        return False

611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
    @classmethod
    def use_sync_weight_loader(cls) -> bool:
        """
        Returns if the current platform needs to sync weight loader.
        """
        return False

    @classmethod
    def make_synced_weight_loader(cls, original_weight_loader):
        """
        Wrap the original weight loader to make it synced.
        """
        if not cls.use_sync_weight_loader():
            return original_weight_loader

        def _synced_weight_loader(param, *args, **kwargs):
            out = original_weight_loader(param, *args, **kwargs)
            if param.device != torch.device("cpu"):
                torch._sync(param)
            return out

        return _synced_weight_loader

634
635
636
    @classmethod
    def get_nixl_supported_devices(cls) -> dict[str, tuple[str, ...]]:
        """
637
        Returns a mapping from device_type to a tuple of supported
638
639
640
641
642
        kv_buffer_device for nixl.
        """
        return {}

    @classmethod
643
    def get_nixl_memory_type(cls) -> str | None:
644
645
646
647
648
        """
        Returns the nixl memory type for the current platform.
        """
        return None

649
650
651
652
653
654
655
    @classmethod
    def check_max_model_len(cls, max_model_len: int) -> int:
        """
        Check max_model_len for the current platform.
        """
        return max_model_len

656
657
658

class UnspecifiedPlatform(Platform):
    _enum = PlatformEnum.UNSPECIFIED
659
    device_type = ""