Fix NVML capability lookup for visible devices (#47892)
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Tyler Michael Smith committed
Jul 8, 2026 at 09:07 UTC
68b4a1d582818e67adc903bf1b8fc5a5447da2fa
3 files changed
+55
-2
tests/cuda/test_cuda_context.py
+38
@@ -77,5 +77,43 @@ class TestSetCudaContext:
77
current_platform.set_device(torch.device("cpu"))
78
79
80
+def test_get_device_capability_uses_visible_device_ordinal(monkeypatch):
81
+ import vllm.platforms.interface as platform_interface
82
+ from vllm.platforms.cuda import NvmlCudaPlatform, pynvml
83
+
84
+ seen_indices: list[int] = []
85
+
86
+ def record_handle(index: int) -> str:
87
+ seen_indices.append(index)
88
+ return f"handle-{index}"
89
+
90
+ monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [1])
91
+ monkeypatch.setenv(NvmlCudaPlatform.device_control_env_var, "0,1")
92
+ monkeypatch.setattr(
93
+ NvmlCudaPlatform,
94
+ "device_control_id_to_physical_device_id",
95
+ classmethod(lambda _cls, device_id: int(device_id)),
96
+ )
97
+ monkeypatch.setattr(pynvml, "nvmlInit", lambda: None)
98
+ monkeypatch.setattr(pynvml, "nvmlShutdown", lambda: None)
99
+ monkeypatch.setattr(
100
+ pynvml,
101
+ "nvmlDeviceGetHandleByIndex",
102
+ record_handle,
103
+ )
104
+ monkeypatch.setattr(
105
+ pynvml,
106
+ "nvmlDeviceGetCudaComputeCapability",
107
+ lambda _handle: (9, 0),
108
+ )
109
+ NvmlCudaPlatform.get_device_capability.cache_clear()
110
+
111
+ capability = NvmlCudaPlatform.get_device_capability(device_id=1)
112
+
113
+ assert capability is not None
114
+ assert capability.to_int() == 90
115
+ assert seen_indices == [1]
116
+
117
+
118
if __name__ == "__main__":
119
pytest.main([__file__, "-v"])
vllm/platforms/cuda.py
+1
-1
@@ -727,7 +727,7 @@ class NvmlCudaPlatform(CudaPlatformBase):
727
@with_nvml_context
728
def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
729
try:
730
- physical_device_id = cls.device_id_to_physical_device_id(device_id)
730
+ physical_device_id = cls.visible_device_id_to_physical_device_id(device_id)
731
handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
732
major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle)
733
return DeviceCapability(major=major, minor=minor)
vllm/platforms/interface.py
+16
-1
@@ -271,6 +271,16 @@ class Platform:
271
f"{cls.device_name}."
272
) from e
273
274
+ # GPU device IDs can refer to three distinct namespaces:
275
+ # - logical: vLLM-local IDs such as local ranks. These index
276
+ # assigned_physical_gpu_ids when it is set.
277
+ # - visible: torch/CUDA ordinals in the current process after applying
278
+ # the device-control env var, e.g. CUDA_VISIBLE_DEVICES.
279
+ # - physical: global GPU IDs used by topology and management APIs such as
280
+ # NVML, which are not remapped by CUDA_VISIBLE_DEVICES.
281
+ # Keep conversions explicit. In particular, torch device indices are
282
+ # visible IDs, not vLLM logical IDs.
283
+
284
@classmethod
285
def device_id_to_physical_device_id(cls, device_id: int):
286
"""Map a vLLM-local logical device ID to a physical device ID.
@@ -411,7 +421,12 @@ class Platform:
421
cls,
422
device_id: int = 0,
423
) -> DeviceCapability | None:
414
- """Stateless version of [torch.cuda.get_device_capability][]."""
424
+ """Stateless version of [torch.cuda.get_device_capability][].
425
+
426
+ Args:
427
+ device_id: Device index in the visible device namespace, matching
428
+ the argument accepted by torch.cuda.
429
+ """
430
return None
431
432
@classmethod