cuda.py 26 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
6
"""Code inside this file can safely assume cuda platform, e.g. importing
pynvml. However, it should not initialize cuda context.
"""

7
import os
8
from collections.abc import Callable
9
from datetime import timedelta
10
from functools import cache, wraps
11
from typing import TYPE_CHECKING, TypeVar
12

13
import torch
14
15
from torch.distributed import PrefixStore, ProcessGroup
from torch.distributed.distributed_c10d import is_nccl_available
16
from typing_extensions import ParamSpec
17

18
19
# import custom ops, trigger op registration
import vllm._C  # noqa
20
from vllm.logger import init_logger
21
from vllm.utils.import_utils import import_pynvml
22
from vllm.utils.torch_utils import cuda_device_count_stateless
23
from vllm.v1.attention.backends.registry import AttentionBackendEnum
24

25
from .interface import DeviceCapability, Platform, PlatformEnum
26

27
if TYPE_CHECKING:
28
    from vllm.config import VllmConfig
29
    from vllm.config.cache import CacheDType
30
    from vllm.v1.attention.selector import AttentionSelectorConfig
31
else:
32
33
    VllmConfig = None
    CacheDType = None
34

35
36
logger = init_logger(__name__)

37
38
39
_P = ParamSpec("_P")
_R = TypeVar("_R")

40
pynvml = import_pynvml()
41

42
43
44
45
# pytorch 2.5 uses cudnn sdpa by default, which will cause crash on some models
# see https://github.com/huggingface/diffusers/issues/9704 for details
torch.backends.cuda.enable_cudnn_sdp(False)

46

47
48
49
50
@cache
def _get_backend_priorities(
    use_mla: bool,
    device_capability: DeviceCapability,
51
    num_heads: int | None = None,
52
53
54
55
) -> list[AttentionBackendEnum]:
    """Get backend priorities with lazy import to avoid circular dependency."""
    if use_mla:
        if device_capability.major == 10:
56
57
58
59
60
61
62
63
64
65
66
            # Prefer FlashInfer at low head counts (FlashMLA uses padding)
            if num_heads is not None and num_heads <= 16:
                sparse_backends = [
                    AttentionBackendEnum.FLASHINFER_MLA_SPARSE,
                    AttentionBackendEnum.FLASHMLA_SPARSE,
                ]
            else:
                sparse_backends = [
                    AttentionBackendEnum.FLASHMLA_SPARSE,
                    AttentionBackendEnum.FLASHINFER_MLA_SPARSE,
                ]
67
            return [
68
                AttentionBackendEnum.FLASHINFER_MLA,
69
                AttentionBackendEnum.CUTLASS_MLA,
70
                AttentionBackendEnum.FLASH_ATTN_MLA,
71
                AttentionBackendEnum.FLASHMLA,
72
                AttentionBackendEnum.TRITON_MLA,
73
                *sparse_backends,
74
75
76
77
            ]
        else:
            return [
                AttentionBackendEnum.FLASH_ATTN_MLA,
78
                AttentionBackendEnum.FLASHMLA,
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
                AttentionBackendEnum.FLASHINFER_MLA,
                AttentionBackendEnum.TRITON_MLA,
                AttentionBackendEnum.FLASHMLA_SPARSE,
            ]
    else:
        if device_capability.major == 10:
            return [
                AttentionBackendEnum.FLASHINFER,
                AttentionBackendEnum.FLASH_ATTN,
                AttentionBackendEnum.TRITON_ATTN,
                AttentionBackendEnum.FLEX_ATTENTION,
            ]
        else:
            return [
                AttentionBackendEnum.FLASH_ATTN,
                AttentionBackendEnum.FLASHINFER,
                AttentionBackendEnum.TRITON_ATTN,
                AttentionBackendEnum.FLEX_ATTENTION,
            ]


100
def with_nvml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]:
101
    @wraps(fn)
102
    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
103
104
105
106
107
108
109
110
111
        pynvml.nvmlInit()
        try:
            return fn(*args, **kwargs)
        finally:
            pynvml.nvmlShutdown()

    return wrapper


112
113
class CudaPlatformBase(Platform):
    _enum = PlatformEnum.CUDA
114
    device_name: str = "cuda"
115
116
    device_type: str = "cuda"
    dispatch_key: str = "CUDA"
117
    ray_device_key: str = "GPU"
