Coverage for cuda / core / system / _nvlink.pxi: 15.79%
19 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-22 01:37 +0000
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-22 01:37 +0000
1# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2#
3# SPDX-License-Identifier: Apache-2.0
6_NVLINK_VERSION_MAPPING = {
7 nvml.NvlinkVersion.VERSION_1_0: (1, 0),
8 nvml.NvlinkVersion.VERSION_2_0: (2, 0),
9 nvml.NvlinkVersion.VERSION_2_2: (2, 2),
10 nvml.NvlinkVersion.VERSION_3_0: (3, 0),
11 nvml.NvlinkVersion.VERSION_3_1: (3, 1),
12 nvml.NvlinkVersion.VERSION_4_0: (4, 0),
13 nvml.NvlinkVersion.VERSION_5_0: (5, 0),
14}
17cdef class NvlinkInfo:
18 """
19 Nvlink information for a device.
20 """
21 cdef Device _device
22 cdef int _link
24 def __init__(self, device: Device, link: int):
25 self._device = device 1a
26 self._link = link 1a
28 @property
29 def version(self) -> tuple[int, int]:
30 """
31 Retrieves the NvLink version for the device and link.
33 For all products with NvLink support.
35 Returns
36 -------
37 tuple[int, int]
38 The Nvlink version as a tuple of (major, minor).
39 """
40 version = nvml.device_get_nvlink_version(self._device._handle, self._link) 1a
41 if version == nvml.NvlinkVersion.VERSION_INVALID:
42 raise RuntimeError("Invalid NvLink version returned for device")
43 try:
44 return _NVLINK_VERSION_MAPPING[version]
45 except KeyError:
46 raise RuntimeError(f"Unknown NvLink version {version} returned for device") from None
48 @property
49 def state(self) -> bool:
50 """
51 Retrieves the state of the device's Nvlink for the device and link specified.
53 For Pascal™ or newer fully supported devices.
55 For all products with Nvlink support.
57 Returns
58 -------
59 bool
60 `True` if the Nvlink is active.
61 """
62 return (
63 nvml.device_get_nvlink_state(self._device._handle, self._link) == nvml.EnableState.FEATURE_ENABLED
64 )
66 max_links = nvml.NVLINK_MAX_LINKS