cuda.py 23.6 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
8
from __future__ import annotations

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

15
import torch
16
17
from torch.distributed import PrefixStore, ProcessGroup
from torch.distributed.distributed_c10d import is_nccl_available
18
from typing_extensions import ParamSpec
19

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

28
from .interface import DeviceCapability, Platform, PlatformEnum
29

30
if TYPE_CHECKING:
31
    from vllm.config import VllmConfig
32
    from vllm.config.cache import CacheDType
33
    from vllm.v1.attention.selector import AttentionSelectorConfig
34
else:
35
36
    VllmConfig = None
    CacheDType = None
37

38
39
logger = init_logger(__name__)

40
41
42
_P = ParamSpec("_P")
_R = TypeVar("_R")

43
pynvml = import_pynvml()
44

45
46
47
48
# 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)

49

50
51
52
53
@cache
def _get_backend_priorities(
    use_mla: bool,
    device_capability: DeviceCapability,
54
    num_heads: int | None = None,
55
    kv_cache_dtype: CacheDType | None = None,
56
57
58
59
) -> list[AttentionBackendEnum]:
    """Get backend priorities with lazy import to avoid circular dependency."""
    if use_mla:
        if device_capability.major == 10:
60
61
62
63
64
            # Sparse MLA backend priorities
            # See https://github.com/vllm-project/vllm/issues/35807 for
            # benchmark results
            if kv_cache_dtype is not None and kv_cache_dtype.startswith("fp8"):
                # Prefer FlashInfer for fp8 kv cache
65
66
67
68
69
                sparse_backends = [
                    AttentionBackendEnum.FLASHINFER_MLA_SPARSE,
                    AttentionBackendEnum.FLASHMLA_SPARSE,
                ]
            else:
70
71
72
73
74
75
76
77
78
79
80
81
82
                # BF16 KV Cache
                # 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,
                    ]

83
            return [
84
                AttentionBackendEnum.FLASHINFER_MLA,
85
                AttentionBackendEnum.CUTLASS_MLA,
86
                AttentionBackendEnum.FLASH_ATTN_MLA,
87
                AttentionBackendEnum.FLASHMLA,
88
                AttentionBackendEnum.TRITON_MLA,
89
                *sparse_backends,
90
91
92
93
            ]
        else:
            return [
                AttentionBackendEnum.FLASH_ATTN_MLA,
94
                AttentionBackendEnum.FLASHMLA,
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
                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,
            ]


116
def with_nvml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]:
117
    @wraps(fn)
118
    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
119
120
121
122
123
124
125
126
127
        pynvml.nvmlInit()
        try:
            return fn(*args, **kwargs)
        finally:
            pynvml.nvmlShutdown()

    return wrapper


128
129
class CudaPlatformBase(Platform):
    _enum = PlatformEnum.CUDA
130
    device_name: str = "cuda"
131
132
    device_type: str = "cuda"
    dispatch_key: str = "CUDA"
133
    ray_device_key: str = "GPU"
134
    dist_backend: str = "nccl"
135
    device_control_env_var: str = "CUDA_VISIBLE_DEVICES"
136
137
138
    ray_noset_device_env_vars: list[str] = [
        "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES",
    ]
139

140
    @property
141
    def supported_dtypes(self) -> list[torch.dtype]:
142
143
144
        if self.has_device_capability(80):
            # Ampere and Hopper or later NVIDIA GPUs.
            return [torch.bfloat16, torch.float16, torch.float32]
145
        if self.has_device_capability(60):
146
147
148
149
150
151
            # 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]

152
153
154
155
156
    @classmethod
    def set_device(cls, device: torch.device) -> None:
        """
        Set the device for the current platform.
        """
157
        torch.cuda.set_device(device)
158
159
160
161
162
        # 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)

163
    @classmethod
164
    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
165
        raise NotImplementedError
166

167
168
169
    @classmethod
    def get_device_name(cls, device_id: int = 0) -> str:
        raise NotImplementedError