118
    dist_backend: str = "nccl"
119
    device_control_env_var: str = "CUDA_VISIBLE_DEVICES"
120
121
122
    ray_noset_device_env_vars: list[str] = [
        "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES",
    ]
123

124
    @property
125
    def supported_dtypes(self) -> list[torch.dtype]:
126
127
128
        if self.has_device_capability(80):
            # Ampere and Hopper or later NVIDIA GPUs.
            return [torch.bfloat16, torch.float16, torch.float32]
129
        if self.has_device_capability(60):
130
131
132
133
134
135
            # Pascal, Volta and Turing NVIDIA GPUs, BF16 is not supported
            return [torch.float16, torch.float32]
        # Kepler and Maxwell NVIDIA GPUs, only FP32 is supported,
        # though vLLM doesn't support these GPUs.
        return [torch.float32]

136
137
138
139
140
    @classmethod
    def set_device(cls, device: torch.device) -> None:
        """
        Set the device for the current platform.
        """
141
        torch.cuda.set_device(device)
142
143
144
145
146
        # With this trick we can force the device to be set eagerly
        # see https://github.com/pytorch/pytorch/issues/155668
        # for why and when it is needed
        _ = torch.zeros(1, device=device)

147
    @classmethod
148
    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
149
        raise NotImplementedError
150

151
152
153
    @classmethod
    def get_device_name(cls, device_id: int = 0) -> str:
        raise NotImplementedError
154

155
156
157
    @classmethod
    def get_device_total_memory(cls, device_id: int = 0) -> int:
        raise NotImplementedError
158

159
    @classmethod
160
    def is_fully_connected(cls, device_ids: list[int]) -> bool:
161
        raise NotImplementedError
162

163
164
165
    @classmethod
    def log_warnings(cls):
        pass
166

167
    @classmethod
168
    def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
169
170
        from vllm.v1.attention.backends.registry import AttentionBackendEnum

171
        parallel_config = vllm_config.parallel_config
172
        model_config = vllm_config.model_config
173

174
        if parallel_config.worker_cls == "auto":
175
            parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"
176

177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
        cache_config = vllm_config.cache_config
        if cache_config and cache_config.block_size is None:
            cache_config.block_size = 16

        # TODO(lucas): handle this more gracefully
        # Note: model_config may be None during testing
        # Note: block_size is initialized in
        # HybridAttentionMambaModelConfig.verify_and_update_config
        # for models with both attention and mamba,
        # and doesn't need to be reinitialized here
        if (
            model_config is not None
            and model_config.use_mla
            and cache_config.block_size is not None
        ):
            use_sparse = hasattr(vllm_config.model_config.hf_config, "index_topk")
            # If `--attention-config.backend` is not set and we are using MLA,
            # then we default to FlashMLA backend for non-blackwell GPUs,
            # else we default to CutlassMLA. For each case, we force the
            # required block_size.
            use_flashmla = False
            use_cutlass_mla = False
            use_flashinfer_mla = False
            use_flashmla_sparse = False
            use_flashinfer_mla_sparse = False

            from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported

            if vllm_config.attention_config.backend is None:
                # Default case
                hf_text_config = model_config.hf_text_config
                qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)
                if (
                    cls.is_device_capability_family(100)
                    and not use_sparse
                    and qk_nope_head_dim == 128
                ):
                    # Blackwell => Force FlashInfer MLA (unless sparse, i.e. DSv3.2)
                    # and only if qk_nope_head_dim == 128 (kernel constraint)
                    use_flashinfer_mla = True
                    # Set the backend in AttentionConfig so it's used during
                    # backend selection
                    vllm_config.attention_config.backend = (
                        AttentionBackendEnum.FLASHINFER_MLA
                    )
                elif cls.is_device_capability_family(100) and not use_sparse:
                    # Fall back to CUTLASS_MLA as 2nd priority on Blackwell
                    use_cutlass_mla = True
                elif is_flashmla_dense_supported()[0]:
                    # Non-Blackwell with FlashMLA support
                    use_flashmla = True
                else:
                    # Fallback: will use Triton MLA or other compatible backend
                    pass
            else:
                # Forced case
                backend = vllm_config.attention_config.backend
                use_flashmla = backend == AttentionBackendEnum.FLASHMLA
                use_cutlass_mla = backend == AttentionBackendEnum.CUTLASS_MLA
                use_flashinfer_mla = backend == AttentionBackendEnum.FLASHINFER_MLA
                use_flashmla_sparse = backend == AttentionBackendEnum.FLASHMLA_SPARSE
                use_flashinfer_mla_sparse = (
                    backend == AttentionBackendEnum.FLASHINFER_MLA_SPARSE
                )

            if (
                use_flashmla
                and is_flashmla_dense_supported()[0]
                and cache_config.block_size % 64 != 0
            ):
                cache_config.block_size = 64
                logger.info("Forcing kv cache block size to 64 for FlashMLA backend.")

            if use_cutlass_mla and cache_config.block_size % 128 != 0:
                cache_config.block_size = 128
                logger.info(
                    "Forcing kv cache block size to 128 for CUTLASS_MLA backend."
                )

            if (
                use_flashinfer_mla
                and cache_config.block_size != 32
                and cache_config.block_size % 64 != 0
            ):
                cache_config.block_size = 64
                logger.info(
                    "Forcing kv cache block size to 64 for FlashInferMLA backend."
                )

            if use_sparse:
                if not (use_flashmla_sparse or use_flashinfer_mla_sparse):
                    use_flashmla_sparse = True

                if use_flashmla_sparse and cache_config.block_size != 64:
                    cache_config.block_size = 64
                    logger.info(
                        "Forcing kv cache block size to 64 for FlashMLASparse backend."
                    )
                elif use_flashinfer_mla_sparse and cache_config.block_size not in (
                    32,
                    64,
                ):
                    cache_config.block_size = 64
                    logger.info(
                        "Forcing kv cache block size to 64 for FlashInferMLASparse "
                        "backend."
                    )

285
286
287
288
289
290
291
292
293
294
295
296
297
298
        scheduler_config = vllm_config.scheduler_config
        # Note: model_config may be None during testing
        if (
            model_config is not None
            and model_config.is_mm_prefix_lm
            and scheduler_config.is_multimodal_model
            and not scheduler_config.disable_chunked_mm_input
        ):
            logger.warning(
                "Forcing --disable_chunked_mm_input for models "
                "with multimodal-bidirectional attention."
            )
            scheduler_config.disable_chunked_mm_input = True

299
    @classmethod
300
    def get_current_memory_usage(
301
        cls, device: torch.types.Device | None = None
302
    ) -> float:
303
        torch.cuda.empty_cache()
304
305
306
        torch.cuda.reset_peak_memory_stats(device)
        return torch.cuda.max_memory_allocated(device)

307
    @classmethod
308
    def get_valid_backends(
309
        cls,
310
311
        device_capability: DeviceCapability,
        attn_selector_config: "AttentionSelectorConfig",
312
        num_heads: int | None = None,
313
314
    ) -> tuple[
        list[tuple["AttentionBackendEnum", int]],
315
        dict["AttentionBackendEnum", list[str]],
316
317
    ]:
        valid_backends_priorities = []
318
        invalid_reasons = {}
319

320
        backend_priorities = _get_backend_priorities(
321
322
323
            attn_selector_config.use_mla,
            device_capability,
            num_heads,
324
        )
325
326
327
328
        for priority, backend in enumerate(backend_priorities):
            try:
                backend_class = backend.get_class()
                invalid_reasons_i = backend_class.validate_configuration(
329
330
                    device_capability=device_capability,
                    **attn_selector_config._asdict(),
331
                )
332
333
334
            except ImportError:
                invalid_reasons_i = ["ImportError"]
            if invalid_reasons_i:
335
                invalid_reasons[backend] = invalid_reasons_i
336
337
            else:
                valid_backends_priorities.append((backend, priority))
338

339
        return valid_backends_priorities, invalid_reasons
340

341
    @classmethod
342
    def get_attn_backend_cls(
343
        cls,
344
        selected_backend: "AttentionBackendEnum",
345
        attn_selector_config: "AttentionSelectorConfig",
346
        num_heads: int | None = None,
347
348
349
350
351
    ) -> str:
        device_capability = cls.get_device_capability()
        assert device_capability is not None

        attn_selector_config = attn_selector_config._replace(block_size=None)
352
353
354
355
        # First try checking just the selected backend, if there is one.
        if selected_backend is not None:
            try:
                backend_class = selected_backend.get_class()