170

171
172
173
    @classmethod
    def get_device_total_memory(cls, device_id: int = 0) -> int:
        raise NotImplementedError
174

175
    @classmethod
176
    def is_fully_connected(cls, device_ids: list[int]) -> bool:
177
        raise NotImplementedError
178

179
180
181
    @classmethod
    def log_warnings(cls):
        pass
182

183
    @classmethod
184
    def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
185
        parallel_config = vllm_config.parallel_config
186
        model_config = vllm_config.model_config
187

188
        if parallel_config.worker_cls == "auto":
189
            parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"
190

191
192
193
194
195
196
197
198
199
200
201
202
203
204
        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

205
    @classmethod
206
    def get_current_memory_usage(
207
        cls, device: torch.types.Device | None = None
208
    ) -> float:
209
        torch.cuda.empty_cache()
210
211
212
        torch.cuda.reset_peak_memory_stats(device)
        return torch.cuda.max_memory_allocated(device)

213
    @classmethod
214
    def get_valid_backends(
215
        cls,
216
        device_capability: DeviceCapability,
217
        attn_selector_config: AttentionSelectorConfig,
218
        num_heads: int | None = None,
219
    ) -> tuple[
220
221
        list[tuple[AttentionBackendEnum, int]],
        dict[AttentionBackendEnum, tuple[int, list[str]]],
222
223
    ]:
        valid_backends_priorities = []
224
        invalid_reasons: dict[AttentionBackendEnum, tuple[int, list[str]]] = {}
225

226
        backend_priorities = _get_backend_priorities(
227
228
229
            attn_selector_config.use_mla,
            device_capability,
            num_heads,
230
            attn_selector_config.kv_cache_dtype,
231
        )
232
233
234
235
        for priority, backend in enumerate(backend_priorities):
            try:
                backend_class = backend.get_class()
                invalid_reasons_i = backend_class.validate_configuration(
236
237
                    device_capability=device_capability,
                    **attn_selector_config._asdict(),
238
                )
239
240
241
            except ImportError:
                invalid_reasons_i = ["ImportError"]
            if invalid_reasons_i:
242
                invalid_reasons[backend] = (priority, invalid_reasons_i)
243
244
            else:
                valid_backends_priorities.append((backend, priority))
245

246
        return valid_backends_priorities, invalid_reasons
247

248
    @classmethod
249
    def get_attn_backend_cls(
250
        cls,
251
252
        selected_backend: AttentionBackendEnum | None,
        attn_selector_config: AttentionSelectorConfig,
253
        num_heads: int | None = None,
254
255
256
257
    ) -> str:
        device_capability = cls.get_device_capability()
        assert device_capability is not None

258
259
260
261
        # First try checking just the selected backend, if there is one.
        if selected_backend is not None:
            try:
                backend_class = selected_backend.get_class()
262
                invalid_reasons = backend_class.validate_configuration(
263
264
                    device_capability=device_capability,
                    **attn_selector_config._asdict(),
265
                )
266
            except ImportError:
267
268
269
270
271
272
273
274
275
                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()
276

277
278
        # No selected backend or the selected backend is invalid,
        # so we try finding a valid backend.
279
        valid_backends_priorities, all_invalid_reasons = cls.get_valid_backends(
280
281
            device_capability=device_capability,
            attn_selector_config=attn_selector_config,
282
            num_heads=num_heads,
283
        )
284
285
286
287
        reasons_str = (
            "{"
            + ", ".join(
                f"{backend.name}: [{', '.join(reasons)}]"
288
                for backend, (_, reasons) in all_invalid_reasons.items()
289
290
291
292
293
294
295
296
            )
            + "}"
        )
        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}."
        )
297
        if len(valid_backends_priorities) == 0:
298
299
300
301
            raise ValueError(
                f"No valid attention backend found for {cls.device_name} "
                f"with {config_str}. Reasons: {reasons_str}."
            )