356
                invalid_reasons = backend_class.validate_configuration(
357
358
                    device_capability=device_capability,
                    **attn_selector_config._asdict(),
359
                )
360
            except ImportError:
361
362
363
364
365
366
367
368
369
                invalid_reasons = ["ImportError"]
            if invalid_reasons:
                raise ValueError(
                    f"Selected backend {selected_backend} is not valid for "
                    f"this configuration. Reason: {invalid_reasons}"
                )
            else:
                logger.info("Using %s backend.", selected_backend)
                return selected_backend.get_path()
370

371
372
        # No selected backend or the selected backend is invalid,
        # so we try finding a valid backend.
373
        valid_backends_priorities, invalid_reasons = cls.get_valid_backends(
374
375
            device_capability=device_capability,
            attn_selector_config=attn_selector_config,
376
            num_heads=num_heads,
377
        )
378
379
380
381
382
383
384
385
386
387
388
389
390
        reasons_str = (
            "{"
            + ", ".join(
                f"{backend.name}: [{', '.join(reasons)}]"
                for backend, reasons in invalid_reasons.items()
            )
            + "}"
        )
        config_str = attn_selector_config.__repr__()
        logger.debug_once(
            f"Some attention backends are not valid for {cls.device_name} with "
            f"{config_str}. Reasons: {reasons_str}."
        )
391
        if len(valid_backends_priorities) == 0:
392
393
394
395
            raise ValueError(
                f"No valid attention backend found for {cls.device_name} "
                f"with {config_str}. Reasons: {reasons_str}."
            )
396

397
398
399
400
401
402
403
404
405
406
407
408
409
        # We have found some valid backends. Select the one with the
        # highest priority.
        sorted_indices = sorted(
            range(len(valid_backends_priorities)),
            key=lambda i: valid_backends_priorities[i][1],
        )
        selected_index = sorted_indices[0]
        selected_backend = valid_backends_priorities[selected_index][0]
        logger.info_once(
            "Using %s attention backend out of potential backends: %s.",
            selected_backend.name,
            "[" + ", ".join(f"'{b[0].name}'" for b in valid_backends_priorities) + "]",
            scope="local",
410
411
        )

412
        return selected_backend.get_path()
413

414
415
416
417
    @classmethod
    def get_supported_vit_attn_backends(cls) -> list["AttentionBackendEnum"]:
        return [
            AttentionBackendEnum.FLASH_ATTN,
418
419
            AttentionBackendEnum.TRITON_ATTN,
            AttentionBackendEnum.TORCH_SDPA,
420
            AttentionBackendEnum.FLASHINFER,
421
422
423
424
425
426
427
        ]

    @classmethod
    def get_vit_attn_backend(
        cls,
        head_size: int,
        dtype: torch.dtype,
428
        backend: "AttentionBackendEnum | None" = None,
429
430
431
432
433
434
435
436
437
    ) -> "AttentionBackendEnum":
        if backend is not None:
            assert backend in cls.get_supported_vit_attn_backends(), (
                f"Backend {backend} is not supported for vit attention. "
                f"Supported backends are: {cls.get_supported_vit_attn_backends()}"
            )
            logger.info_once(f"Using backend {backend} for vit attention")
            return backend

438
439
440
441
        cc = cls.get_device_capability()
        for vit_attn_backend in cls.get_supported_vit_attn_backends():
            if vit_attn_backend == AttentionBackendEnum.TORCH_SDPA:
                continue
442
            try:
443
444
                backend_class = vit_attn_backend.get_class()
                is_backend_supported = backend_class.supports_head_size(
445
                    head_size
446
447
448
449
450
451
452
453
454
455
456
                ) and backend_class.supports_dtype(dtype)
                if cc is not None:
                    is_backend_supported = (
                        is_backend_supported
                        and backend_class.supports_compute_capability(cc)
                    )
                if is_backend_supported:
                    logger.info_once(
                        f"Using backend {vit_attn_backend} for vit attention"
                    )
                    return vit_attn_backend
457
458
459
460
461
            except ImportError:
                pass

        return AttentionBackendEnum.TORCH_SDPA

462
463
464
465
    @classmethod
    def get_punica_wrapper(cls) -> str:
        return "vllm.lora.punica_wrapper.punica_gpu.PunicaWrapperGPU"

466
467
    @classmethod
    def get_device_communicator_cls(cls) -> str:
468
469
470
        return (
            "vllm.distributed.device_communicators.cuda_communicator.CudaCommunicator"  # noqa
        )
471

472
473
474
475
    @classmethod
    def supports_fp8(cls) -> bool:
        return cls.has_device_capability(89)

476
477
478
479
    @classmethod
    def use_custom_allreduce(cls) -> bool:
        return True

480
481
482
483
    @classmethod
    def opaque_attention_op(cls) -> bool:
        return True

484
    @classmethod
485
486
    def get_static_graph_wrapper_cls(cls) -> str:
        return "vllm.compilation.cuda_graph.CUDAGraphWrapper"
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
517
518
    @classmethod
    def stateless_init_device_torch_dist_pg(
        cls,
        backend: str,
        prefix_store: PrefixStore,
        group_rank: int,
        group_size: int,
        timeout: timedelta,
    ) -> ProcessGroup:
        assert is_nccl_available()
        pg: ProcessGroup = ProcessGroup(
            prefix_store,
            group_rank,
            group_size,
        )
        from torch.distributed.distributed_c10d import ProcessGroupNCCL

        backend_options = ProcessGroupNCCL.Options()
        backend_options._timeout = timeout

        backend_class = ProcessGroupNCCL(
            prefix_store, group_rank, group_size, backend_options
        )
        backend_type = ProcessGroup.BackendType.NCCL
        device = torch.device("cuda")
        pg._set_default_backend(backend_type)
        backend_class._set_sequence_number_for_group()

        pg._register_backend(device, backend_type, backend_class)
        return pg

519
520
521
522
    @classmethod
    def device_count(cls) -> int:
        return cuda_device_count_stateless()

523
    @classmethod
524
525
    def check_if_supports_dtype(cls, dtype: torch.dtype):
        if dtype == torch.bfloat16:  # noqa: SIM102
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
            if not cls.has_device_capability(80):
                capability = cls.get_device_capability()
                gpu_name = cls.get_device_name()

                if capability is None:
                    compute_str = "does not have a compute capability"
                else:
                    version_str = capability.as_version_str()
                    compute_str = f"has compute capability {version_str}"

                raise ValueError(
                    "Bfloat16 is only supported on GPUs "
                    "with compute capability of at least 8.0. "
                    f"Your {gpu_name} GPU {compute_str}. "
                    "You can use float16 instead by explicitly setting the "
541
542
                    "`dtype` flag in CLI, for example: --dtype=half."
                )
543

544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
    @classmethod
    def insert_blocks_to_device(
        cls,
        src_cache: torch.Tensor,
        dst_cache: torch.Tensor,
        src_block_indices: torch.Tensor,
        dst_block_indices: torch.Tensor,
    ) -> None:
        """Copy blocks from src_cache to dst_cache on GPU."""
        _src_cache = src_cache[:, src_block_indices]
        dst_cache[:, dst_block_indices] = _src_cache.to(dst_cache.device)

    @classmethod
    def swap_out_blocks_to_host(
        cls,
        src_cache: torch.Tensor,
        dst_cache: torch.Tensor,
        src_block_indices: torch.Tensor,
        dst_block_indices: torch.Tensor,
    ) -> None:
        """Copy blocks from GPU to host (CPU)."""
        _src_cache = src_cache[:, src_block_indices]
        dst_cache[:, dst_block_indices] = _src_cache.cpu()

568
569
570
571
    @classmethod
    def support_hybrid_kv_cache(cls) -> bool:
        return True

572
573
574
575
    @classmethod
    def support_static_graph_mode(cls) -> bool:
        return True

576
577
578
579
    @classmethod
    def num_compute_units(cls, device_id=0):
        return torch.cuda.get_device_properties(device_id).multi_processor_count

580

581
582
583
584
585
# NVML utils
# Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`,
# all the related functions work on real physical device ids.
# the major benefit of using NVML is that it will not initialize CUDA
class NvmlCudaPlatform(CudaPlatformBase):
586
    @classmethod
587
    @cache
588
    @with_nvml_context
589
    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
590
        try:
591
            physical_device_id = cls.device_id_to_physical_device_id(device_id)