302

303
304
305
306
307
308
309
310
        # 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]
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
        selected_priority = valid_backends_priorities[selected_index][1]

        # If the user specified --block-size (but not --attention-backend),
        # check whether that constraint precluded any higher-priority backends.
        if attn_selector_config.block_size is not None:
            excluded = [
                backend
                for backend, (priority, reasons) in all_invalid_reasons.items()
                if priority < selected_priority
                and reasons == ["block_size not supported"]
            ]
            if excluded:
                names = ", ".join(b.name for b in excluded)
                logger.warning(
                    "--block-size %d precluded higher-priority backend(s) "
                    "%s. Using %s instead, which may result in reduced "
                    "performance. Consider removing --block-size to "
                    "auto-select the optimal block size.",
                    attn_selector_config.block_size,
                    names,
                    selected_backend.name,
                )

334
335
336
337
338
        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",
339
340
        )

341
        return selected_backend.get_path()
342

343
    @classmethod
344
    def get_supported_vit_attn_backends(cls) -> list[AttentionBackendEnum]:
345
346
347
348
349
350
351
352
353
354
355
356
357
358
        if cls.has_device_capability(80):
            return [
                AttentionBackendEnum.FLASH_ATTN,
                AttentionBackendEnum.TRITON_ATTN,
                AttentionBackendEnum.TORCH_SDPA,
                AttentionBackendEnum.FLASHINFER,
            ]
        else:
            return [
                AttentionBackendEnum.FLASH_ATTN,
                AttentionBackendEnum.TORCH_SDPA,
                AttentionBackendEnum.TRITON_ATTN,
                AttentionBackendEnum.FLASHINFER,
            ]
359
360
361
362
363
364

    @classmethod
    def get_vit_attn_backend(
        cls,
        head_size: int,
        dtype: torch.dtype,
365
366
        backend: AttentionBackendEnum | None = None,
    ) -> AttentionBackendEnum:
367
368
369
370
371
372
373
374
        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

375
376
377
        cc = cls.get_device_capability()
        for vit_attn_backend in cls.get_supported_vit_attn_backends():
            if vit_attn_backend == AttentionBackendEnum.TORCH_SDPA:
378
                return vit_attn_backend
379
            try:
380
381
                backend_class = vit_attn_backend.get_class()
                is_backend_supported = backend_class.supports_head_size(
382
                    head_size
383
384
385
386
387
388
389
390
                ) 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(
391
392
                        f"Using backend {vit_attn_backend} for vit attention",
                        scope="local",
393
394
                    )
                    return vit_attn_backend
395
396
397
398
399
            except ImportError:
                pass

        return AttentionBackendEnum.TORCH_SDPA

400
401
402
403
    @classmethod
    def get_punica_wrapper(cls) -> str:
        return "vllm.lora.punica_wrapper.punica_gpu.PunicaWrapperGPU"

404
405
    @classmethod
    def get_device_communicator_cls(cls) -> str:
406
407
408
        return (
            "vllm.distributed.device_communicators.cuda_communicator.CudaCommunicator"  # noqa
        )
409

410
411
412
413
    @classmethod
    def supports_fp8(cls) -> bool:
        return cls.has_device_capability(89)

414
415
416
417
    @classmethod
    def use_custom_allreduce(cls) -> bool:
        return True

418
419
420
421
    @classmethod
    def opaque_attention_op(cls) -> bool:
        return True

422
    @classmethod
423
424
    def get_static_graph_wrapper_cls(cls) -> str:
        return "vllm.compilation.cuda_graph.CUDAGraphWrapper"
425

426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
    @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

457
458
459
460
    @classmethod
    def device_count(cls) -> int:
        return cuda_device_count_stateless()

461
    @classmethod
462
463
    def check_if_supports_dtype(cls, dtype: torch.dtype):
        if dtype == torch.bfloat16:  # noqa: SIM102
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
            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 "
479
480
                    "`dtype` flag in CLI, for example: --dtype=half."
                )
481

482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
    @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()

506
507
508
509
    @classmethod
    def support_hybrid_kv_cache(cls) -> bool:
        return True

510
511
512
513
    @classmethod
    def support_static_graph_mode(cls) -> bool:
        return True

514
515
516
517
518
    @classmethod
    def support_deep_gemm(cls) -> bool:
        """Currently, only Hopper and Blackwell GPUs are supported."""
        return cls.is_device_capability(90) or cls.is_device_capability_family(100)

519
    @classmethod
520
    def num_compute_units(cls, device_id: int = 0) -> int:
521
522
        return torch.cuda.get_device_properties(device_id).multi_processor_count

523
524
525
526
    @classmethod
    def use_custom_op_collectives(cls) -> bool:
        return True

527

528
529
530
531
532
# 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):
533
    @classmethod
534
    @cache
535
    @with_nvml_context
536
    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
537
        try:
538
            physical_device_id = cls.device_id_to_physical_device_id(device_id)
539
540
541
542
543
544
545
546
547
548
            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,
549
        capability: tuple[int, int] | int,
550
551
552
553
554
555
        device_id: int = 0,
    ) -> bool:
        try:
            return super().has_device_capability(capability, device_id)
        except RuntimeError:
            return False
556

557
    @classmethod
558
    @with_nvml_context
559
    def get_device_name(cls, device_id: int = 0) -> str:
560
        physical_device_id = cls.device_id_to_physical_device_id(device_id)
561
        return cls._get_physical_device_name(physical_device_id)
562

563
564
565
    @classmethod
    @with_nvml_context
    def get_device_uuid(cls, device_id: int = 0) -> str:
566
        physical_device_id = cls.device_id_to_physical_device_id(device_id)
567
568
569
        handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
        return pynvml.nvmlDeviceGetUUID(handle)

570
    @classmethod
571
    @with_nvml_context
572
    def get_device_total_memory(cls, device_id: int = 0) -> int:
573
        physical_device_id = cls.device_id_to_physical_device_id(device_id)
574
575
        handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
        return int(pynvml.nvmlDeviceGetMemoryInfo(handle).total)
576

577
    @classmethod
578
    @with_nvml_context
579
    def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:
580
581
582
        """
        query if the set of gpus are fully connected by nvlink (1 hop)
        """
583
        handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_ids]
584
585
586
587
588
        for i, handle in enumerate(handles):
            for j, peer_handle in enumerate(handles):
                if i < j:
                    try:
                        p2p_status = pynvml.nvmlDeviceGetP2PStatus(
589
590
591
592
                            handle,
                            peer_handle,
                            pynvml.NVML_P2P_CAPS_INDEX_NVLINK,
                        )
593
594
                        if p2p_status != pynvml.NVML_P2P_STATUS_OK:
                            return False
595
596
                    except pynvml.NVMLError:
                        logger.exception(
597
                            "NVLink detection failed. This is normal if"
598
599
                            " your machine has no NVLink equipped."
                        )
600
601
                        return False
        return True
602
603

    @classmethod
604
605
606
607
608
609
610
611
612
    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:
613
614
615
616
617
            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"
            ):
618
                logger.warning(
619
                    "Detected different devices in the system: %s. Please"
620
621
                    " make sure to set `CUDA_DEVICE_ORDER=PCI_BUS_ID` to "
                    "avoid unexpected behavior.",
622
                    ", ".join(device_names),
623
624
625
626
627
                )


class NonNvmlCudaPlatform(CudaPlatformBase):
    @classmethod
628
    @cache
629
630
631
632
633
634
635
636
637
638
639
640
641
642
    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
643
    def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:
644
645
        logger.exception(
            "NVLink detection not possible, as context support was"
646
647
            " not found. Assuming no NVLink available."
        )
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
        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

667
CudaPlatform.log_warnings()