592
593
594
595
596
597
598
599
600
601
            handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
            major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle)
            return DeviceCapability(major=major, minor=minor)
        except RuntimeError:
            return None

    @classmethod
    @with_nvml_context
    def has_device_capability(
        cls,
602
        capability: tuple[int, int] | int,
603
604
605
606
607
608
        device_id: int = 0,
    ) -> bool:
        try:
            return super().has_device_capability(capability, device_id)
        except RuntimeError:
            return False
609

610
    @classmethod
611
    @with_nvml_context
612
    def get_device_name(cls, device_id: int = 0) -> str:
613
        physical_device_id = cls.device_id_to_physical_device_id(device_id)
614
        return cls._get_physical_device_name(physical_device_id)
615

616
617
618
    @classmethod
    @with_nvml_context
    def get_device_uuid(cls, device_id: int = 0) -> str:
619
        physical_device_id = cls.device_id_to_physical_device_id(device_id)
620
621
622
        handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
        return pynvml.nvmlDeviceGetUUID(handle)

623
    @classmethod
624
    @with_nvml_context
625
    def get_device_total_memory(cls, device_id: int = 0) -> int:
626
        physical_device_id = cls.device_id_to_physical_device_id(device_id)
627
628
        handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
        return int(pynvml.nvmlDeviceGetMemoryInfo(handle).total)
629

630
    @classmethod
631
    @with_nvml_context
632
    def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:
633
634
635
        """
        query if the set of gpus are fully connected by nvlink (1 hop)
        """
636
        handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_ids]
637
638
639
640
641
        for i, handle in enumerate(handles):
            for j, peer_handle in enumerate(handles):
                if i < j:
                    try:
                        p2p_status = pynvml.nvmlDeviceGetP2PStatus(
642
643
644
645
                            handle,
                            peer_handle,
                            pynvml.NVML_P2P_CAPS_INDEX_NVLINK,
                        )
646
647
                        if p2p_status != pynvml.NVML_P2P_STATUS_OK:
                            return False
648
649
                    except pynvml.NVMLError:
                        logger.exception(
650
                            "NVLink detection failed. This is normal if"
651
652
                            " your machine has no NVLink equipped."
                        )
653
654
                        return False
        return True
655
656

    @classmethod
657
658
659
660
661
662
663
664
665
    def _get_physical_device_name(cls, device_id: int = 0) -> str:
        handle = pynvml.nvmlDeviceGetHandleByIndex(device_id)
        return pynvml.nvmlDeviceGetName(handle)

    @classmethod
    @with_nvml_context
    def log_warnings(cls):
        device_ids: int = pynvml.nvmlDeviceGetCount()
        if device_ids > 1:
666
667
668
669
670
            device_names = [cls._get_physical_device_name(i) for i in range(device_ids)]
            if (
                len(set(device_names)) > 1
                and os.environ.get("CUDA_DEVICE_ORDER") != "PCI_BUS_ID"
            ):
671
                logger.warning(
672
                    "Detected different devices in the system: %s. Please"
673
674
                    " make sure to set `CUDA_DEVICE_ORDER=PCI_BUS_ID` to "
                    "avoid unexpected behavior.",
675
                    ", ".join(device_names),
676
677
678
679
680
                )


class NonNvmlCudaPlatform(CudaPlatformBase):
    @classmethod
681
    @cache
682
683
684
685
686
687
688
689
690
691
692
693
694
695
    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability:
        major, minor = torch.cuda.get_device_capability(device_id)
        return DeviceCapability(major=major, minor=minor)

    @classmethod
    def get_device_name(cls, device_id: int = 0) -> str:
        return torch.cuda.get_device_name(device_id)

    @classmethod
    def get_device_total_memory(cls, device_id: int = 0) -> int:
        device_props = torch.cuda.get_device_properties(device_id)
        return device_props.total_memory

    @classmethod
696
    def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:
697
698
        logger.exception(
            "NVLink detection not possible, as context support was"
699
700
            " not found. Assuming no NVLink available."
        )
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
        return False


# Autodetect either NVML-enabled or non-NVML platform
# based on whether NVML is available.
nvml_available = False
try:
    try:
        pynvml.nvmlInit()
        nvml_available = True
    except Exception:
        # On Jetson, NVML is not supported.
        nvml_available = False
finally:
    if nvml_available:
        pynvml.nvmlShutdown()

CudaPlatform = NvmlCudaPlatform if nvml_available else NonNvmlCudaPlatform

720
CudaPlatform.log_warnings()